Transpose functions in the API

4 files changed, 27 insertions, 2 deletions

diff --git a/src/Data/Array/Mixed.hs b/src/Data/Array/Mixed.hs
index 2ad1d26..12c247f 100644
--- a/src/Data/Array/Mixed.hs
+++ b/src/Data/Array/Mixed.hs
@@ -251,13 +251,20 @@ rerank ssh ssh1 ssh2 f (XArray arr)
   , Refl <- lemRankApp ssh ssh1
   , Refl <- lemRankApp ssh ssh2
   , Dict <- lemKnownNatRankSSX (ssxAppend ssh ssh2)  -- these two should be redundant but the
-  , Dict <- gknownNat (Proxy @(Rank (sh ++ sh2)))  -- solver is not clever enough
+  , Dict <- gknownNat (Proxy @(Rank (sh ++ sh2)))    -- solver is not clever enough
   = XArray (U.rerank @(GNat (Rank sh)) @(GNat (Rank sh1)) @(GNat (Rank sh2))
                 (\a -> unXArray (f (XArray a)))
                 arr)
   where
     unXArray (XArray a) = a
 
+rerankTop :: forall sh sh1 sh2 a b.
+             (U.Unbox a, U.Unbox b)
+          => StaticShapeX sh1 -> StaticShapeX sh2 -> StaticShapeX sh
+          -> (XArray sh1 a -> XArray sh2 b)
+          -> XArray (sh1 ++ sh) a -> XArray (sh2 ++ sh) b
+rerankTop ssh1 ssh2 ssh f = transpose2 ssh ssh2 . rerank ssh ssh1 ssh2 f . transpose2 ssh1 ssh
+
 rerank2 :: forall sh sh1 sh2 a b c.
            (U.Unbox a, U.Unbox b, U.Unbox c)
         => StaticShapeX sh -> StaticShapeX sh1 -> StaticShapeX sh2
diff --git a/src/Data/Array/Nested.hs b/src/Data/Array/Nested.hs
index 983a636..aa7c7f9 100644
--- a/src/Data/Array/Nested.hs
+++ b/src/Data/Array/Nested.hs
@@ -4,6 +4,7 @@ module Data.Array.Nested (
   Ranked,
   IxR(..),
   rshape, rindex, rindexPartial, rgenerate, rsumOuter1,
+  rtranspose,
   -- ** Lifting orthotope operations to 'Ranked' arrays
   rlift,
 
@@ -12,6 +13,7 @@ module Data.Array.Nested (
   IxS(..),
   KnownShape(..), SShape(..),
   sshape, sindex, sindexPartial, sgenerate, ssumOuter1,
+  stranspose,
   -- ** Lifting orthotope operations to 'Shaped' arrays
   slift,
 
diff --git a/src/Data/Array/Nested/Internal.hs b/src/Data/Array/Nested/Internal.hs
index fa40921..d92c79c 100644
--- a/src/Data/Array/Nested/Internal.hs
+++ b/src/Data/Array/Nested/Internal.hs
@@ -315,6 +315,11 @@ mgenerate sh f
     checkBounds (n ::@ sh') (n' :$@ ssh') = n == fromIntegral (unSNat n') && checkBounds sh' ssh'
     checkBounds (_ ::? sh') (() :$? ssh') = checkBounds sh' ssh'
 
+mtranspose :: forall sh a. (KnownShapeX sh, Elt a) => [Int] -> Mixed sh a -> Mixed sh a
+mtranspose perm =
+  mlift (\(Proxy @sh') -> X.rerankTop (knownShapeX @sh) (knownShapeX @sh) (knownShapeX @sh')
+                            (X.transpose perm))
+
 
 -- | A rank-typed array: the number of dimensions of the array (its /rank/) is
 -- represented on the type level as a 'Nat'.
@@ -578,6 +583,11 @@ rsumOuter1 (Ranked arr)
     . coerce @(Mixed (Replicate (S n) Nothing) a) @(XArray (Replicate (S n) Nothing) a)
     $ arr
 
+rtranspose :: forall n a. (KnownNat n, Elt a) => [Int] -> Ranked n a -> Ranked n a
+rtranspose perm (Ranked arr)
+  | Dict <- lemKnownReplicate (Proxy @n)
+  = Ranked (mtranspose perm arr)
+
 
 -- ====== API OF SHAPED ARRAYS ====== --
 
@@ -640,3 +650,8 @@ ssumOuter1 (Shaped arr)
     . X.sumOuter (knownNat @n :$@ SZX) (knownShapeX @(MapJust sh))
     . coerce @(Mixed (Just n : MapJust sh) a) @(XArray (Just n : MapJust sh) a)
     $ arr
+
+stranspose :: forall sh a. (KnownShape sh, Elt a) => [Int] -> Shaped sh a -> Shaped sh a
+stranspose perm (Shaped arr)
+  | Dict <- lemKnownMapJust (Proxy @sh)
+  = Shaped (mtranspose perm arr)
diff --git a/test/Main.hs b/test/Main.hs
index 156e0a5..8257ff0 100644
--- a/test/Main.hs
+++ b/test/Main.hs
@@ -9,7 +9,7 @@ import Data.Array.Nested
 arr :: Ranked N2 (Shaped [N2, N3] (Double, Int))
 arr = rgenerate (3 ::: 4 ::: IZR) $ \(i ::: j ::: IZR) ->
         sgenerate @[N2, N3] (2 ::$ 3 ::$ IZS) $ \(k ::$ l ::$ IZS) ->
-          let s = i + j + k + l
+          let s = 24*i + 6*j + 3*k + l
           in (fromIntegral s, s)
 
 foo :: (Double, Int)
@@ -19,3 +19,4 @@ main :: IO ()
 main = do
   print arr
   print foo
+  print (rtranspose [1,0] arr)