Simplify and slightly speed up internal fromList functions

1 files changed, 25 insertions, 23 deletions

diff --git a/src/Data/Array/XArray.hs b/src/Data/Array/XArray.hs
index 42aed6e..77c0dc0 100644
--- a/src/Data/Array/XArray.hs
+++ b/src/Data/Array/XArray.hs
@@ -310,22 +310,19 @@ sumOuter ssh ssh' arr
 -- the list's spine must be fully materialised to compute its length before
 -- constructing the array. The list can't be empty (not enough information
 -- in the given shape to guess the shape of the empty array, in general).
-fromListOuter :: forall n sh a. Storable a
-              => StaticShX (n : sh) -> [XArray sh a] -> XArray (n : sh) a
-fromListOuter ssh l
-  | Dict <- lemKnownNatRankSSX (ssxTail ssh)
+{-# INLINE fromListOuterSN #-}
+fromListOuterSN :: forall n sh a. Storable a
+                => SNat n -> IShX sh -> [XArray sh a] -> XArray (Just n : sh) a
+fromListOuterSN m sh l
+  | Dict <- lemKnownNatRank sh
   , let l' = coerce @[XArray sh a] @[S.Array (Rank sh) a] l
-  = case ssh of
-      _ :!% ZKX ->
-        fromList1 ssh (map S.unScalar l')
-      _ ->
-        let n = case ssh of
-              SKnown m :!% _ -> fromSNat' m
-              _ -> length l
-        in XArray (ravelOuterN n l')
+  = case sh of
+      ZSX -> fromList1SN m (map S.unScalar l')
+      _ -> XArray (ravelOuterN (fromSNat' m) l')
 
 -- | This checks that the list has the given length and that all shapes in the
 -- list are equal. The list must be non-empty, and is streamed.
+{-# INLINEABLE ravelOuterN #-}
 ravelOuterN :: (KnownNat k, Storable a)
             => Int -> [S.Array k a] -> S.Array (1 + k) a
 ravelOuterN 0 _ = error "ravelOuterN: N == 0"
@@ -351,8 +348,8 @@ ravelOuterN k as@(a0 : _) = runST $ do
   else error $ "ravelOuterN: list too short " ++ show (nFinal, k)
 
 toListOuter :: forall a n sh. Storable a => XArray (n : sh) a -> [XArray sh a]
-toListOuter (XArray arr@(ORS.A (ORG.A _ t))) =
-  case S.shapeL arr of
+toListOuter (XArray arr@(ORS.A (ORG.A shArr t))) =
+  case shArr of
     [] -> error "impossible"
     0 : _ -> []
     -- using orthotope's functions here would entail using rerank, which is slow, so we don't
@@ -362,15 +359,20 @@ toListOuter (XArray arr@(ORS.A (ORG.A _ t))) =
 -- | If @n@ is an 'SKnown' dimension, the list is streamed. If @n@ is unknown,
 -- the list's spine must be fully materialised to compute its length before
 -- constructing the array.
-fromList1 :: Storable a => StaticShX '[n] -> [a] -> XArray '[n] a
-fromList1 ssh l =
-  case ssh of
-    SKnown m :!% _ ->
-      let n = fromSNat' m  -- do length check and vector construction simultaneously so that l can be streamed
-      in XArray (S.fromVector [n] (VGC.fromListNChecked n l))
-    _ ->
-      let n = length l  -- avoid S.fromList because it takes a length _and_ does another length check itself
-      in XArray (S.fromVector [n] (VS.fromListN         n l))
+{-# INLINE fromList1 #-}
+fromList1 :: Storable a => [a] -> XArray '[Nothing] a
+fromList1 l =
+  let n = length l  -- avoid S.fromList because it takes a length _and_ does another length check itself
+  in XArray (S.fromVector [n] (VS.fromListN n l))
+
+-- | If @n@ is an 'SKnown' dimension, the list is streamed. If @n@ is unknown,
+-- the list's spine must be fully materialised to compute its length before
+-- constructing the array.
+{-# INLINE fromList1SN #-}
+fromList1SN :: Storable a => SNat n -> [a] -> XArray '[Just n] a
+fromList1SN m l =
+  let n = fromSNat' m  -- do length check and vector construction simultaneously so that l can be streamed
+  in XArray (S.fromVector [n] (VGC.fromListNChecked n l))
 
 toList1 :: Storable a => XArray '[n] a -> [a]
 toList1 (XArray arr) = S.toList arr