Add utility function

1 files changed, 8 insertions, 0 deletions

diff --git a/src/Data/Array/Mixed.hs b/src/Data/Array/Mixed.hs
index 7293914..748914c 100644
--- a/src/Data/Array/Mixed.hs
+++ b/src/Data/Array/Mixed.hs
@@ -433,6 +433,14 @@ type family Flatten' acc sh where
   Flatten' acc (Nothing : sh) = Nothing
   Flatten' acc (Just n : sh) = Flatten' (acc * n) sh
 
+flattenSSX :: StaticShX sh -> SMayNat () SNat (Flatten sh)
+flattenSSX = go (SNat @1)
+  where
+    go :: SNat acc -> StaticShX sh -> SMayNat () SNat (Flatten' acc sh)
+    go acc ZKX = SKnown acc
+    go _ (SUnknown () :!% _) = SUnknown ()
+    go acc (SKnown sn :!% sh) = go (mulSNat acc sn) sh
+
 flattenSh :: IShX sh -> SMayNat Int SNat (Flatten sh)
 flattenSh = go (SNat @1)
   where