applyPerm* functions

1 files changed, 83 insertions, 10 deletions

diff --git a/src/Data/Array/Nested/Internal.hs b/src/Data/Array/Nested/Internal.hs
index f8d16aa..712c5f1 100644
--- a/src/Data/Array/Nested/Internal.hs
+++ b/src/Data/Array/Nested/Internal.hs
@@ -153,6 +153,9 @@ lemReplicatePlusApp sn _ _ = go sn
 lemLeqPlus :: n <= m => Proxy n -> Proxy m -> Proxy k -> (n <=? (m + k)) :~: 'True
 lemLeqPlus _ _ _ = Refl
 
+lemLeqSuccSucc :: (k + 1 <= n) => Proxy k -> Proxy n -> (k <=? n - 1) :~: True
+lemLeqSuccSucc _ _ = unsafeCoerce Refl
+
 lemDropLenApp :: X.Rank l1 <= X.Rank l2
               => Proxy l1 -> Proxy l2 -> Proxy rest
               -> X.DropLen l1 l2 ++ rest :~: X.DropLen l1 (l2 ++ rest)
@@ -197,6 +200,19 @@ showListR f l = showString "[" . go "" l . showString "]"
     go _ ZR = id
     go prefix (x ::: xs) = showString prefix . f x . go "," xs
 
+listrAppend :: ListR n i -> ListR m i -> ListR (n + m) i
+listrAppend ZR sh = sh
+listrAppend (x ::: xs) sh = x ::: listrAppend xs sh
+
+listrFromList :: [i] -> (forall n. ListR n i -> r) -> r
+listrFromList [] k = k ZR
+listrFromList (x : xs) k = listrFromList xs $ \l -> k (x ::: l)
+
+listrIndex :: forall k n i. (k + 1 <= n) => SNat k -> ListR n i -> i
+listrIndex SZ (x ::: _) = x
+listrIndex (SS i) (_ ::: xs) | Refl <- lemLeqSuccSucc (Proxy @k) (Proxy @n) = listrIndex i xs
+listrIndex _ ZR = error "k + 1 <= 0"
+
 
 -- | An index into a rank-typed array.
 type role IxR nominal representational
@@ -1497,6 +1513,36 @@ rsumOuter1 :: forall n a. (NumElt a, PrimElt a)
            => Ranked (n + 1) a -> Ranked n a
 rsumOuter1 = rfromPrimitive . rsumOuter1P . rtoPrimitive
 
