From 36a89c7a26d73b0583ac765d29dadb1b918007f6 Mon Sep 17 00:00:00 2001
From: Tom Smeding <t.j.smeding@uu.nl>
Date: Tue, 9 Apr 2024 18:35:35 +0200
Subject: Transpose functions in the API

---
 src/Data/Array/Mixed.hs           |  9 ++++++++-
 src/Data/Array/Nested.hs          |  2 ++
 src/Data/Array/Nested/Internal.hs | 15 +++++++++++++++
 3 files changed, 25 insertions(+), 1 deletion(-)

(limited to 'src/Data')

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)
-- 
cgit v1.2.3-70-g09d2