diff options
Diffstat (limited to 'src/Data/Array/XArray.hs')
| -rw-r--r-- | src/Data/Array/XArray.hs | 10 |
1 files changed, 10 insertions, 0 deletions
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))) |