1 files changed, 13 insertions, 4 deletions
diff --git a/src/Data/Array/Nested/Ranked/Shape.hs b/src/Data/Array/Nested/Ranked/Shape.hs
index 200fac5..50338d2 100644
--- a/src/Data/Array/Nested/Ranked/Shape.hs
+++ b/src/Data/Array/Nested/Ranked/Shape.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DeriveFoldable #-}
@@ -8,6 +9,7 @@
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE NoStarIsType #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE PolyKinds #-}
@@ -21,6 +23,7 @@
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE UndecidableInstances #-}
 {-# LANGUAGE ViewPatterns #-}
 {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
@@ -33,6 +36,7 @@ import Data.Foldable qualified as Foldable
 import Data.Kind (Type)
 import Data.Proxy
 import Data.Type.Equality
+import GHC.Exts (Int(..), Int#, quotRemInt#)
 import GHC.Generics (Generic)
 import GHC.IsList (IsList)
 import GHC.IsList qualified as IsList
@@ -335,11 +339,16 @@ shrEnum = shrEnum'
 
 {-# INLINABLE shrEnum' #-}  -- ensure this can be specialised at use site
 shrEnum' :: Num i => IShR sh -> [IxR sh i]
-shrEnum' = \sh -> go sh id []
+shrEnum' sh = [fromLin sh suffixes li# | I# li# <- [0 .. shrSize sh - 1]]
   where
-    go :: Num i => IShR sh -> (IxR sh i -> a) -> [a] -> [a]
-    go ZSR f = (f ZIR :)
-    go (n :$: sh) f = foldr (.) id [go sh (f . (fromIntegral i :.:)) | i <- [0 .. n - 1]]
+    suffixes = drop 1 (scanr (*) 1 (Foldable.toList sh))
+
+    fromLin :: Num i => IShR sh -> [Int] -> Int# -> IxR sh i
+    fromLin ZSR _ _ = ZIR
+    fromLin (_ :$: sh') (I# suff# : suffs) i# =
+      let !(# q#, r# #) = i# `quotRemInt#` suff#  -- suff == shrSize sh'
+      in fromIntegral (I# q#) :.: fromLin sh' suffs r#
+    fromLin _ _ _ = error "impossible"
 
 
 -- | Untyped: length is checked at runtime.