{-# LANGUAGE TypeInType, TypeFamilies, TypeOperators, GADTs,
   StandaloneDeriving #-}

import Data.Kind

data Nat where
  Z :: Nat
  S :: Nat -> Nat

data Vector :: Nat -> Type -> Type where
  Nil  :: Vector Z a
  (:|) :: a -> Vector x a -> Vector (S x) a

infixr 7 :|

deriving instance Show a => Show (Vector n a)

data SNat :: Nat -> Type where
  SZ :: SNat Z
  SS :: SNat x -> SNat (S x)

type family Min n m where
  Min Z x = Z
  Min x Z = Z
  Min (S x) (S y) = S (Min x y)

take' :: SNat n -> Vector k a -> Vector (Min n k) a
take' SZ xs = Nil
take' (SS k) Nil = Nil
take' (SS k) (x :| xs) = x :| take' k xs

-- finite sets of size k
data Fin :: Nat -> Type where
  FZ :: Fin (S n)
  FS :: Fin n -> Fin (S n)

-- this is half of an isomorphism
index :: Vector n a -> Fin n -> a
index (x :| xs) FZ = x
index (_ :| xs) (FS k) = index xs k

-- this is the other half, but we need the length of the nat to be
-- statically known to do this (could use a class a la KnownNat to
-- automatically make SNat)
tabulate :: SNat n -> (Fin n -> a) -> Vector n a
tabulate SZ f = Nil
tabulate (SS k) f = f FZ :| tabulate k (f . FS)

type family Add n m where
  Add Z x = x
  Add (S k) n = S (Add k n)

-- note that the pattern of definition of the recursive function on
-- vectors (append, take', etc) matches the definition of the
-- corresponding type family on naturals. this is because natural
-- numbers are essentially lists of ()
append :: Vector n a -> Vector k a -> Vector (Add n k) a
append Nil xs = xs
append (x :| xs) ys = x :| append xs ys

-- need to use an existential because the length of the filtered vector
-- depends on the given function in a way that can not be captured
-- statically
data SomeVect a where
  Mk :: Vector n a -> SomeVect a

filter' :: (a -> Bool) -> Vector n a -> SomeVect a
filter' f Nil = Mk Nil
filter' f (x :| xs)
  | Mk tail <- filter' f xs
  , f x = Mk (x :| tail)
  | Mk tail <- filter' f xs = Mk tail

-- ghc will literally not let you get this wrong:
length' :: Vector n a -> SNat n
length' Nil = SZ
length' (_ :| xs) = SS (length' xs)