1 files changed, 56 insertions, 30 deletions
diff --git a/src/Data/Array/Nested/Permutation.hs b/src/Data/Array/Nested/Permutation.hs
index 8b46d81..e520e0f 100644
--- a/src/Data/Array/Nested/Permutation.hs
+++ b/src/Data/Array/Nested/Permutation.hs
@@ -172,6 +172,62 @@ type family DropLen ref l where
   DropLen '[] l = l
   DropLen (_ : ref) (_ : xs) = DropLen ref xs
 
+listhTakeLen :: forall f is sh. Perm is -> ListH sh f -> ListH (TakeLen is sh) f
+listhTakeLen PNil _ = ZH
+listhTakeLen (_ `PCons` is) (n ::# sh) = n ::# listhTakeLen is sh
+listhTakeLen (_ `PCons` _) ZH = error "Permutation longer than shape"
+
+listhDropLen :: forall f is sh. Perm is -> ListH sh f -> ListH (DropLen is sh) f
+listhDropLen PNil sh = sh
+listhDropLen (_ `PCons` is) (_ ::# sh) = listhDropLen is sh
+listhDropLen (_ `PCons` _) ZH = error "Permutation longer than shape"
+
+listhPermute :: forall f is sh. Perm is -> ListH sh f -> ListH (Permute is sh) f
+listhPermute PNil _ = ZH
+listhPermute (i `PCons` (is :: Perm is')) (sh :: ListH sh f) =
+  listhIndex (Proxy @is') (Proxy @sh) i sh ::# listhPermute is sh
+
+listhIndex :: forall f is shT i sh. Proxy is -> Proxy shT -> SNat i -> ListH sh f -> f (Index i sh)
+listhIndex _ _ SZ (n ::# _) = n
+listhIndex p pT (SS (i :: SNat i')) ((_ :: f n) ::# (sh :: ListH sh' f))
+  | Refl <- lemIndexSucc (Proxy @i') (Proxy @n) (Proxy @sh')
+  = listhIndex p pT i sh
+listhIndex _ _ _ ZH = error "Index into empty shape"
+
+listhPermutePrefix :: forall f is sh. Perm is -> ListH sh f -> ListH (PermutePrefix is sh) f
+listhPermutePrefix perm sh = listhAppend (listhPermute perm (listhTakeLen perm sh)) (listhDropLen perm sh)
+
+ssxTakeLen :: forall is sh. Perm is -> StaticShX sh -> StaticShX (TakeLen is sh)
+ssxTakeLen = coerce (listhTakeLen @(SMayNat ()))
+
+ssxDropLen :: Perm is -> StaticShX sh -> StaticShX (DropLen is sh)
+ssxDropLen = coerce (listhDropLen @(SMayNat ()))
+
+ssxPermute :: Perm is -> StaticShX sh -> StaticShX (Permute is sh)
+ssxPermute = coerce (listhPermute @(SMayNat ()))
+
+ssxIndex :: Proxy is -> Proxy shT -> SNat i -> StaticShX sh -> SMayNat () (Index i sh)
+ssxIndex p1 p2 i = coerce (listhIndex @(SMayNat ()) p1 p2 i)
+
+ssxPermutePrefix :: Perm is -> StaticShX sh -> StaticShX (PermutePrefix is sh)
+ssxPermutePrefix = coerce (listhPermutePrefix @(SMayNat ()))
+
+shxTakeLen :: forall is sh. Perm is -> IShX sh -> IShX (TakeLen is sh)
+shxTakeLen = coerce (listhTakeLen @(SMayNat Int))
+
+shxDropLen :: Perm is -> IShX sh -> IShX (DropLen is sh)
+shxDropLen = coerce (listhDropLen @(SMayNat Int))
+
+shxPermute :: Perm is -> IShX sh -> IShX (Permute is sh)
+shxPermute = coerce (listhPermute @(SMayNat Int))
+
+shxIndex :: Proxy is -> Proxy shT -> SNat i -> IShX sh -> SMayNat Int (Index i sh)
+shxIndex p1 p2 i = coerce (listhIndex @(SMayNat Int) p1 p2 i)
+
+shxPermutePrefix :: Perm is -> IShX sh -> IShX (PermutePrefix is sh)
+shxPermutePrefix = coerce (listhPermutePrefix @(SMayNat Int))
+
+
 listxTakeLen :: forall f is sh. Perm is -> ListX sh f -> ListX (TakeLen is sh) f
 listxTakeLen PNil _ = ZX
 listxTakeLen (_ `PCons` is) (n ::% sh) = n ::% listxTakeLen is sh
@@ -200,36 +256,6 @@ listxPermutePrefix perm sh = listxAppend (listxPermute perm (listxTakeLen perm s
 ixxPermutePrefix :: forall i is sh. Perm is -> IxX sh i -> IxX (PermutePrefix is sh) i
 ixxPermutePrefix = coerce (listxPermutePrefix @(Const i))
 
-ssxTakeLen :: forall is sh. Perm is -> StaticShX sh -> StaticShX (TakeLen is sh)
-ssxTakeLen = coerce (listxTakeLen @(SMayNat ()))
-
-ssxDropLen :: Perm is -> StaticShX sh -> StaticShX (DropLen is sh)
-ssxDropLen = coerce (listxDropLen @(SMayNat ()))
-
-ssxPermute :: Perm is -> StaticShX sh -> StaticShX (Permute is sh)
-ssxPermute = coerce (listxPermute @(SMayNat ()))
-
-ssxIndex :: Proxy is -> Proxy shT -> SNat i -> StaticShX sh -> SMayNat () (Index i sh)
-ssxIndex p1 p2 i = coerce (listxIndex @(SMayNat ()) p1 p2 i)
-
-ssxPermutePrefix :: Perm is -> StaticShX sh -> StaticShX (PermutePrefix is sh)
-ssxPermutePrefix = coerce (listxPermutePrefix @(SMayNat ()))
-
-shxTakeLen :: forall is sh. Perm is -> IShX sh -> IShX (TakeLen is sh)
-shxTakeLen = coerce (listxTakeLen @(SMayNat Int))
-
-shxDropLen :: Perm is -> IShX sh -> IShX (DropLen is sh)
-shxDropLen = coerce (listxDropLen @(SMayNat Int))
-
-shxPermute :: Perm is -> IShX sh -> IShX (Permute is sh)
-shxPermute = coerce (listxPermute @(SMayNat Int))
-
-shxIndex :: Proxy is -> Proxy shT -> SNat i -> IShX sh -> SMayNat Int (Index i sh)
-shxIndex p1 p2 i = coerce (listxIndex @(SMayNat Int) p1 p2 i)
-
-shxPermutePrefix :: Perm is -> IShX sh -> IShX (PermutePrefix is sh)
-shxPermutePrefix = coerce (listxPermutePrefix @(SMayNat Int))
-
 
 -- * Operations on permutations