diff options
Diffstat (limited to 'src/Data/Array')
| -rw-r--r-- | src/Data/Array/Mixed.hs | 108 | ||||
| -rw-r--r-- | src/Data/Array/Nested/Internal.hs | 86 |
2 files changed, 97 insertions, 97 deletions
diff --git a/src/Data/Array/Mixed.hs b/src/Data/Array/Mixed.hs index 8d20583..c19fbe5 100644 --- a/src/Data/Array/Mixed.hs +++ b/src/Data/Array/Mixed.hs @@ -47,20 +47,20 @@ lemAppAssoc _ _ _ = unsafeCoerce Refl type IxX :: Type -> [Maybe Nat] -> Type data IxX i sh where - IZX :: IxX i '[] - (::@) :: forall n sh i. i -> IxX i sh -> IxX i (Just n : sh) - (::?) :: forall sh i. i -> IxX i sh -> IxX i (Nothing : sh) + ZIX :: IxX i '[] + (:.@) :: forall n sh i. i -> IxX i sh -> IxX i (Just n : sh) + (:.?) :: forall sh i. i -> IxX i sh -> IxX i (Nothing : sh) deriving instance Show i => Show (IxX i sh) deriving instance Eq i => Eq (IxX i sh) -infixr 3 ::@ -infixr 3 ::? +infixr 3 :.@ +infixr 3 :.? type IIxX = IxX Int -- | The part of a shape that is statically known. type StaticShapeX :: [Maybe Nat] -> Type data StaticShapeX sh where - SZX :: StaticShapeX '[] + ZSX :: StaticShapeX '[] (:$@) :: SNat n -> StaticShapeX sh -> StaticShapeX (Just n : sh) (:$?) :: () -> StaticShapeX sh -> StaticShapeX (Nothing : sh) deriving instance Show (StaticShapeX sh) @@ -72,7 +72,7 @@ type KnownShapeX :: [Maybe Nat] -> Constraint class KnownShapeX sh where knownShapeX :: StaticShapeX sh instance KnownShapeX '[] where - knownShapeX = SZX + knownShapeX = ZSX instance (KnownNat n, KnownShapeX sh) => KnownShapeX (Just n : sh) where knownShapeX = natSing :$@ knownShapeX instance KnownShapeX sh => KnownShapeX (Nothing : sh) where @@ -87,39 +87,39 @@ newtype XArray sh a = XArray (S.Array (FromINat (Rank sh)) a) deriving (Show) zeroIxX :: StaticShapeX sh -> IIxX sh -zeroIxX SZX = IZX -zeroIxX (_ :$@ ssh) = 0 ::@ zeroIxX ssh -zeroIxX (_ :$? ssh) = 0 ::? zeroIxX ssh +zeroIxX ZSX = ZIX +zeroIxX (_ :$@ ssh) = 0 :.@ zeroIxX ssh +zeroIxX (_ :$? ssh) = 0 :.? zeroIxX ssh zeroIxX' :: IIxX sh -> IIxX sh -zeroIxX' IZX = IZX -zeroIxX' (_ ::@ sh) = 0 ::@ zeroIxX' sh -zeroIxX' (_ ::? sh) = 0 ::? zeroIxX' sh +zeroIxX' ZIX = ZIX +zeroIxX' (_ :.@ sh) = 0 :.@ zeroIxX' sh +zeroIxX' (_ :.? sh) = 0 :.? zeroIxX' sh ixAppend :: IIxX sh -> IIxX sh' -> IIxX (sh ++ sh') -ixAppend IZX idx' = idx' -ixAppend (i ::@ idx) idx' = i ::@ ixAppend idx idx' -ixAppend (i ::? idx) idx' = i ::? ixAppend idx idx' +ixAppend ZIX idx' = idx' +ixAppend (i :.@ idx) idx' = i :.@ ixAppend idx idx' +ixAppend (i :.? idx) idx' = i :.? ixAppend idx idx' ixDrop :: IIxX (sh ++ sh') -> IIxX sh -> IIxX sh' -ixDrop sh IZX = sh -ixDrop (_ ::@ sh) (_ ::@ idx) = ixDrop sh idx -ixDrop (_ ::? sh) (_ ::? idx) = ixDrop sh idx +ixDrop sh ZIX = sh +ixDrop (_ :.@ sh) (_ :.@ idx) = ixDrop sh idx +ixDrop (_ :.? sh) (_ :.? idx) = ixDrop sh idx ssxAppend :: StaticShapeX sh -> StaticShapeX sh' -> StaticShapeX (sh ++ sh') -ssxAppend SZX sh' = sh' +ssxAppend ZSX sh' = sh' ssxAppend (n :$@ sh) sh' = n :$@ ssxAppend sh sh' ssxAppend (() :$? sh) sh' = () :$? ssxAppend sh sh' shapeSize :: IIxX sh -> Int -shapeSize IZX = 1 -shapeSize (n ::@ sh) = n * shapeSize sh -shapeSize (n ::? sh) = n * shapeSize sh +shapeSize ZIX = 1 +shapeSize (n :.@ sh) = n * shapeSize sh +shapeSize (n :.? sh) = n * shapeSize sh -- | This may fail if @sh@ has @Nothing@s in it. ssxToShape' :: StaticShapeX sh -> Maybe (IIxX sh) -ssxToShape' SZX = Just IZX -ssxToShape' (n :$@ sh) = (fromIntegral (fromSNat n) ::@) <$> ssxToShape' sh +ssxToShape' ZSX = Just ZIX +ssxToShape' (n :$@ sh) = (fromIntegral (fromSNat n) :.@) <$> ssxToShape' sh ssxToShape' (_ :$? _) = Nothing fromLinearIdx :: IIxX sh -> Int -> IIxX sh @@ -130,26 +130,26 @@ fromLinearIdx = \sh i -> case go sh i of where -- returns (index in subarray, remaining index in enclosing array) go :: IIxX sh -> Int -> (IIxX sh, Int) - go IZX i = (IZX, i) - go (n ::@ sh) i = + go ZIX i = (ZIX, i) + go (n :.@ sh) i = let (idx, i') = go sh i (upi, locali) = i' `quotRem` n - in (locali ::@ idx, upi) - go (n ::? sh) i = + in (locali :.@ idx, upi) + go (n :.? sh) i = let (idx, i') = go sh i (upi, locali) = i' `quotRem` n - in (locali ::? idx, upi) + in (locali :.? idx, upi) toLinearIdx :: IIxX sh -> IIxX sh -> Int toLinearIdx = \sh i -> fst (go sh i) where -- returns (index in subarray, size of subarray) go :: IIxX sh -> IIxX sh -> (Int, Int) - go IZX IZX = (0, 1) - go (n ::@ sh) (i ::@ ix) = + go ZIX ZIX = (0, 1) + go (n :.@ sh) (i :.@ ix) = let (lidx, sz) = go sh ix in (sz * i + lidx, n * sz) - go (n ::? sh) (i ::? ix) = + go (n :.? sh) (i :.? ix) = let (lidx, sz) = go sh ix in (sz * i + lidx, n * sz) @@ -157,22 +157,22 @@ enumShape :: IIxX sh -> [IIxX sh] enumShape = \sh -> go sh id [] where go :: IIxX sh -> (IIxX sh -> a) -> [a] -> [a] - go IZX f = (f IZX :) - go (n ::@ sh) f = foldr (.) id [go sh (f . (i ::@)) | i <- [0 .. n-1]] - go (n ::? sh) f = foldr (.) id [go sh (f . (i ::?)) | i <- [0 .. n-1]] + go ZIX f = (f ZIX :) + go (n :.@ sh) f = foldr (.) id [go sh (f . (i :.@)) | i <- [0 .. n-1]] + go (n :.? sh) f = foldr (.) id [go sh (f . (i :.?)) | i <- [0 .. n-1]] shapeLshape :: IIxX sh -> S.ShapeL -shapeLshape IZX = [] -shapeLshape (n ::@ sh) = n : shapeLshape sh -shapeLshape (n ::? sh) = n : shapeLshape sh +shapeLshape ZIX = [] +shapeLshape (n :.@ sh) = n : shapeLshape sh +shapeLshape (n :.? sh) = n : shapeLshape sh ssxLength :: StaticShapeX sh -> Int -ssxLength SZX = 0 +ssxLength ZSX = 0 ssxLength (_ :$@ ssh) = 1 + ssxLength ssh ssxLength (_ :$? ssh) = 1 + ssxLength ssh ssxIotaFrom :: Int -> StaticShapeX sh -> [Int] -ssxIotaFrom _ SZX = [] +ssxIotaFrom _ ZSX = [] ssxIotaFrom i (_ :$@ ssh) = i : ssxIotaFrom (i+1) ssh ssxIotaFrom i (_ :$? ssh) = i : ssxIotaFrom (i+1) ssh @@ -185,22 +185,22 @@ lemRankAppComm :: StaticShapeX sh1 -> StaticShapeX sh2 lemRankAppComm _ _ = unsafeCoerce Refl -- TODO improve this lemKnownINatRank :: IIxX sh -> Dict KnownINat (Rank sh) -lemKnownINatRank IZX = Dict -lemKnownINatRank (_ ::@ sh) | Dict <- lemKnownINatRank sh = Dict -lemKnownINatRank (_ ::? sh) | Dict <- lemKnownINatRank sh = Dict +lemKnownINatRank ZIX = Dict +lemKnownINatRank (_ :.@ sh) | Dict <- lemKnownINatRank sh = Dict +lemKnownINatRank (_ :.? sh) | Dict <- lemKnownINatRank sh = Dict lemKnownINatRankSSX :: StaticShapeX sh -> Dict KnownINat (Rank sh) -lemKnownINatRankSSX SZX = Dict +lemKnownINatRankSSX ZSX = Dict lemKnownINatRankSSX (_ :$@ ssh) | Dict <- lemKnownINatRankSSX ssh = Dict lemKnownINatRankSSX (_ :$? ssh) | Dict <- lemKnownINatRankSSX ssh = Dict lemKnownShapeX :: StaticShapeX sh -> Dict KnownShapeX sh -lemKnownShapeX SZX = Dict +lemKnownShapeX ZSX = Dict lemKnownShapeX (GHC_SNat :$@ ssh) | Dict <- lemKnownShapeX ssh = Dict lemKnownShapeX (() :$? ssh) | Dict <- lemKnownShapeX ssh = Dict lemAppKnownShapeX :: StaticShapeX sh1 -> StaticShapeX sh2 -> Dict KnownShapeX (sh1 ++ sh2) -lemAppKnownShapeX SZX ssh' = lemKnownShapeX ssh' +lemAppKnownShapeX ZSX ssh' = lemKnownShapeX ssh' lemAppKnownShapeX (GHC_SNat :$@ ssh) ssh' | Dict <- lemAppKnownShapeX ssh ssh' = Dict @@ -212,9 +212,9 @@ shape :: forall sh a. KnownShapeX sh => XArray sh a -> IIxX sh shape (XArray arr) = go (knownShapeX @sh) (S.shapeL arr) where go :: StaticShapeX sh' -> [Int] -> IIxX sh' - go SZX [] = IZX - go (n :$@ ssh) (_ : l) = fromIntegral (fromSNat n) ::@ go ssh l - go (() :$? ssh) (n : l) = n ::? go ssh l + go ZSX [] = ZIX + go (n :$@ ssh) (_ : l) = fromIntegral (fromSNat n) :.@ go ssh l + go (() :$? ssh) (n : l) = n :.? go ssh l go _ _ = error "Invalid shapeL" fromVector :: forall sh a. Storable a => IIxX sh -> VS.Vector a -> XArray sh a @@ -247,9 +247,9 @@ generate sh f = fromVector sh $ VS.generate (shapeSize sh) (f . fromLinearIdx sh -- <$> VS.generateM (shapeSize sh) (f . fromLinearIdx sh) indexPartial :: Storable a => XArray (sh ++ sh') a -> IIxX sh -> XArray sh' a -indexPartial (XArray arr) IZX = XArray arr -indexPartial (XArray arr) (i ::@ idx) = indexPartial (XArray (S.index arr i)) idx -indexPartial (XArray arr) (i ::? idx) = indexPartial (XArray (S.index arr i)) idx +indexPartial (XArray arr) ZIX = XArray arr +indexPartial (XArray arr) (i :.@ idx) = indexPartial (XArray (S.index arr i)) idx +indexPartial (XArray arr) (i :.? idx) = indexPartial (XArray (S.index arr i)) idx index :: forall sh a. Storable a => XArray sh a -> IIxX sh -> a index xarr i @@ -344,7 +344,7 @@ sumInner :: forall sh sh' a. (Storable a, Num a) => StaticShapeX sh -> StaticShapeX sh' -> XArray (sh ++ sh') a -> XArray sh a sumInner ssh ssh' | Refl <- lemAppNil @sh - = rerank ssh ssh' SZX (scalar . sumFull) + = rerank ssh ssh' ZSX (scalar . sumFull) sumOuter :: forall sh sh' a. (Storable a, Num a) => StaticShapeX sh -> StaticShapeX sh' -> XArray (sh ++ sh') a -> XArray sh' a diff --git a/src/Data/Array/Nested/Internal.hs b/src/Data/Array/Nested/Internal.hs index 9a87389..0582a14 100644 --- a/src/Data/Array/Nested/Internal.hs +++ b/src/Data/Array/Nested/Internal.hs @@ -59,7 +59,7 @@ import Data.INat -- array. For example, for this array: -- -- arr :: Ranked I3 (Ranked I2 Int, Ranked I1 Float) --- rshape arr == 0 ::: 0 ::: 0 ::: IZR +-- rshape arr == 0 :.: 0 :.: 0 :.: ZIR -- -- the two underlying XArrays have a shape, and those shapes might be anything. -- The invariant is that these element shapes are unobservable in the API. @@ -87,7 +87,7 @@ lemKnownReplicate :: forall n. KnownINat n => Proxy n -> Dict KnownShapeX (Repli lemKnownReplicate _ = X.lemKnownShapeX (go (inatSing @n)) where go :: SINat m -> StaticShapeX (Replicate m Nothing) - go SZ = SZX + go SZ = ZSX go (SS n) = () :$? go n lemRankReplicate :: forall n. KnownINat n => Proxy n -> X.Rank (Replicate n (Nothing @Nat)) :~: n @@ -106,9 +106,9 @@ lemReplicatePlusApp _ _ _ = go (inatSing @n) go (SS n) | Refl <- go n = Refl ixAppSplit :: Proxy sh' -> StaticShapeX sh -> IIxX (sh ++ sh') -> (IIxX sh, IIxX sh') -ixAppSplit _ SZX idx = (IZX, idx) -ixAppSplit p (_ :$@ ssh) (i ::@ idx) = first (i ::@) (ixAppSplit p ssh idx) -ixAppSplit p (_ :$? ssh) (i ::? idx) = first (i ::?) (ixAppSplit p ssh idx) +ixAppSplit _ ZSX idx = (ZIX, idx) +ixAppSplit p (_ :$@ ssh) (i :.@ idx) = first (i :.@) (ixAppSplit p ssh idx) +ixAppSplit p (_ :$? ssh) (i :.? idx) = first (i :.?) (ixAppSplit p ssh idx) -- | Wrapper type used as a tag to attach instances on. The instances on arrays @@ -429,9 +429,9 @@ instance (Elt a, KnownShapeX sh') => Elt (Mixed sh' a) where -- | Check whether a given shape corresponds on the statically-known components of the shape. checkBounds :: IIxX sh' -> StaticShapeX sh' -> Bool -checkBounds IZX SZX = True -checkBounds (n ::@ sh') (n' :$@ ssh') = n == fromIntegral (fromSNat n') && checkBounds sh' ssh' -checkBounds (_ ::? sh') (() :$? ssh') = checkBounds sh' ssh' +checkBounds ZIX ZSX = True +checkBounds (n :.@ sh') (n' :$@ ssh') = n == fromIntegral (fromSNat n') && checkBounds sh' ssh' +checkBounds (_ :.? sh') (() :$? ssh') = checkBounds sh' ssh' -- | Create an array given a size and a function that computes the element at a -- given index. @@ -441,8 +441,8 @@ checkBounds (_ ::? sh') (() :$? ssh') = checkBounds sh' ssh' -- runtime error: -- -- > foo :: Mixed [Nothing] (Mixed [Nothing] Double) --- > foo = mgenerate (10 ::: IZR) $ \(i ::: IZR) -> --- > mgenerate (i ::: IZR) $ \(j ::: IZR) -> +-- > foo = mgenerate (10 :.: ZIR) $ \(i :.: ZIR) -> +-- > mgenerate (i :.: ZIR) $ \(j :.: ZIR) -> -- > ... -- -- because the size of the inner 'mgenerate' is not always the same (it depends @@ -502,7 +502,7 @@ mtoList1 :: Elt a => Mixed '[n] a -> [a] mtoList1 = map munScalar . mtoList munScalar :: Elt a => Mixed '[] a -> a -munScalar arr = mindex arr IZX +munScalar arr = mindex arr ZIX mconstantP :: forall sh a. (KnownShapeX sh, Storable a) => IIxX sh -> a -> Mixed sh (Primitive a) mconstantP sh x @@ -708,7 +708,7 @@ lemKnownMapJust :: forall sh. KnownShape sh => Proxy sh -> Dict KnownShapeX (Map lemKnownMapJust _ = X.lemKnownShapeX (go (knownShape @sh)) where go :: SShape sh' -> StaticShapeX (MapJust sh') - go ShNil = SZX + go ShNil = ZSX go (ShCons n sh) = n :$@ go sh lemMapJustPlusApp :: forall sh1 sh2. KnownShape sh1 => Proxy sh1 -> Proxy sh2 @@ -852,34 +852,34 @@ deriving via Ranked n (Primitive Double) instance KnownINat n => Num (Ranked n D -- | An index into a rank-typed array. type IxR :: Type -> INat -> Type data IxR i n where - IZR :: IxR i Z - (:::) :: forall n i. i -> IxR i n -> IxR i (S n) + ZIR :: IxR i Z + (:.:) :: forall n i. i -> IxR i n -> IxR i (S n) deriving instance Show i => Show (IxR i n) deriving instance Eq i => Eq (IxR i n) -infixr 3 ::: +infixr 3 :.: type IIxR = IxR Int zeroIxR :: SINat n -> IIxR n -zeroIxR SZ = IZR -zeroIxR (SS n) = 0 ::: zeroIxR n +zeroIxR SZ = ZIR +zeroIxR (SS n) = 0 :.: zeroIxR n ixCvtXR :: IIxX sh -> IIxR (X.Rank sh) -ixCvtXR IZX = IZR -ixCvtXR (n ::@ idx) = n ::: ixCvtXR idx -ixCvtXR (n ::? idx) = n ::: ixCvtXR idx +ixCvtXR ZIX = ZIR +ixCvtXR (n :.@ idx) = n :.: ixCvtXR idx +ixCvtXR (n :.? idx) = n :.: ixCvtXR idx ixCvtRX :: IIxR n -> IIxX (Replicate n Nothing) -ixCvtRX IZR = IZX -ixCvtRX (n ::: idx) = n ::? ixCvtRX idx +ixCvtRX ZIR = ZIX +ixCvtRX (n :.: idx) = n :.? ixCvtRX idx knownIxR :: IIxR n -> Dict KnownINat n -knownIxR IZR = Dict -knownIxR (_ ::: idx) | Dict <- knownIxR idx = Dict +knownIxR ZIR = Dict +knownIxR (_ :.: idx) | Dict <- knownIxR idx = Dict shapeSizeR :: IIxR n -> Int -shapeSizeR IZR = 1 -shapeSizeR (n ::: sh) = n * shapeSizeR sh +shapeSizeR ZIR = 1 +shapeSizeR (n :.: sh) = n * shapeSizeR sh rshape :: forall n a. (KnownINat n, Elt a) => Ranked n a -> IIxR n @@ -921,7 +921,7 @@ rsumOuter1 (Ranked arr) | Dict <- lemKnownReplicate (Proxy @n) = Ranked . coerce @(XArray (Replicate n 'Nothing) a) @(Mixed (Replicate n 'Nothing) (Primitive a)) - . X.sumOuter (() :$? SZX) (knownShapeX @(Replicate n Nothing)) + . X.sumOuter (() :$? ZSX) (knownShapeX @(Replicate n Nothing)) . coerce @(Mixed (Replicate (S n) Nothing) (Primitive a)) @(XArray (Replicate (S n) Nothing) a) $ arr @@ -957,7 +957,7 @@ rtoList1 :: Elt a => Ranked I1 a -> [a] rtoList1 = map runScalar . rtoList runScalar :: Elt a => Ranked I0 a -> a -runScalar arr = rindex arr IZR +runScalar arr = rindex arr ZIR rconstantP :: forall n a. (KnownINat n, Storable a) => IIxR n -> a -> Ranked n (Primitive a) rconstantP sh x @@ -1004,33 +1004,33 @@ deriving via Shaped sh (Primitive Double) instance KnownShape sh => Num (Shaped -- from a 'KnownShape' dictionary. type IxS :: Type -> [Nat] -> Type data IxS i sh where - IZS :: IxS i '[] - (::$) :: forall n sh i. i -> IxS i sh -> IxS i (n : sh) + ZIS :: IxS i '[] + (:.$) :: forall n sh i. i -> IxS i sh -> IxS i (n : sh) deriving instance Show i => Show (IxS i n) deriving instance Eq i => Eq (IxS i n) -infixr 3 ::$ +infixr 3 :.$ type IIxS = IxS Int zeroIxS :: SShape sh -> IIxS sh -zeroIxS ShNil = IZS -zeroIxS (ShCons _ sh) = 0 ::$ zeroIxS sh +zeroIxS ShNil = ZIS +zeroIxS (ShCons _ sh) = 0 :.$ zeroIxS sh cvtSShapeIxS :: SShape sh -> IIxS sh -cvtSShapeIxS ShNil = IZS -cvtSShapeIxS (ShCons n sh) = fromIntegral (fromSNat n) ::$ cvtSShapeIxS sh +cvtSShapeIxS ShNil = ZIS +cvtSShapeIxS (ShCons n sh) = fromIntegral (fromSNat n) :.$ cvtSShapeIxS sh ixCvtXS :: SShape sh -> IIxX (MapJust sh) -> IIxS sh -ixCvtXS ShNil IZX = IZS -ixCvtXS (ShCons _ sh) (n ::@ idx) = n ::$ ixCvtXS sh idx +ixCvtXS ShNil ZIX = ZIS +ixCvtXS (ShCons _ sh) (n :.@ idx) = n :.$ ixCvtXS sh idx ixCvtSX :: IIxS sh -> IIxX (MapJust sh) -ixCvtSX IZS = IZX -ixCvtSX (n ::$ sh) = n ::@ ixCvtSX sh +ixCvtSX ZIS = ZIX +ixCvtSX (n :.$ sh) = n :.@ ixCvtSX sh shapeSizeS :: IIxS sh -> Int -shapeSizeS IZS = 1 -shapeSizeS (n ::$ sh) = n * shapeSizeS sh +shapeSizeS ZIS = 1 +shapeSizeS (n :.$ sh) = n * shapeSizeS sh -- | This does not touch the passed array, all information comes from 'KnownShape'. @@ -1068,7 +1068,7 @@ ssumOuter1 (Shaped arr) | Dict <- lemKnownMapJust (Proxy @sh) = Shaped . coerce @(XArray (MapJust sh) a) @(Mixed (MapJust sh) (Primitive a)) - . X.sumOuter (natSing @n :$@ SZX) (knownShapeX @(MapJust sh)) + . X.sumOuter (natSing @n :$@ ZSX) (knownShapeX @(MapJust sh)) . coerce @(Mixed (Just n : MapJust sh) (Primitive a)) @(XArray (Just n : MapJust sh) a) $ arr @@ -1105,7 +1105,7 @@ stoList1 :: Elt a => Shaped '[n] a -> [a] stoList1 = map sunScalar . stoList sunScalar :: Elt a => Shaped '[] a -> a -sunScalar arr = sindex arr IZS +sunScalar arr = sindex arr ZIS sconstantP :: forall sh a. (KnownShape sh, Storable a) => a -> Shaped sh (Primitive a) sconstantP x |