diff options
Diffstat (limited to 'src/Data/Array/Mixed.hs')
| -rw-r--r-- | src/Data/Array/Mixed.hs | 306 |
1 files changed, 174 insertions, 132 deletions
diff --git a/src/Data/Array/Mixed.hs b/src/Data/Array/Mixed.hs index 69c44ab..df506d6 100644 --- a/src/Data/Array/Mixed.hs +++ b/src/Data/Array/Mixed.hs @@ -8,6 +8,7 @@ {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE PolyKinds #-} {-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE QuantifiedConstraints #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE StandaloneDeriving #-} @@ -24,6 +25,7 @@ module Data.Array.Mixed where import qualified Data.Array.RankedS as S import qualified Data.Array.Ranked as ORB import Data.Coerce +import Data.Functor.Const import Data.Kind import Data.Proxy import Data.Type.Bool @@ -32,6 +34,7 @@ import qualified Data.Vector.Storable as VS import Foreign.Storable (Storable) import GHC.TypeError import GHC.TypeLits +import qualified GHC.TypeNats as TypeNats import Unsafe.Coerce (unsafeCoerce) @@ -87,45 +90,115 @@ type family Replicate n a where Replicate n a = a : Replicate (n - 1) a +type ListX :: [Maybe Nat] -> (Maybe Nat -> Type) -> Type +data ListX sh f where + ZX :: ListX '[] f + (::%) :: f n -> ListX sh f -> ListX (n : sh) f +deriving instance (forall n. Show (f n)) => Show (ListX sh f) +deriving instance (forall n. Eq (f n)) => Eq (ListX sh f) +deriving instance (forall n. Ord (f n)) => Ord (ListX sh f) +infixr 3 ::% + +data UnconsListXRes f sh1 = + forall n sh. (n : sh ~ sh1) => UnconsListXRes (ListX sh f) (f n) +unconsListX :: ListX sh1 f -> Maybe (UnconsListXRes f sh1) +unconsListX (i ::% shl') = Just (UnconsListXRes shl' i) +unconsListX ZX = Nothing + +fmapListX :: (forall n. f n -> g n) -> ListX sh f -> ListX sh g +fmapListX _ ZX = ZX +fmapListX f (x ::% xs) = f x ::% fmapListX f xs + +foldListX :: Monoid m => (forall n. f n -> m) -> ListX sh f -> m +foldListX _ ZX = mempty +foldListX f (x ::% xs) = f x <> foldListX f xs + + type IxX :: [Maybe Nat] -> Type -> Type -data IxX sh i where - ZIX :: IxX '[] i - (:.@) :: forall n sh i. i -> IxX sh i -> IxX (Just n : sh) i - (:.?) :: forall sh i. i -> IxX sh i -> IxX (Nothing : sh) i -deriving instance Show i => Show (IxX sh i) -deriving instance Eq i => Eq (IxX sh i) -deriving instance Ord i => Ord (IxX sh i) -deriving instance Functor (IxX sh) -deriving instance Foldable (IxX sh) -infixr 3 :.@ -infixr 3 :.? +newtype IxX sh i = IxX (ListX sh (Const i)) + deriving (Show, Eq, Ord) + +pattern ZIX :: forall sh i. () => sh ~ '[] => IxX sh i +pattern ZIX = IxX ZX + +pattern (:.%) + :: forall {sh1} {i}. + forall n sh. (n : sh ~ sh1) + => i -> IxX sh i -> IxX sh1 i +pattern i :.% shl <- IxX (unconsListX -> Just (UnconsListXRes (IxX -> shl) (getConst -> i))) + where i :.% IxX shl = IxX (Const i ::% shl) +infixr 3 :.% + +{-# COMPLETE ZIX, (:.%) #-} type IIxX sh = IxX sh Int +instance Functor (IxX sh) where + fmap f (IxX l) = IxX (fmapListX (Const . f . getConst) l) + +instance Foldable (IxX sh) where + foldMap f (IxX l) = foldListX (f . getConst) l + + +data SMayNat i f n where + SUnknown :: i -> SMayNat i f Nothing + SKnown :: f n -> SMayNat i f (Just n) +deriving instance (Show i, forall m. Show (f m)) => Show (SMayNat i f n) +deriving instance (Eq i, forall m. Eq (f m)) => Eq (SMayNat i f n) +deriving instance (Ord i, forall m. Ord (f m)) => Ord (SMayNat i f n) + +fromSMayNat :: (n ~ Nothing => i -> r) -> (forall m. n ~ Just m => f m -> r) -> SMayNat i f n -> r +fromSMayNat f _ (SUnknown i) = f i +fromSMayNat _ g (SKnown s) = g s + +fromSMayNat' :: SMayNat Int SNat n -> Int +fromSMayNat' = fromSMayNat id fromSNat' + type ShX :: [Maybe Nat] -> Type -> Type -data ShX sh i where - ZSX :: ShX '[] i - (:$@) :: forall n sh i. SNat n -> ShX sh i -> ShX (Just n : sh) i - (:$?) :: forall sh i. i -> ShX sh i -> ShX (Nothing : sh) i -deriving instance Show i => Show (ShX sh i) -deriving instance Eq i => Eq (ShX sh i) -deriving instance Ord i => Ord (ShX sh i) -deriving instance Functor (ShX sh) -deriving instance Foldable (ShX sh) -infixr 3 :$@ -infixr 3 :$? +newtype ShX sh i = ShX (ListX sh (SMayNat i SNat)) + deriving (Show, Eq, Ord) + +pattern ZSX :: forall sh i. () => sh ~ '[] => ShX sh i +pattern ZSX = ShX ZX + +pattern (:$%) + :: forall {sh1} {i}. + forall n sh. (n : sh ~ sh1) + => SMayNat i SNat n -> ShX sh i -> ShX sh1 i +pattern i :$% shl <- ShX (unconsListX -> Just (UnconsListXRes (ShX -> shl) i)) + where i :$% ShX shl = ShX (i ::% shl) +infixr 3 :$% + +{-# COMPLETE ZSX, (:$%) #-} type IShX sh = ShX sh Int +instance Functor (ShX sh) where + fmap f (ShX l) = ShX (fmapListX (fromSMayNat (SUnknown . f) SKnown) l) + +lengthShX :: ShX sh i -> Int +lengthShX ZSX = 0 +lengthShX (_ :$% sh) = 1 + lengthShX sh + + -- | The part of a shape that is statically known. type StaticShX :: [Maybe Nat] -> Type -data StaticShX sh where - ZKSX :: StaticShX '[] - (:!$@) :: SNat n -> StaticShX sh -> StaticShX (Just n : sh) - (:!$?) :: () -> StaticShX sh -> StaticShX (Nothing : sh) -deriving instance Show (StaticShX sh) -infixr 3 :!$@ -infixr 3 :!$? +newtype StaticShX sh = StaticShX (ListX sh (SMayNat () SNat)) + deriving (Show, Eq, Ord) + +pattern ZKX :: forall sh. () => sh ~ '[] => StaticShX sh +pattern ZKX = StaticShX ZX + +pattern (:!%) + :: forall {sh1}. + forall n sh. (n : sh ~ sh1) + => SMayNat () SNat n -> StaticShX sh -> StaticShX sh1 +pattern i :!% shl <- StaticShX (unconsListX -> Just (UnconsListXRes (StaticShX -> shl) i)) + where i :!% StaticShX shl = StaticShX (i ::% shl) +infixr 3 :!% + +{-# COMPLETE ZKX, (:!%) #-} + type family Rank sh where Rank '[] = 0 @@ -136,70 +209,68 @@ newtype XArray sh a = XArray (S.Array (Rank sh) a) deriving (Show) zeroIxX :: StaticShX sh -> IIxX sh -zeroIxX ZKSX = ZIX -zeroIxX (_ :!$@ ssh) = 0 :.@ zeroIxX ssh -zeroIxX (_ :!$? ssh) = 0 :.? zeroIxX ssh +zeroIxX ZKX = ZIX +zeroIxX (_ :!% ssh) = 0 :.% zeroIxX ssh zeroIxX' :: IShX sh -> IIxX sh zeroIxX' ZSX = ZIX -zeroIxX' (_ :$@ sh) = 0 :.@ zeroIxX' sh -zeroIxX' (_ :$? sh) = 0 :.? zeroIxX' sh +zeroIxX' (_ :$% sh) = 0 :.% zeroIxX' sh -- This is a weird operation, so it has a long name completeShXzeros :: StaticShX sh -> IShX sh -completeShXzeros ZKSX = ZSX -completeShXzeros (n :!$@ ssh) = n :$@ completeShXzeros ssh -completeShXzeros (_ :!$? ssh) = 0 :$? completeShXzeros ssh +completeShXzeros ZKX = ZSX +completeShXzeros (SUnknown () :!% ssh) = SUnknown 0 :$% completeShXzeros ssh +completeShXzeros (SKnown n :!% ssh) = SKnown n :$% completeShXzeros ssh -- TODO: generalise all these things to arbitrary @i@ ixAppend :: IIxX sh -> IIxX sh' -> IIxX (sh ++ sh') ixAppend ZIX idx' = idx' -ixAppend (i :.@ idx) idx' = i :.@ ixAppend idx idx' -ixAppend (i :.? idx) idx' = i :.? ixAppend idx idx' +ixAppend (i :.% idx) idx' = i :.% ixAppend idx idx' shAppend :: IShX sh -> IShX sh' -> IShX (sh ++ sh') shAppend ZSX sh' = sh' -shAppend (n :$@ sh) sh' = n :$@ shAppend sh sh' -shAppend (n :$? sh) sh' = n :$? shAppend sh sh' +shAppend (n :$% sh) sh' = n :$% shAppend sh sh' ixDrop :: IIxX (sh ++ sh') -> IIxX sh -> IIxX sh' -ixDrop sh ZIX = sh -ixDrop (_ :.@ sh) (_ :.@ idx) = ixDrop sh idx -ixDrop (_ :.? sh) (_ :.? idx) = ixDrop sh idx +ixDrop long ZIX = long +ixDrop long (_ :.% short) = case long of _ :.% long' -> ixDrop long' short shDropIx :: IShX (sh ++ sh') -> IIxX sh -> IShX sh' shDropIx sh ZIX = sh -shDropIx (_ :$@ sh) (_ :.@ idx) = shDropIx sh idx -shDropIx (_ :$? sh) (_ :.? idx) = shDropIx sh idx +shDropIx sh (_ :.% idx) = case sh of _ :$% sh' -> shDropIx sh' idx + +ssxDropIx :: StaticShX (sh ++ sh') -> IIxX sh -> StaticShX sh' +ssxDropIx ssh ZIX = ssh +ssxDropIx ssh (_ :.% idx) = case ssh of _ :!% ssh' -> ssxDropIx ssh' idx shTail :: IShX (n : sh) -> IShX sh -shTail (_ :$@ sh) = sh -shTail (_ :$? sh) = sh +shTail (_ :$% sh) = sh + +ssxTail :: StaticShX (n : sh) -> StaticShX sh +ssxTail (_ :!% ssh) = ssh ssxAppend :: StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh') -ssxAppend ZKSX sh' = sh' -ssxAppend (n :!$@ sh) sh' = n :!$@ ssxAppend sh sh' -ssxAppend (() :!$? sh) sh' = () :!$? ssxAppend sh sh' +ssxAppend ZKX sh' = sh' +ssxAppend (n :!% sh) sh' = n :!% ssxAppend sh sh' shapeSize :: IShX sh -> Int shapeSize ZSX = 1 -shapeSize (n :$@ sh) = fromSNat' n * shapeSize sh -shapeSize (n :$? sh) = n * shapeSize sh +shapeSize (n :$% sh) = fromSMayNat' n * shapeSize sh -- | This may fail if @sh@ has @Nothing@s in it. ssxToShape' :: StaticShX sh -> Maybe (IShX sh) -ssxToShape' ZKSX = Just ZSX -ssxToShape' (n :!$@ sh) = (n :$@) <$> ssxToShape' sh -ssxToShape' (_ :!$? _) = Nothing +ssxToShape' ZKX = Just ZSX +ssxToShape' (SKnown n :!% sh) = (SKnown n :$%) <$> ssxToShape' sh +ssxToShape' (SUnknown _ :!% _) = Nothing lemReplicateSucc :: (a : Replicate n a) :~: Replicate (n + 1) a lemReplicateSucc = unsafeCoerce Refl ssxReplicate :: SNat n -> StaticShX (Replicate n Nothing) -ssxReplicate SZ = ZKSX +ssxReplicate SZ = ZKX ssxReplicate (SS (n :: SNat n')) | Refl <- lemReplicateSucc @(Nothing @Nat) @n' - = () :!$? ssxReplicate n + = SUnknown () :!% ssxReplicate n fromLinearIdx :: IShX sh -> Int -> IIxX sh fromLinearIdx = \sh i -> case go sh i of @@ -210,14 +281,10 @@ fromLinearIdx = \sh i -> case go sh i of -- returns (index in subarray, remaining index in enclosing array) go :: IShX sh -> Int -> (IIxX sh, Int) go ZSX i = (ZIX, i) - go (n :$@ sh) i = + go (n :$% sh) i = let (idx, i') = go sh i - (upi, locali) = i' `quotRem` fromSNat' n - in (locali :.@ idx, upi) - go (n :$? sh) i = - let (idx, i') = go sh i - (upi, locali) = i' `quotRem` n - in (locali :.? idx, upi) + (upi, locali) = i' `quotRem` fromSMayNat' n + in (locali :.% idx, upi) toLinearIdx :: IShX sh -> IIxX sh -> Int toLinearIdx = \sh i -> fst (go sh i) @@ -225,40 +292,32 @@ toLinearIdx = \sh i -> fst (go sh i) -- returns (index in subarray, size of subarray) go :: IShX sh -> IIxX sh -> (Int, Int) go ZSX ZIX = (0, 1) - go (n :$@ sh) (i :.@ ix) = - let (lidx, sz) = go sh ix - in (sz * i + lidx, fromSNat' n * sz) - go (n :$? sh) (i :.? ix) = + go (n :$% sh) (i :.% ix) = let (lidx, sz) = go sh ix - in (sz * i + lidx, n * sz) + in (sz * i + lidx, fromSMayNat' n * sz) enumShape :: IShX sh -> [IIxX sh] enumShape = \sh -> go sh id [] where go :: IShX sh -> (IIxX sh -> a) -> [a] -> [a] go ZSX f = (f ZIX :) - go (n :$@ sh) f = foldr (.) id [go sh (f . (i :.@)) | i <- [0 .. fromSNat' n - 1]] - 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 .. fromSMayNat' n - 1]] shapeLshape :: IShX sh -> S.ShapeL shapeLshape ZSX = [] -shapeLshape (n :$@ sh) = fromSNat' n : shapeLshape sh -shapeLshape (n :$? sh) = n : shapeLshape sh +shapeLshape (n :$% sh) = fromSMayNat' n : shapeLshape sh ssxLength :: StaticShX sh -> Int -ssxLength ZKSX = 0 -ssxLength (_ :!$@ ssh) = 1 + ssxLength ssh -ssxLength (_ :!$? ssh) = 1 + ssxLength ssh +ssxLength ZKX = 0 +ssxLength (_ :!% ssh) = 1 + ssxLength ssh ssxIotaFrom :: Int -> StaticShX sh -> [Int] -ssxIotaFrom _ ZKSX = [] -ssxIotaFrom i (_ :!$@ ssh) = i : ssxIotaFrom (i+1) ssh -ssxIotaFrom i (_ :!$? ssh) = i : ssxIotaFrom (i+1) ssh +ssxIotaFrom _ ZKX = [] +ssxIotaFrom i (_ :!% ssh) = i : ssxIotaFrom (i+1) ssh staticShapeFrom :: IShX sh -> StaticShX sh -staticShapeFrom ZSX = ZKSX -staticShapeFrom (n :$@ sh) = n :!$@ staticShapeFrom sh -staticShapeFrom (_ :$? sh) = () :!$? staticShapeFrom sh +staticShapeFrom ZSX = ZKX +staticShapeFrom (n :$% sh) = fromSMayNat (\_ -> SUnknown ()) SKnown n :!% staticShapeFrom sh lemRankApp :: StaticShX sh1 -> StaticShX sh2 -> Rank (sh1 ++ sh2) :~: Rank sh1 + Rank sh2 @@ -270,35 +329,18 @@ lemRankAppComm _ _ = unsafeCoerce Refl -- TODO improve this lemKnownNatRank :: IShX sh -> Dict KnownNat (Rank sh) lemKnownNatRank ZSX = Dict -lemKnownNatRank (_ :$@ sh) | Dict <- lemKnownNatRank sh = Dict -lemKnownNatRank (_ :$? sh) | Dict <- lemKnownNatRank sh = Dict +lemKnownNatRank (_ :$% sh) | Dict <- lemKnownNatRank sh = Dict lemKnownNatRankSSX :: StaticShX sh -> Dict KnownNat (Rank sh) -lemKnownNatRankSSX ZKSX = Dict -lemKnownNatRankSSX (_ :!$@ ssh) | Dict <- lemKnownNatRankSSX ssh = Dict -lemKnownNatRankSSX (_ :!$? ssh) | Dict <- lemKnownNatRankSSX ssh = Dict - --- lemKnownShapeX :: StaticShX sh -> Dict KnownShapeX sh --- lemKnownShapeX ZKSX = Dict --- lemKnownShapeX (GHC_SNat :!$@ ssh) | Dict <- lemKnownShapeX ssh = Dict --- lemKnownShapeX (() :!$? ssh) | Dict <- lemKnownShapeX ssh = Dict - --- lemAppKnownShapeX :: StaticShX sh1 -> StaticShX sh2 -> Dict KnownShapeX (sh1 ++ sh2) --- lemAppKnownShapeX ZKSX ssh' = lemKnownShapeX ssh' --- lemAppKnownShapeX (GHC_SNat :!$@ ssh) ssh' --- | Dict <- lemAppKnownShapeX ssh ssh' --- = Dict --- lemAppKnownShapeX (() :!$? ssh) ssh' --- | Dict <- lemAppKnownShapeX ssh ssh' --- = Dict +lemKnownNatRankSSX ZKX = Dict +lemKnownNatRankSSX (_ :!% ssh) | Dict <- lemKnownNatRankSSX ssh = Dict shape :: forall sh a. StaticShX sh -> XArray sh a -> IShX sh shape = \ssh (XArray arr) -> go ssh (S.shapeL arr) where go :: StaticShX sh' -> [Int] -> IShX sh' - go ZKSX [] = ZSX - go (n :!$@ ssh) (_ : l) = n :$@ go ssh l - go (() :!$? ssh) (n : l) = n :$? go ssh l + go ZKX [] = ZSX + go (n :!% ssh) (i : l) = fromSMayNat (\_ -> SUnknown i) SKnown n :$% go ssh l go _ _ = error "Invalid shapeL" fromVector :: forall sh a. Storable a => IShX sh -> VS.Vector a -> XArray sh a @@ -330,8 +372,7 @@ generate sh f = fromVector sh $ VS.generate (shapeSize sh) (f . fromLinearIdx sh indexPartial :: Storable a => XArray (sh ++ sh') a -> IIxX sh -> XArray sh' a 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 +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,6 +385,15 @@ type family AddMaybe n m where AddMaybe (Just _) Nothing = Nothing AddMaybe (Just n) (Just m) = Just (n + m) +-- This should be a function in base +plusSNat :: SNat n -> SNat m -> SNat (n + m) +plusSNat n m = TypeNats.withSomeSNat (TypeNats.fromSNat n + TypeNats.fromSNat m) unsafeCoerce + +smnAddMaybe :: SMayNat Int SNat n -> SMayNat Int SNat m -> SMayNat Int SNat (AddMaybe n m) +smnAddMaybe (SUnknown n) m = SUnknown (n + fromSMayNat' m) +smnAddMaybe (SKnown n) (SUnknown m) = SUnknown (fromSNat' n + m) +smnAddMaybe (SKnown n) (SKnown m) = SKnown (plusSNat n m) + append :: forall n m sh a. Storable a => StaticShX sh -> XArray (n : sh) a -> XArray (m : sh) a -> XArray (AddMaybe n m : sh) a append ssh (XArray a) (XArray b) @@ -442,42 +492,34 @@ lemRankPermute p (_ `HCons` is) | Refl <- lemRankPermute p is = Refl lemRankDropLen :: forall is sh. (Rank is <= Rank sh) => StaticShX sh -> HList SNat is -> Rank (DropLen is sh) :~: Rank sh - Rank is -lemRankDropLen ZKSX HNil = Refl -lemRankDropLen (_ :!$@ sh) (_ `HCons` is) | Refl <- lemRankDropLen sh is = Refl -lemRankDropLen (_ :!$? sh) (_ `HCons` is) | Refl <- lemRankDropLen sh is = Refl -lemRankDropLen (_ :!$@ _) HNil = Refl -lemRankDropLen (_ :!$? _) HNil = Refl -lemRankDropLen ZKSX (_ `HCons` _) = error "1 <= 0" +lemRankDropLen ZKX HNil = Refl +lemRankDropLen (_ :!% sh) (_ `HCons` is) | Refl <- lemRankDropLen sh is = Refl +lemRankDropLen (_ :!% _) HNil = Refl +lemRankDropLen ZKX (_ `HCons` _) = error "1 <= 0" lemIndexSucc :: Proxy i -> Proxy a -> Proxy l -> Index (i + 1) (a : l) :~: Index i l lemIndexSucc _ _ _ = unsafeCoerce Refl ssxTakeLen :: HList SNat is -> StaticShX sh -> StaticShX (TakeLen is sh) -ssxTakeLen HNil _ = ZKSX -ssxTakeLen (_ `HCons` is) (n :!$@ sh) = n :!$@ ssxTakeLen is sh -ssxTakeLen (_ `HCons` is) (n :!$? sh) = n :!$? ssxTakeLen is sh -ssxTakeLen (_ `HCons` _) ZKSX = error "Permutation longer than shape" +ssxTakeLen HNil _ = ZKX +ssxTakeLen (_ `HCons` is) (n :!% sh) = n :!% ssxTakeLen is sh +ssxTakeLen (_ `HCons` _) ZKX = error "Permutation longer than shape" ssxDropLen :: HList SNat is -> StaticShX sh -> StaticShX (DropLen is sh) ssxDropLen HNil sh = sh -ssxDropLen (_ `HCons` is) (_ :!$@ sh) = ssxDropLen is sh -ssxDropLen (_ `HCons` is) (_ :!$? sh) = ssxDropLen is sh -ssxDropLen (_ `HCons` _) ZKSX = error "Permutation longer than shape" +ssxDropLen (_ `HCons` is) (_ :!% sh) = ssxDropLen is sh +ssxDropLen (_ `HCons` _) ZKX = error "Permutation longer than shape" ssxPermute :: HList SNat is -> StaticShX sh -> StaticShX (Permute is sh) -ssxPermute HNil _ = ZKSX +ssxPermute HNil _ = ZKX ssxPermute (i `HCons` (is :: HList SNat is')) (sh :: StaticShX sh) = ssxIndex (Proxy @is') (Proxy @sh) i sh (ssxPermute is sh) ssxIndex :: Proxy is -> Proxy shT -> SNat i -> StaticShX sh -> StaticShX (Permute is shT) -> StaticShX (Index i sh : Permute is shT) -ssxIndex _ _ SZ (n :!$@ _) rest = n :!$@ rest -ssxIndex _ _ SZ (n :!$? _) rest = n :!$? rest -ssxIndex p pT (SS (i :: SNat i')) ((_ :: SNat n) :!$@ (sh :: StaticShX sh')) rest - | Refl <- lemIndexSucc (Proxy @i') (Proxy @(Just n)) (Proxy @sh') - = ssxIndex p pT i sh rest -ssxIndex p pT (SS (i :: SNat i')) (() :!$? (sh :: StaticShX sh')) rest - | Refl <- lemIndexSucc (Proxy @i') (Proxy @Nothing) (Proxy @sh') +ssxIndex _ _ SZ (n :!% _) rest = n :!% rest +ssxIndex p pT (SS (i :: SNat i')) ((_ :: SMayNat () SNat n) :!% (sh :: StaticShX sh')) rest + | Refl <- lemIndexSucc (Proxy @i') (Proxy @n) (Proxy @sh') = ssxIndex p pT i sh rest -ssxIndex _ _ _ ZKSX _ = error "Index into empty shape" +ssxIndex _ _ _ ZKX _ = error "Index into empty shape" -- | The list argument gives indices into the original dimension list. transpose :: forall is sh a. (Permutation is, Rank is <= Rank sh) @@ -526,7 +568,7 @@ sumInner :: forall sh sh' a. (Storable a, Num a) => StaticShX sh -> StaticShX sh' -> XArray (sh ++ sh') a -> XArray sh a sumInner ssh ssh' | Refl <- lemAppNil @sh - = rerank ssh ssh' ZKSX (scalar . sumFull) + = rerank ssh ssh' ZKX (scalar . sumFull) sumOuter :: forall sh sh' a. (Storable a, Num a) => StaticShX sh -> StaticShX sh' -> XArray (sh ++ sh') a -> XArray sh' a @@ -539,7 +581,7 @@ fromList1 :: forall n sh a. Storable a fromList1 ssh l | Dict <- lemKnownNatRankSSX ssh = case ssh of - m@GHC_SNat :!$@ _ | natVal m /= fromIntegral (length l) -> + SKnown m@GHC_SNat :!% _ | natVal m /= fromIntegral (length l) -> error $ "Data.Array.Mixed.fromList: length of list (" ++ show (length l) ++ ")" ++ "does not match the type (" ++ show (natVal m) ++ ")" _ -> XArray (S.ravel (ORB.fromList [length l] (coerce @[XArray sh a] @[S.Array (Rank sh) a] l))) @@ -575,4 +617,4 @@ reshapePartial :: forall sh1 sh2 sh' a. Storable a => StaticShX sh1 -> StaticShX reshapePartial ssh1 ssh' sh2 (XArray arr) | Dict <- lemKnownNatRankSSX (ssxAppend ssh1 ssh') , Dict <- lemKnownNatRankSSX (ssxAppend (staticShapeFrom sh2) ssh') - = XArray (S.reshape (shapeLshape sh2 ++ drop (length sh2) (S.shapeL arr)) arr) + = XArray (S.reshape (shapeLshape sh2 ++ drop (lengthShX sh2) (S.shapeL arr)) arr) |