From 9560d0f26420409afd2230fb7e5e111eafcced06 Mon Sep 17 00:00:00 2001
From: Mikolaj Konarski <mikolaj.konarski@funktory.com>
Date: Mon, 1 Dec 2025 01:44:05 +0100
Subject: Expose the unfolding of the indexing operations

---
 src/Data/Array/XArray.hs | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

(limited to 'src/Data/Array/XArray.hs')

diff --git a/src/Data/Array/XArray.hs b/src/Data/Array/XArray.hs
index 0f87168..ee83654 100644
--- a/src/Data/Array/XArray.hs
+++ b/src/Data/Array/XArray.hs
@@ -114,10 +114,20 @@ generate sh f = fromVector sh $ VS.generate (shxSize sh) (f . ixxFromLinear sh)
 --   XArray . S.fromVector (shxShapeL sh)
 --     <$> VS.generateM (shxSize sh) (f . ixxFromLinear sh)
 
+{-# INLINEABLE indexPartial #-}
 indexPartial :: Storable a => XArray (sh ++ sh') a -> IIxX sh -> XArray sh' a
 indexPartial (XArray arr) ZIX = XArray arr
 indexPartial (XArray arr) (i :.% idx) = indexPartial (XArray (S.index arr i)) idx
-
+{- Strangely, this increases allocation and there's no noticeable speedup:
+indexPartial (XArray (ORS.A (ORG.A sh t))) ix =
+  let linear = OI.offset t + sum (zipWith (*) (ixxToList ix) (OI.strides t))
+      len = ixxLength ix
+  in XArray (ORS.A (ORG.A (drop len sh)
+                          OI.T{ OI.strides = drop len (OI.strides t)
+                              , OI.offset = linear
+                              , OI.values = OI.values t })) -}
+
+{-# INLINEABLE index #-}
 index :: forall sh a. Storable a => XArray sh a -> IIxX sh -> a
 index (XArray (ORS.A (ORG.A _ t))) i =
   OI.values t VS.! (OI.offset t + sum (zipWith (*) (toList i) (OI.strides t)))
-- 
cgit v1.2.3-70-g09d2