+applyPermR :: forall i n. [Int] -> ListR n i -> ListR n i
+applyPermR = \perm sh ->
+  listrFromList perm $ \sperm ->
+    case (snatFromListR sperm, snatFromListR sh) of
+      (permlen@SNat, shlen@SNat) -> case cmpNat permlen shlen of
+        LTI -> let (pre, post) = listrSplitAt permlen sh in listrAppend (applyPermRFull permlen sperm pre) post
+        EQI -> let (pre, post) = listrSplitAt permlen sh in listrAppend (applyPermRFull permlen sperm pre) post
+        GTI -> error $ "Length of permutation (" ++ show (fromSNat' permlen) ++ ")"
+                       ++ " > length of shape (" ++ show (fromSNat' shlen) ++ ")"
+  where
+    listrSplitAt :: m <= n' => SNat m -> ListR n' i -> (ListR m i, ListR (n' - m) i)
+    listrSplitAt SZ sh = (ZR, sh)
+    listrSplitAt (SS m) (n ::: sh) = (\(pre, post) -> (n ::: pre, post)) (listrSplitAt m sh)
+    listrSplitAt SS{} ZR = error "m' + 1 <= 0"
+
+    applyPermRFull :: SNat m -> ListR k Int -> ListR m i -> ListR k i
+    applyPermRFull _ ZR _ = ZR
+    applyPermRFull sm@SNat (i ::: perm) l =
+      TypeNats.withSomeSNat (fromIntegral i) $ \si@(SNat :: SNat idx) ->
+        case cmpNat (SNat @(idx + 1)) sm of
+          LTI -> listrIndex si l ::: applyPermRFull sm perm l
+          EQI -> listrIndex si l ::: applyPermRFull sm perm l
+          GTI -> error "applyPermR: Index in permutation out of range"
+
+applyPermIxR :: forall n i. [Int] -> IxR n i -> IxR n i
+applyPermIxR = coerce (applyPermR @i)
+
+applyPermShR :: forall n i. [Int] -> ShR n i -> ShR n i
+applyPermShR = coerce (applyPermR @i)
+
 rtranspose :: forall n a. Elt a => [Int] -> Ranked n a -> Ranked n a
 rtranspose perm arr
   | sn@SNat <- snatFromShR (rshape arr)
@@ -1835,21 +1881,48 @@ lemCommMapJustPermute (i `HCons` is) sh
   , Refl <- lemCommMapJustIndex i sh
   = Refl
 
+listsAppend :: ListS sh f -> ListS sh' f -> ListS (sh ++ sh') f
+listsAppend ZS idx' = idx'
+listsAppend (i ::$ idx) idx' = i ::$ listsAppend idx idx'
+
+listsTakeLen :: forall f is sh. HList SNat is -> ListS sh f -> ListS (X.TakeLen is sh) f
+listsTakeLen HNil _ = ZS
+listsTakeLen (_ `HCons` is) (n ::$ sh) = n ::$ listsTakeLen is sh
+listsTakeLen (_ `HCons` _) ZS = error "Permutation longer than shape"
+
+listsDropLen :: forall f is sh. HList SNat is -> ListS sh f -> ListS (DropLen is sh) f
+listsDropLen HNil sh = sh
+listsDropLen (_ `HCons` is) (_ ::$ sh) = listsDropLen is sh
+listsDropLen (_ `HCons` _) ZS = error "Permutation longer than shape"
+
+listsPermute :: forall f is sh. HList SNat is -> ListS sh f -> ListS (X.Permute is sh) f
+listsPermute HNil _ = ZS
+listsPermute (i `HCons` (is :: HList SNat is')) (sh :: ListS sh f) = listsIndex (Proxy @is') (Proxy @sh) i sh (listsPermute is sh)
+
+listsIndex :: forall f i is sh shT. Proxy is -> Proxy shT -> SNat i -> ListS sh f -> ListS (X.Permute is shT) f -> ListS (X.Index i sh : X.Permute is shT) f
+listsIndex _ _ SZ (n ::$ _) rest = n ::$ rest
+listsIndex p pT (SS (i :: SNat i')) ((_ :: f n) ::$ (sh :: ListS sh' f)) rest
+  | Refl <- X.lemIndexSucc (Proxy @i') (Proxy @n) (Proxy @sh')
+  = listsIndex p pT i sh rest
+listsIndex _ _ _ ZS _ = error "Index into empty shape"
+
 shsTakeLen :: HList SNat is -> ShS sh -> ShS (X.TakeLen is sh)
-shsTakeLen HNil _ = ZSS
-shsTakeLen (_ `HCons` is) (n :$$ sh) = n :$$ shsTakeLen is sh
-shsTakeLen (_ `HCons` _) ZSS = error "Permutation longer than shape"
+shsTakeLen = coerce (listsTakeLen @SNat)
 
 shsPermute :: HList SNat is -> ShS sh -> ShS (X.Permute is sh)
-shsPermute HNil _ = ZSS
-shsPermute (i `HCons` (is :: HList SNat is')) (sh :: ShS sh) = shsIndex (Proxy @is') (Proxy @sh) i sh (shsPermute is sh)
+shsPermute = coerce (listsPermute @SNat)
 
 shsIndex :: Proxy is -> Proxy shT -> SNat i -> ShS sh -> ShS (X.Permute is shT) -> ShS (X.Index i sh : X.Permute is shT)
-shsIndex _ _ SZ (n :$$ _) rest = n :$$ rest
-shsIndex p pT (SS (i :: SNat i')) ((_ :: SNat n) :$$ (sh :: ShS sh')) rest
-  | Refl <- X.lemIndexSucc (Proxy @i') (Proxy @n) (Proxy @sh')
-  = shsIndex p pT i sh rest
-shsIndex _ _ _ ZSS _ = error "Index into empty shape"
+shsIndex pis pshT = coerce (listsIndex @SNat pis pshT)
+
+applyPermS :: forall f is sh. HList SNat is -> ListS sh f -> ListS (PermutePrefix is sh) f
+applyPermS perm sh = listsAppend (listsPermute perm (listsTakeLen perm sh)) (listsDropLen perm sh)
+
+applyPermIxS :: forall i is sh. HList SNat is -> IxS sh i -> IxS (PermutePrefix is sh) i
+applyPermIxS = coerce (applyPermS @(Const i))
+
+applyPermShS :: forall is sh. HList SNat is -> ShS sh -> ShS (PermutePrefix is sh)
+applyPermShS = coerce (applyPermS @SNat)
 
 stranspose :: forall is sh a. (X.Permutation is, X.Rank is <= X.Rank sh, Elt a)
            => HList SNat is -> Shaped sh a -> Shaped (X.PermutePrefix is sh) a