From 9560d0f26420409afd2230fb7e5e111eafcced06 Mon Sep 17 00:00:00 2001 From: Mikolaj Konarski Date: Mon, 1 Dec 2025 01:44:05 +0100 Subject: Expose the unfolding of the indexing operations --- src/Data/Array/Nested/Mixed.hs | 2 ++ src/Data/Array/Nested/Ranked.hs | 2 ++ src/Data/Array/Nested/Shaped.hs | 2 ++ 3 files changed, 6 insertions(+) (limited to 'src/Data/Array/Nested') diff --git a/src/Data/Array/Nested/Mixed.hs b/src/Data/Array/Nested/Mixed.hs index 6d4ffd6..e3aa7a1 100644 --- a/src/Data/Array/Nested/Mixed.hs +++ b/src/Data/Array/Nested/Mixed.hs @@ -399,7 +399,9 @@ class Elt a => KnownElt a where -- Arrays of scalars are basically just arrays of scalars. instance Storable a => Elt (Primitive a) where mshape (M_Primitive sh _) = sh + {-# INLINEABLE mindex #-} mindex (M_Primitive _ a) i = Primitive (X.index a i) + {-# INLINEABLE mindexPartial #-} mindexPartial (M_Primitive sh a) i = M_Primitive (shxDropIx i sh) (X.indexPartial a i) mscalar (Primitive x) = M_Primitive ZSX (X.scalar x) mfromListOuterSN sn l@(arr1 :| _) = diff --git a/src/Data/Array/Nested/Ranked.hs b/src/Data/Array/Nested/Ranked.hs index 2fbfdd8..bf35cc4 100644 --- a/src/Data/Array/Nested/Ranked.hs +++ b/src/Data/Array/Nested/Ranked.hs @@ -49,9 +49,11 @@ remptyArray = mtoRanked (memptyArray ZSX) rsize :: Elt a => Ranked n a -> Int rsize = shrSize . rshape +{-# INLINEABLE rindex #-} rindex :: Elt a => Ranked n a -> IIxR n -> a rindex (Ranked arr) idx = mindex arr (ixxFromIxR idx) +{-# INLINEABLE rindexPartial #-} rindexPartial :: forall n m a. Elt a => Ranked (n + m) a -> IIxR n -> Ranked m a rindexPartial (Ranked arr) idx = Ranked (mindexPartial @a @(Replicate n Nothing) @(Replicate m Nothing) diff --git a/src/Data/Array/Nested/Shaped.hs b/src/Data/Array/Nested/Shaped.hs index 8957549..82dfc91 100644 --- a/src/Data/Array/Nested/Shaped.hs +++ b/src/Data/Array/Nested/Shaped.hs @@ -52,6 +52,7 @@ srank = shsRank . sshape ssize :: Elt a => Shaped sh a -> Int ssize = shsSize . sshape +{-# INLINEABLE sindex #-} sindex :: Elt a => Shaped sh a -> IIxS sh -> a sindex (Shaped arr) idx = mindex arr (ixxFromIxS idx) @@ -59,6 +60,7 @@ shsTakeIx :: Proxy sh' -> ShS (sh ++ sh') -> IIxS sh -> ShS sh shsTakeIx _ _ ZIS = ZSS shsTakeIx p sh (_ :.$ idx) = case sh of n :$$ sh' -> n :$$ shsTakeIx p sh' idx +{-# INLINEABLE sindexPartial #-} sindexPartial :: forall sh1 sh2 a. Elt a => Shaped (sh1 ++ sh2) a -> IIxS sh1 -> Shaped sh2 a sindexPartial sarr@(Shaped arr) idx = Shaped (mindexPartial @a @(MapJust sh1) @(MapJust sh2) -- cgit v1.2.3-70-g09d2