Vector.hs

{-# LANGUAGE PartialTypeSignatures, ConstraintKinds, InstanceSigs, ScopedTypeVariables, RankNTypes, UndecidableInstances, TypeFamilies, DataKinds, PolyKinds, MultiParamTypeClasses, FunctionalDependencies, FlexibleInstances, FlexibleContexts, GADTs, DeriveTraversable, StandaloneDeriving #-}

-- | Quick(slow) and dirty typesafe vectors
module Vector where

import Prelude hiding (id, (.))
import Control.Category
import Data.Semigroupoid
import Data.Proxy
import Data.Kind

data Natural = Z | S Natural
data Vec (n :: Natural) a where
  VNil :: Vec Z a
  VCons :: a -> Vec n a -> Vec (S n) a

deriving instance Eq a => Eq (Vec n a)
deriving instance Ord a => Ord (Vec n a)
deriving instance Show a => Show (Vec n a)
deriving instance Functor (Vec n)
deriving instance Foldable (Vec n)
deriving instance Traversable (Vec n)

dotProduct :: Num a => Vec n a -> Vec n a -> a
dotProduct VNil VNil = 0
dotProduct (VCons x xs) (VCons y ys) = x * y + dotProduct xs ys

vmult :: Num a => Vec i (Vec k a) -> Vec j (Vec k a) -> Vec j (Vec i a)
vmult _ VNil = VNil
vmult xs (VCons y ys) = VCons (dotProduct y <$> xs) $ vmult xs ys

data KNat n where
  KZ :: KNat Z
  KS :: Finite n => KNat n -> KNat (S n)

class Finite (a :: Natural) where toFNat :: proxy a -> KNat a
instance Finite Z where toFNat _ = KZ
instance Finite n => Finite (S n) where toFNat _= KS (toFNat (Proxy :: Proxy n))

vdiag :: forall a i j. a -> a -> a -> KNat i -> KNat j -> Vec i (Vec j a)
vdiag u d l = go
  where
    go :: forall i' j'. KNat i' -> KNat j' -> Vec i' (Vec j' a)
    go (KS i) (KS j) =
      VCons (VCons d  $  vpure u j)
            (VCons l <$> go i j)
    go KZ _ = VNil
    go (KS i) KZ = vpure VNil (KS i)

vpure :: a -> KNat m -> Vec m a
vpure x KZ = VNil
vpure x (KS n) = VCons x $ vpure x n


instance Finite n => Applicative (Vec n) where
  pure :: forall a. a -> Vec n a
  pure x = vpure x (toFNat (Proxy :: Proxy n))

  (<*>) :: forall a b. Vec n (a -> b) -> Vec n a -> Vec n b
  nfs <*> nxs = go (toFNat (Proxy :: Proxy n)) nfs nxs
    where
      go :: forall (m :: Natural). KNat m -> Vec m (a -> b) -> Vec m a -> Vec m b
      go KZ VNil VNil = VNil
      go (KS n) (VCons f fs) (VCons x xs) = VCons (f x) (go n fs xs)

data Matrix a i j where
  DiagonalMatrix :: (Finite i => KNat i -> Vec i (Vec i a)) -> Matrix a i i
  Matrix :: (Finite i, Finite j) => Vec i (Vec j a) -> Matrix a i j

instance Num a => Semigroupoid (Matrix a) where
  o :: forall a i j k. Num a => Matrix a k j -> Matrix a i k -> Matrix a i j
  Matrix x          `o` Matrix y          = Matrix (vmult (sequenceA   x                             )   y)
  DiagonalMatrix fx `o` Matrix y          = Matrix (vmult (sequenceA (fx (toFNat (Proxy :: Proxy k))))   y)
  Matrix x          `o` DiagonalMatrix fy = Matrix (vmult (sequenceA   x                             ) (fy (toFNat (Proxy :: Proxy k))))
  DiagonalMatrix fx `o` DiagonalMatrix fy = DiagonalMatrix $ \i -> vmult (sequenceA (fx i)) (fy i)

instance Num a => Category (Matrix a) where
  (.) = o
  id = DiagonalMatrix $ \i -> vdiag 0 1 0 i i

unMatrix :: (Finite i, Finite j) => Matrix a i j -> Vec i (Vec j a)
unMatrix (Matrix x) = x
unMatrix (DiagonalMatrix fx) = fx (toFNat (Proxy))

type Zero = Z
type One = S Z
type Two = S One
type Three = S Two
type Four = S Three

f :: Vec Two (Vec Three Int)
f = VCons (VCons 1 $ VCons 2 $ VCons 3 VNil)
  $ VCons (VCons 4 $ VCons 5 $ VCons 6 VNil)
  $ VNil

g :: Vec Four (Vec Two Int)
g = VCons (VCons 1 $ VCons 2 VNil)
  $ VCons (VCons 3 $ VCons 4 VNil)
  $ VCons (VCons 5 $ VCons 6 VNil)
  $ VCons (VCons 7 $ VCons 8 VNil)
  $ VNil

fg = unMatrix $ Matrix f . id . Matrix g

data Fin n where
  ZZ :: Fin (S n)
  SS :: Fin n -> Fin (S n)

vlookup :: Fin n -> Vec n a -> a
vlookup ZZ (VCons x _) = x
vlookup (SS i) (VCons _ xs) = vlookup i xs