Add *flatten

1 files changed, 10 insertions, 0 deletions

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