XArray: Further optimise toListOuter

Again, thanks Mikolaj

1 files changed, 5 insertions, 3 deletions

diff --git a/src/Data/Array/XArray.hs b/src/Data/Array/XArray.hs
index 8b4fd60..6ad48cc 100644
--- a/src/Data/Array/XArray.hs
+++ b/src/Data/Array/XArray.hs
@@ -302,12 +302,14 @@ fromListOuter ssh l
       _ -> XArray (S.ravel (ORB.fromList [length l] (coerce @[XArray sh a] @[S.Array (Rank sh) a] l)))
 
 toListOuter :: forall a n sh. Storable a => XArray (n : sh) a -> [XArray sh a]
-toListOuter (XArray arr) =
+toListOuter (XArray arr@(ORS.A (ORG.A _ t))) =
   case S.shapeL arr of
+    [] -> error "impossible"
     0 : _ -> []
-    -- orthotope doesn't seem to optimise this case, so we need to do it here
+    -- using orthotope's functions here would entail using rerank, which is slow, so we don't
     [_] | Refl <- (unsafeCoerceRefl :: sh :~: '[]) -> coerce (map S.scalar $ S.toList arr)
-    _ -> coerce (ORB.toList (S.unravel arr))
+    -- TODO: this shares the backing array among all the results; is that what we want?
+    n : sh -> coerce $ map (ORG.A sh . OI.indexT t) [0 .. n - 1]
 
 fromList1 :: Storable a => StaticShX '[n] -> [a] -> XArray '[n] a
 fromList1 ssh l =