diff options
| author | Mikolaj Konarski <mikolaj.konarski@funktory.com> | 2025-12-16 13:31:36 +0100 |
|---|---|---|
| committer | Mikolaj Konarski <mikolaj.konarski@funktory.com> | 2026-02-18 14:23:02 +0100 |
| commit | bf9b6b8f147c1cb26bba4459d760723648f72b9d (patch) | |
| tree | 250081b82813f3d4d67b9bc8f05218754ad45587 /src/Data/Array/Nested/Mixed | |
| parent | 1dc51cedf3e6c01a366753f48ce2b7c842b05185 (diff) | |
De-triplicate sh?Enum
Diffstat (limited to 'src/Data/Array/Nested/Mixed')
| -rw-r--r-- | src/Data/Array/Nested/Mixed/Shape.hs | 33 |
1 files changed, 17 insertions, 16 deletions
diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs index 057418e..7422b0a 100644 --- a/src/Data/Array/Nested/Mixed/Shape.hs +++ b/src/Data/Array/Nested/Mixed/Shape.hs @@ -328,6 +328,23 @@ ixxFromLinear = \sh -> -- give this function arity 1 so that suffixes is shared outrange sh i = error $ "ixxFromLinear: out of range (" ++ show i ++ " in array of shape " ++ show sh ++ ")" +shxEnum :: IShX sh -> [IIxX sh] +shxEnum = shxEnum' + +{-# INLINABLE shxEnum' #-} -- ensure this can be specialised at use site +shxEnum' :: Num i => IShX sh -> [IxX sh i] +shxEnum' sh = [fromLin sh suffixes li# | I# li# <- [0 .. shxSize sh - 1]] + where + suffixes = drop 1 (scanr (*) 1 (shxToList sh)) + + fromLin :: Num i => IShX sh -> [Int] -> Int# -> IxX sh i + fromLin ZSX _ _ = ZIX + 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" + + -- * Mixed shape-like lists to be used for ShX and StaticShX data SMayNat i n where @@ -649,22 +666,6 @@ shxSplitApp :: proxy sh' -> StaticShX sh -> ShX (sh ++ sh') i -> (ShX sh i, ShX shxSplitApp _ ZKX idx = (ZSX, idx) shxSplitApp p (_ :!% ssh) (i :$% idx) = first (i :$%) (shxSplitApp p ssh idx) -shxEnum :: IShX sh -> [IIxX sh] -shxEnum = shxEnum' - -{-# INLINABLE shxEnum' #-} -- ensure this can be specialised at use site -shxEnum' :: Num i => IShX sh -> [IxX sh i] -shxEnum' sh = [fromLin sh suffixes li# | I# li# <- [0 .. shxSize sh - 1]] - where - suffixes = drop 1 (scanr (*) 1 (shxToList sh)) - - fromLin :: Num i => IShX sh -> [Int] -> Int# -> IxX sh i - fromLin ZSX _ _ = ZIX - 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" - shxCast :: StaticShX sh' -> IShX sh -> Maybe (IShX sh') shxCast ZKX ZSX = Just ZSX shxCast (SKnown m :!% ssh) (SKnown n :$% sh) | Just Refl <- testEquality n m = (SKnown n :$%) <$> shxCast ssh sh |