Flip some index/shape-related functions

This ensures that the argument order consistently puts the main thing
being operated on at the end, and supporting singletons at the start.

1 files changed, 24 insertions, 27 deletions

diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs
index 8e0c274..2ee3600 100644
--- a/src/Data/Array/Nested/Mixed/Shape.hs
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -145,9 +145,9 @@ listxAppend :: ListX sh f -> ListX sh' f -> ListX (sh ++ sh') f
 listxAppend ZX idx' = idx'
 listxAppend (i ::% idx) idx' = i ::% listxAppend idx idx'
 
-listxDrop :: forall f g sh sh'. ListX (sh ++ sh') f -> ListX sh g -> ListX sh' f
-listxDrop long ZX = long
-listxDrop long (_ ::% short) = case long of _ ::% long' -> listxDrop long' short
+listxDrop :: forall f g sh sh'. ListX sh g -> ListX (sh ++ sh') f -> ListX sh' f
+listxDrop ZX long = long
+listxDrop (_ ::% short) long = case long of _ ::% long' -> listxDrop short long'
 
 listxInit :: forall f n sh. ListX (n : sh) f -> ListX (Init (n : sh)) f
 listxInit (i ::% sh@(_ ::% _)) = i ::% listxInit sh
@@ -235,7 +235,7 @@ ixxTail (IxX list) = IxX (listxTail list)
 ixxAppend :: forall sh sh' i. IxX sh i -> IxX sh' i -> IxX (sh ++ sh') i
 ixxAppend = coerce (listxAppend @_ @(Const i))
 
-ixxDrop :: forall sh sh' i. IxX (sh ++ sh') i -> IxX sh i -> IxX sh' i
+ixxDrop :: forall sh sh' i. IxX sh i -> IxX (sh ++ sh') i -> IxX sh' i
 ixxDrop = coerce (listxDrop @(Const i) @(Const i))
 
 ixxInit :: forall n sh i. IxX (n : sh) i -> IxX (Init (n : sh)) i
@@ -424,13 +424,13 @@ shxHead (ShX list) = listxHead list
 shxTail :: ShX (n : sh) i -> ShX sh i
 shxTail (ShX list) = ShX (listxTail list)
 
-shxDropSSX :: forall sh sh' i. ShX (sh ++ sh') i -> StaticShX sh -> ShX sh' i
+shxDropSSX :: forall sh sh' i. StaticShX sh -> ShX (sh ++ sh') i -> ShX sh' i
 shxDropSSX = coerce (listxDrop @(SMayNat i SNat) @(SMayNat () SNat))
 
-shxDropIx :: forall sh sh' i j. ShX (sh ++ sh') i -> IxX sh j -> ShX sh' i
+shxDropIx :: forall sh sh' i j. IxX sh j -> ShX (sh ++ sh') i -> ShX sh' i
 shxDropIx = coerce (listxDrop @(SMayNat i SNat) @(Const j))
 
-shxDropSh :: forall sh sh' i. ShX (sh ++ sh') i -> ShX sh i -> ShX sh' i
+shxDropSh :: forall sh sh' i. ShX sh i -> ShX (sh ++ sh') i -> ShX sh' i
 shxDropSh = coerce (listxDrop @(SMayNat i SNat) @(SMayNat i SNat))
 
 shxInit :: forall n sh i. ShX (n : sh) i -> ShX (Init (n : sh)) i
@@ -439,12 +439,9 @@ shxInit = coerce (listxInit @(SMayNat i SNat))
 shxLast :: forall n sh i. ShX (n : sh) i -> SMayNat i SNat (Last (n : sh))
 shxLast = coerce (listxLast @(SMayNat i SNat))
 
