diff options
Diffstat (limited to 'src/Data/Array/Nested/Mixed')
| -rw-r--r-- | src/Data/Array/Nested/Mixed/Shape.hs | 41 | ||||
| -rw-r--r-- | src/Data/Array/Nested/Mixed/Shape/Internal.hs | 59 |
2 files changed, 35 insertions, 65 deletions
diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs index 5ffd40c..ebf0a07 100644 --- a/src/Data/Array/Nested/Mixed/Shape.hs +++ b/src/Data/Array/Nested/Mixed/Shape.hs @@ -46,7 +46,6 @@ import GHC.TypeLits import GHC.TypeLits.Orphans () #endif -import Data.Array.Nested.Mixed.Shape.Internal import Data.Array.Nested.Types @@ -293,6 +292,41 @@ ixxToLinear = \sh i -> go sh i 0 go ZSX ZIX a = a go (n :$% sh) (i :.% ix) a = go sh ix (fromIntegral (fromSMayNat' n) * a + i) +{-# INLINEABLE ixxFromLinear #-} +ixxFromLinear :: Num i => IShX sh -> Int -> IxX sh i +ixxFromLinear = \sh -> -- give this function arity 1 so that suffixes is shared when it's called many times + let suffixes = drop 1 (scanr (*) 1 (shxToList sh)) + in fromLin0 sh suffixes + where + -- Unfold first iteration of fromLin to do the range check. + -- Don't inline this function at first to allow GHC to inline the outer + -- function and realise that 'suffixes' is shared. But then later inline it + -- anyway, to avoid the function call. Removing the pragma makes GHC + -- somehow unable to recognise that 'suffixes' can be shared in a loop. + {-# NOINLINE [0] fromLin0 #-} + fromLin0 :: Num i => IShX sh -> [Int] -> Int -> IxX sh i + fromLin0 sh suffixes i = + if i < 0 then outrange sh i else + case (sh, suffixes) of + (ZSX, _) | i > 0 -> outrange sh i + | otherwise -> ZIX + ((fromSMayNat' -> n) :$% sh', suff : suffs) -> + let (q, r) = i `quotRem` suff + in if q >= n then outrange sh i else + fromIntegral q :.% fromLin sh' suffs r + _ -> error "impossible" + + fromLin :: Num i => IShX sh -> [Int] -> Int -> IxX sh i + fromLin ZSX _ !_ = ZIX + fromLin (_ :$% sh') (suff : suffs) i = + let (q, r) = i `quotRem` suff -- suff == shrSize sh' + in fromIntegral q :.% fromLin sh' suffs r + fromLin _ _ _ = error "impossible" + + {-# NOINLINE outrange #-} + outrange :: IShX sh -> Int -> a + outrange sh i = error $ "ixxFromLinear: out of range (" ++ show i ++ + " in array of shape " ++ show sh ++ ")" -- * Mixed shape-like lists to be used for ShX and StaticShX @@ -798,8 +832,3 @@ instance KnownShX sh => IsList (IShX sh) where type Item (IShX sh) = Int fromList = shxFromList (knownShX @sh) toList = shxToList - --- This needs to be at the bottom of the file to not split the file into --- pieces; some of the shape/index stuff refers to StaticShX. -$(ixFromLinearStub "ixxFromLinear" [t| IShX |] [t| IxX |] [p| ZSX |] (\a b -> [p| (fromSMayNat' -> $a) :$% $b |]) [| ZIX |] [| (:.%) |] [| shxToList |]) -{-# INLINEABLE ixxFromLinear #-} diff --git a/src/Data/Array/Nested/Mixed/Shape/Internal.hs b/src/Data/Array/Nested/Mixed/Shape/Internal.hs deleted file mode 100644 index 2a86ac1..0000000 --- a/src/Data/Array/Nested/Mixed/Shape/Internal.hs +++ /dev/null @@ -1,59 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -module Data.Array.Nested.Mixed.Shape.Internal where - -import Language.Haskell.TH - - --- | A TH stub function to avoid having to write the same code three times for --- the three kinds of shapes. -ixFromLinearStub :: String - -> TypeQ -> TypeQ - -> PatQ -> (PatQ -> PatQ -> PatQ) - -> ExpQ -> ExpQ - -> ExpQ - -> DecsQ -ixFromLinearStub fname' ishty ixty zshC consshC ixz ixcons shtolist = do - let fname = mkName fname' - typesig <- [t| forall i sh. Num i => $ishty sh -> Int -> $ixty sh i |] - - locals <- [d| - -- Unfold first iteration of fromLin to do the range check. - -- Don't inline this function at first to allow GHC to inline the outer - -- function and realise that 'suffixes' is shared. But then later inline it - -- anyway, to avoid the function call. Removing the pragma makes GHC - -- somehow unable to recognise that 'suffixes' can be shared in a loop. - {-# NOINLINE [0] fromLin0 #-} - fromLin0 :: Num i => $ishty sh -> [Int] -> Int -> $ixty sh i - fromLin0 sh suffixes i = - if i < 0 then outrange sh i else - case (sh, suffixes) of - ($zshC, _) | i > 0 -> outrange sh i - | otherwise -> $ixz - ($(consshC (varP (mkName "n")) (varP (mkName "sh'"))), suff : suffs) -> - let (q, r) = i `quotRem` suff - in if q >= n then outrange sh i else - $ixcons (fromIntegral q) (fromLin sh' suffs r) - _ -> error "impossible" - - fromLin :: Num i => $ishty sh -> [Int] -> Int -> $ixty sh i - fromLin $zshC _ !_ = $ixz - fromLin ($(consshC wildP (varP (mkName "sh'")))) (suff : suffs) i = - let (q, r) = i `quotRem` suff -- suff == shrSize sh' - in $ixcons (fromIntegral q) (fromLin sh' suffs r) - fromLin _ _ _ = error "impossible" - - {-# NOINLINE outrange #-} - outrange :: $ishty sh -> Int -> a - outrange sh i = error $ fname' ++ ": out of range (" ++ show i ++ - " in array of shape " ++ show sh ++ ")" |] - - body <- [| - \sh -> -- give this function arity 1 so that 'suffixes' is shared when - -- it's called many times - let suffixes = drop 1 (scanr (*) 1 ($shtolist sh)) - in fromLin0 sh suffixes |] - - return [SigD fname typesig - ,FunD fname [Clause [] (NormalB body) locals]] |