diff options
| author | Mikolaj Konarski <mikolaj.konarski@funktory.com> | 2025-12-09 10:34:03 +0100 |
|---|---|---|
| committer | Mikolaj Konarski <mikolaj.konarski@funktory.com> | 2026-01-31 10:43:16 +0100 |
| commit | 05ddeb90d8b1be91d8d98c66219b0fac111ec70a (patch) | |
| tree | 2cb28c8286c190f028ffd0428211b8bcf31a84e5 /src/Data/Array/Nested/Mixed | |
| parent | 0a4dae6543a5d5e5515be40d930863ec9fa96e3e (diff) | |
Generalize ix?ToLinear and speed it up a bit
Diffstat (limited to 'src/Data/Array/Nested/Mixed')
| -rw-r--r-- | src/Data/Array/Nested/Mixed/Shape.hs | 15 |
1 files changed, 7 insertions, 8 deletions
diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs index 145ea5f..11ef757 100644 --- a/src/Data/Array/Nested/Mixed/Shape.hs +++ b/src/Data/Array/Nested/Mixed/Shape.hs @@ -284,16 +284,15 @@ ixxZipWith :: (i -> j -> k) -> IxX sh i -> IxX sh j -> IxX sh k ixxZipWith _ ZIX ZIX = ZIX ixxZipWith f (i :.% is) (j :.% js) = f i j :.% ixxZipWith f is js +-- | Given a multidimensional index, get the corresponding linear +-- index into the buffer. {-# INLINEABLE ixxToLinear #-} -ixxToLinear :: IShX sh -> IIxX sh -> Int -ixxToLinear = \sh i -> fst (go sh i) +ixxToLinear :: Num i => IShX sh -> IxX sh i -> i +ixxToLinear = \sh i -> go sh i 0 where - -- returns (index in subarray, size of subarray) - go :: IShX sh -> IIxX sh -> (Int, Int) - go ZSX ZIX = (0, 1) - go (n :$% sh) (i :.% ix) = - let (lidx, sz) = go sh ix - in (sz * i + lidx, fromSMayNat' n * sz) + go :: Num i => IShX sh -> IxX sh i -> i -> i + go ZSX ZIX a = a + go (n :$% sh) (i :.% ix) a = go sh ix (fromIntegral (fromSMayNat' n) * a + i) -- * Mixed shapes |