-shxTakeSSX :: forall sh sh' i proxy. proxy sh' -> ShX (sh ++ sh') i -> StaticShX sh -> ShX sh i
-shxTakeSSX _ = flip go
-  where
-    go :: StaticShX sh1 -> ShX (sh1 ++ sh') i -> ShX sh1 i
-    go ZKX _ = ZSX
-    go (_ :!% ssh1) (n :$% sh) = n :$% go ssh1 sh
+shxTakeSSX :: forall sh sh' i proxy. proxy sh' -> StaticShX sh -> ShX (sh ++ sh') i -> ShX sh i
+shxTakeSSX _ ZKX _ = ZSX
+shxTakeSSX p (_ :!% ssh1) (n :$% sh) = n :$% shxTakeSSX p ssh1 sh
 
 shxZipWith :: (forall n. SMayNat i SNat n -> SMayNat j SNat n -> SMayNat k SNat n)
            -> ShX sh i -> ShX sh j -> ShX sh k
@@ -468,17 +465,17 @@ shxEnum = \sh -> go sh id []
     go ZSX f = (f ZIX :)
     go (n :$% sh) f = foldr (.) id [go sh (f . (i :.%)) | i <- [0 .. fromSMayNat' n - 1]]
 
-shxCast :: IShX sh -> StaticShX sh' -> Maybe (IShX sh')
-shxCast ZSX ZKX = Just ZSX
-shxCast (SKnown n   :$% sh) (SKnown m    :!% ssh) | Just Refl <- testEquality n m = (SKnown n :$%) <$> shxCast sh ssh
-shxCast (SUnknown n :$% sh) (SKnown m    :!% ssh) | n == fromSNat' m              = (SKnown m :$%) <$> shxCast sh ssh
-shxCast (SKnown n   :$% sh) (SUnknown () :!% ssh)                                 = (SUnknown (fromSNat' n) :$%) <$> shxCast sh ssh
-shxCast (SUnknown n :$% sh) (SUnknown () :!% ssh)                                 = (SUnknown n :$%) <$> shxCast sh ssh
+shxCast :: StaticShX sh' -> IShX sh -> Maybe (IShX sh')
+shxCast ZKX ZSX = Just ZSX
+shxCast (SKnown m    :!% ssh) (SKnown n   :$% sh) | Just Refl <- testEquality n m = (SKnown n :$%) <$> shxCast ssh sh
+shxCast (SKnown m    :!% ssh) (SUnknown n :$% sh) | n == fromSNat' m              = (SKnown m :$%) <$> shxCast ssh sh
+shxCast (SUnknown () :!% ssh) (SKnown n   :$% sh)                                 = (SUnknown (fromSNat' n) :$%) <$> shxCast ssh sh
+shxCast (SUnknown () :!% ssh) (SUnknown n :$% sh)                                 = (SUnknown n :$%) <$> shxCast ssh sh
 shxCast _ _ = Nothing
 
 -- | Partial version of 'shxCast'.
-shxCast' :: IShX sh -> StaticShX sh' -> IShX sh'
-shxCast' sh ssh = case shxCast sh ssh of
+shxCast' :: StaticShX sh' -> IShX sh -> IShX sh'
+shxCast' ssh sh = case shxCast ssh sh of
   Just sh' -> sh'
   Nothing -> error $ "shxCast': Mismatch: (" ++ show sh ++ ") does not match (" ++ show ssh ++ ")"
 
@@ -538,13 +535,13 @@ ssxHead (StaticShX list) = listxHead list
 ssxTail :: StaticShX (n : sh) -> StaticShX sh
 ssxTail (_ :!% ssh) = ssh
 
-ssxDropSSX :: forall sh sh'. StaticShX (sh ++ sh') -> StaticShX sh -> StaticShX sh'
+ssxDropSSX :: forall sh sh'. StaticShX sh -> StaticShX (sh ++ sh') -> StaticShX sh'
 ssxDropSSX = coerce (listxDrop @(SMayNat () SNat) @(SMayNat () SNat))
 
-ssxDropIx :: forall sh sh' i. StaticShX (sh ++ sh') -> IxX sh i -> StaticShX sh'
+ssxDropIx :: forall sh sh' i. IxX sh i -> StaticShX (sh ++ sh') -> StaticShX sh'
 ssxDropIx = coerce (listxDrop @(SMayNat () SNat) @(Const i))
 
-ssxDropSh :: forall sh sh' i. StaticShX (sh ++ sh') -> ShX sh i -> StaticShX sh'
+ssxDropSh :: forall sh sh' i. ShX sh i -> StaticShX (sh ++ sh') -> StaticShX sh'
 ssxDropSh = coerce (listxDrop @(SMayNat () SNat) @(SMayNat i SNat))
 
 ssxInit :: forall n sh. StaticShX (n : sh) -> StaticShX (Init (n : sh))
@@ -559,9 +556,9 @@ ssxReplicate (SS (n :: SNat n'))
   | Refl <- lemReplicateSucc @(Nothing @Nat) @n'
   = SUnknown () :!% ssxReplicate n
 
-ssxIotaFrom :: Int -> StaticShX sh -> [Int]
-ssxIotaFrom _ ZKX = []
-ssxIotaFrom i (_ :!% ssh) = i : ssxIotaFrom (i+1) ssh
+ssxIotaFrom :: StaticShX sh -> Int -> [Int]
+ssxIotaFrom ZKX _ = []
+ssxIotaFrom (_ :!% ssh) i = i : ssxIotaFrom ssh (i+1)
 
 ssxFromShX :: ShX sh i -> StaticShX sh
 ssxFromShX ZSX = ZKX