From 5f6a64660b16d8f188caca5216e55debf4264611 Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Mon, 10 Jun 2024 23:28:11 +0200
Subject: Add *flatten

---
 src/Data/Array/Nested/Internal/Mixed.hs  |  3 +++
 src/Data/Array/Nested/Internal/Ranked.hs |  3 +++
 src/Data/Array/Nested/Internal/Shape.hs  | 10 ++++++++++
 src/Data/Array/Nested/Internal/Shaped.hs |  5 +++++
 4 files changed, 21 insertions(+)

(limited to 'src/Data/Array/Nested')

diff --git a/src/Data/Array/Nested/Internal/Mixed.hs b/src/Data/Array/Nested/Internal/Mixed.hs
index 9c2096d..69df52a 100644
--- a/src/Data/Array/Nested/Internal/Mixed.hs
+++ b/src/Data/Array/Nested/Internal/Mixed.hs
@@ -796,6 +796,9 @@ mreshape sh' arr =
         (\sshIn -> X.reshapePartial (ssxFromShape (mshape arr)) sshIn sh')
         arr
 
+mflatten :: Elt a => Mixed sh a -> Mixed '[Flatten sh] a
+mflatten arr = mreshape (shxFlatten (mshape arr) :$% ZSX) arr
+
 miota :: (Enum a, PrimElt a) => SNat n -> Mixed '[Just n] a
 miota sn = fromPrimitive $ M_Primitive (SKnown sn :$% ZSX) (X.iota sn)
 
diff --git a/src/Data/Array/Nested/Internal/Ranked.hs b/src/Data/Array/Nested/Internal/Ranked.hs
index 3e9f528..59c1820 100644
--- a/src/Data/Array/Nested/Internal/Ranked.hs
+++ b/src/Data/Array/Nested/Internal/Ranked.hs
@@ -446,6 +446,9 @@ rreshape sh' rarr@(Ranked arr)
   , Dict <- lemKnownReplicate (shrLengthSNat sh')
   = Ranked (mreshape (shCvtRX sh') arr)
 
+rflatten :: Elt a => Ranked n a -> Ranked 1 a
+rflatten (Ranked arr) = mtoRanked (mflatten arr)
+
 riota :: (Enum a, PrimElt a, Elt a) => Int -> Ranked 1 a
 riota n = TN.withSomeSNat (fromIntegral n) $ mtoRanked . miota
 
diff --git a/src/Data/Array/Nested/Internal/Shape.hs b/src/Data/Array/Nested/Internal/Shape.hs
index 9d718cc..7d95f61 100644
--- a/src/Data/Array/Nested/Internal/Shape.hs
+++ b/src/Data/Array/Nested/Internal/Shape.hs
@@ -6,6 +6,7 @@
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE NoStarIsType #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE QuantifiedConstraints #-}
@@ -18,6 +19,7 @@
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UndecidableInstances #-}
 {-# LANGUAGE ViewPatterns #-}
 {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
 {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
@@ -475,6 +477,14 @@ shsIndex pis pshT i sh = coerce (fst (listsIndex @SNat pis pshT i (coerce sh)))
 shsPermutePrefix :: forall is sh. Perm is -> ShS sh -> ShS (PermutePrefix is sh)
 shsPermutePrefix = coerce (listsPermutePrefix @SNat)
 
+type family Product sh where
+  Product '[] = 1
+  Product (n : ns) = n * Product ns
+
+shsProduct :: ShS sh -> SNat (Product sh)
+shsProduct ZSS = SNat
+shsProduct (n :$$ sh) = n `snatMul` shsProduct sh
+
 -- | Evidence for the static part of a shape. This pops up only when you are
 -- polymorphic in the element type of an array.
 type KnownShS :: [Nat] -> Constraint
diff --git a/src/Data/Array/Nested/Internal/Shaped.hs b/src/Data/Array/Nested/Internal/Shaped.hs
index 863e604..1855015 100644
--- a/src/Data/Array/Nested/Internal/Shaped.hs
+++ b/src/Data/Array/Nested/Internal/Shaped.hs
@@ -370,6 +370,11 @@ srev1 arr = slift (sshape arr) (\_ -> X.rev1) arr
 sreshape :: Elt a => ShS sh' -> Shaped sh a -> Shaped sh' a
 sreshape sh' (Shaped arr) = Shaped (mreshape (shCvtSX sh') arr)
 
+sflatten :: Elt a => Shaped sh a -> Shaped '[Product sh] a
+sflatten arr =
+  case shsProduct (sshape arr) of  -- TODO: simplify when removing the KnownNat stuff
+    n@SNat -> sreshape (n :$$ ZSS) arr
+
 siota :: (Enum a, PrimElt a) => SNat n -> Shaped '[n] a
 siota sn = Shaped (miota sn)
 
-- 
cgit v1.2.3-70-g09d2