Add append

1 files changed, 18 insertions, 0 deletions

diff --git a/src/Data/Array/Nested/Internal.hs b/src/Data/Array/Nested/Internal.hs
index f7c383a..209d594 100644
--- a/src/Data/Array/Nested/Internal.hs
+++ b/src/Data/Array/Nested/Internal.hs
@@ -15,6 +15,7 @@
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE UndecidableInstances #-}
+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
 
 {-|
 TODO:
@@ -29,6 +30,8 @@ TODO:
 
 module Data.Array.Nested.Internal where
 
+import Prelude hiding (mappend)
+
 import Control.Monad (forM_)
 import Control.Monad.ST
 import qualified Data.Array.RankedS as S
@@ -354,6 +357,13 @@ mtranspose perm =
   mlift (\(Proxy @sh') -> X.rerankTop (knownShapeX @sh) (knownShapeX @sh) (knownShapeX @sh')
                             (X.transpose perm))
 
+mappend :: forall n m sh a. (KnownShapeX sh, KnownShapeX (n : sh), KnownShapeX (m : sh), KnownShapeX (X.AddMaybe n m : sh), Elt a)
+        => Mixed (n : sh) a -> Mixed (m : sh) a -> Mixed (X.AddMaybe n m : sh) a
+mappend = mlift2 go
+  where go :: forall sh' b. (KnownShapeX sh', Storable b)
+           => Proxy sh' -> XArray (n : sh ++ sh') b -> XArray (m : sh ++ sh') b -> XArray (X.AddMaybe n m : sh ++ sh') b
+        go Proxy | Dict <- X.lemAppKnownShapeX (knownShapeX @sh) (knownShapeX @sh') = X.append
+
 mliftPrim :: (KnownShapeX sh, Storable a)
           => (a -> a)
           -> Mixed sh (Primitive a) -> Mixed sh (Primitive a)
@@ -693,6 +703,10 @@ rtranspose perm (Ranked arr)
   | Dict <- lemKnownReplicate (Proxy @n)
   = Ranked (mtranspose perm arr)
 
+rappend :: forall n a. (KnownINat n, Elt a)
+        => Ranked (S n) a -> Ranked (S n) a -> Ranked (S n) a
+rappend | Dict <- lemKnownReplicate (Proxy @n) = coerce mappend
+
 
 -- ====== API OF SHAPED ARRAYS ====== --
 
@@ -782,3 +796,7 @@ stranspose :: forall sh a. (KnownShape sh, Elt a) => [Int] -> Shaped sh a -> Sha
 stranspose perm (Shaped arr)
   | Dict <- lemKnownMapJust (Proxy @sh)
   = Shaped (mtranspose perm arr)
+
+sappend :: forall n m sh a. (KnownNat n, KnownNat m, KnownShape sh, Elt a)
+        => Shaped (n : sh) a -> Shaped (m : sh) a -> Shaped (n + m : sh) a
+sappend | Dict <- lemKnownMapJust (Proxy @sh) = coerce mappend