module Shady.Vec
(
Z, S, (:+:), ZeroT, OneT, TwoT, ThreeT, FourT
, Nat(..), zero, one, two, three, four
, withIsNat, natSucc, natIsNat
, natToZ, natEq, natAdd, (:<:)
, Index(..), succI, index0, index1, index2, index3
, Vec(..), IsNat(..), (<+>), indices
, Zero, One, Two, Three, Four, vElems
, vec1, vec2, vec3, vec4
, un1, un2, un3, un4
, get0, get1, get2, get3
, get, swizzle
) where
import Prelude hiding (foldr,sum)
import Control.Applicative (Applicative(..),liftA2,(<$>))
import Data.Foldable (Foldable(..),sum)
import Data.Maybe (isJust)
import Foreign.Storable
import Foreign.Ptr (Ptr,plusPtr,castPtr)
import Control.Compose (result)
import Data.VectorSpace
import Shady.Misc (Sink)
import Data.Proof.EQ
data Z
data S n
infixl 6 :+:
type family a :+: b
type instance Z :+: b = b
type instance S a :+: b = S (a :+: b)
type ZeroT = Z
type OneT = S ZeroT
type TwoT = S OneT
type ThreeT = S TwoT
type FourT = S ThreeT
data Nat :: * -> * where
Zero :: Nat Z
Succ :: IsNat n => Nat n -> Nat (S n)
instance Show (Nat n) where show = show . natToZ
withIsNat :: (IsNat n => Nat n -> a) -> (Nat n -> a)
withIsNat p Zero = p Zero
withIsNat p (Succ n) = p (Succ n)
natSucc :: Nat n -> Nat (S n)
natSucc = withIsNat Succ
natIsNat :: Nat n -> (IsNat n => Nat n)
natIsNat Zero = Zero
natIsNat (Succ n) = Succ n
natToZ :: Nat n -> Integer
natToZ Zero = 0
natToZ (Succ n) = (succ . natToZ) n
natEq :: Nat m -> Nat n -> Maybe (m :=: n)
Zero `natEq` Zero = Just Refl
Succ m `natEq` Succ n = liftEq <$> (m `natEq` n)
_ `natEq` _ = Nothing
natAdd :: Nat m -> Nat n -> Nat (m :+: n)
Zero `natAdd` n = n
Succ m `natAdd` n = natSucc (m `natAdd` n)
zero :: Nat ZeroT
zero = Zero
one :: Nat OneT
one = Succ zero
two :: Nat TwoT
two = Succ one
three :: Nat ThreeT
three = Succ two
four :: Nat FourT
four = Succ three
infix 4 :<:
data m :<: n where
ZLess :: Z :<: S n
SLess :: m :<: n -> S m :<: S n
data Index lim = forall n. IsNat n => Index (n :<: lim) (Nat n)
instance Eq (Index lim) where
Index _ n == Index _ n' = isJust (n `natEq` n')
succI :: Index m -> Index (S m)
succI (Index p m) = Index (SLess p) (Succ m)
index0 :: Index (S n)
index0 = Index ZLess Zero
index1 :: Index (S (S n))
index1 = succI index0
index2 :: Index (S (S (S n)))
index2 = succI index1
index3 :: Index (S (S (S (S n))))
index3 = succI index2
infixr 5 :<
data Vec :: * -> * -> * where
ZVec :: Vec Z a
(:<) :: a -> Vec n a -> Vec (S n) a
vElems :: Vec n a -> [a]
vElems = foldr (:) []
instance Functor (Vec n) where
fmap _ ZVec = ZVec
fmap f (a :< u) = f a :< fmap f u
class IsNat n where
nat :: Nat n
pureV :: a -> Vec n a
elemsV :: [a] -> Vec n a
peekV :: Storable a => Ptr a -> IO (Vec n a)
pokeV :: Storable a => Ptr a -> Sink (Vec n a)
instance IsNat Z where
nat = Zero
pureV _ = ZVec
elemsV [] = ZVec
elemsV (_:_) = error "elemsV: too many elements"
peekV = const (return ZVec)
pokeV = const (const (return ()))
instance IsNat n => IsNat (S n) where
nat = Succ nat
pureV a = a :< pureV a
elemsV [] = error "elemsV: too few elements"
elemsV (a : as) = a :< elemsV as
peekV p = do a <- peek p
as <- peekV (p `plusPtr` sizeOf a)
return (a :< as)
pokeV p (a :< as) = do poke p a
pokeV (p `plusPtr` sizeOf a) as
applyV :: Vec n (a -> b) -> Vec n a -> Vec n b
ZVec `applyV` ZVec = ZVec
(f :< fs) `applyV` (x :< xs) = f x :< (fs `applyV` xs)
instance IsNat n => Applicative (Vec n) where
pure = pureV
(<*>) = applyV
instance Foldable (Vec n) where
foldr _ b ZVec = b
foldr h b (a :< as) = a `h` foldr h b as
infixl 1 <+>
(<+>) :: Vec m a -> Vec n a -> Vec (m :+: n) a
ZVec <+> v = v
(a :< u) <+> v = a :< (u <+> v)
indices :: Nat n -> Vec n (Index n)
indices Zero = ZVec
indices (Succ n) = index0 :< fmap succI (indices n)
type Zero = Vec ZeroT
type One = Vec OneT
type Two = Vec TwoT
type Three = Vec ThreeT
type Four = Vec FourT
vec1 :: a -> One a
vec1 a = a :< ZVec
vec2 :: a -> a -> Two a
vec2 a b = a :< vec1 b
vec3 :: a -> a -> a -> Three a
vec3 a b c = a :< vec2 b c
vec4 :: a -> a -> a -> a -> Four a
vec4 a b c d = a :< vec3 b c d
un1 :: One a -> a
un1 (a :< ZVec) = a
un2 :: Two a -> (a,a)
un2 (a :< b :< ZVec) = (a,b)
un3 :: Three a -> (a,a,a)
un3 (a :< b :< c :< ZVec) = (a,b,c)
un4 :: Four a -> (a,a,a,a)
un4 (a :< b :< c :< d :< ZVec) = (a,b,c,d)
instance (IsNat n, Num a) => AdditiveGroup (Vec n a) where
{ zeroV = pure 0; (^+^) = liftA2 (+) ; negateV = fmap negate }
instance (IsNat n, Num a) => VectorSpace (Vec n a) where
type Scalar (Vec n a) = One a
(*^) (s :< ZVec) = fmap (s *)
instance (IsNat n, Num a) => InnerSpace (Vec n a) where
(<.>) = (result.result) (vec1 . sum) (liftA2 (*))
get :: Index n -> Vec n a -> One a
get (Index ZLess Zero ) (a :< _) = vec1 a
get (Index (SLess p) (Succ m)) (_ :< as) = get (Index p m) as
get0 :: Vec (S n) a -> One a
get1 :: Vec (S (S n)) a -> One a
get2 :: Vec (S (S (S n))) a -> One a
get3 :: Vec (S (S (S (S n)))) a -> One a
get0 = get index0
get1 = get index1
get2 = get index2
get3 = get index3
swizzle :: Vec n (Index m) -> Vec m a -> Vec n a
swizzle ZVec _ = ZVec
swizzle (ix :< ixs) v = un1 (get ix v) :< swizzle ixs v
instance Eq a => Eq (Vec n a) where
ZVec == ZVec = True
a :< as == b :< bs = a==b && as==bs
instance Ord a => Ord (Vec n a) where
ZVec `compare` ZVec = EQ
(a :< as) `compare` (b :< bs) =
case a `compare` b of
LT -> LT
GT -> GT
EQ -> as `compare` bs
instance (IsNat n, Storable a) => Storable (Vec n a) where
sizeOf = const (fromIntegral (natToZ (nat :: Nat n))
* sizeOf (undefined :: a))
alignment = const (alignment (undefined :: a))
peek = peekV . castPtr
poke = pokeV . castPtr