diff options
Diffstat (limited to 'src/Data/Array/Nested/Mixed')
| -rw-r--r-- | src/Data/Array/Nested/Mixed/Shape/Internal.hs | 47 |
1 files changed, 27 insertions, 20 deletions
diff --git a/src/Data/Array/Nested/Mixed/Shape/Internal.hs b/src/Data/Array/Nested/Mixed/Shape/Internal.hs index cf44522..9997b0f 100644 --- a/src/Data/Array/Nested/Mixed/Shape/Internal.hs +++ b/src/Data/Array/Nested/Mixed/Shape/Internal.hs @@ -19,32 +19,39 @@ ixFromLinearStub fname' ishty ixty zshC consshC ixz ixcons shtolist = do typesig <- [t| forall i sh. Num i => $ishty sh -> Int -> $ixty sh i |] locals <- [d| - 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" + -- Unfold first iteration of fromLin to do the range check. + -- Don't inline because if this is inlined, GHC seems to stop sharing + -- 'suffixes' over multiple calls, which breaks performance in sh*Enum. + {-# NOINLINE 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" - {-# NOINLINE outrange #-} - outrange :: $ishty sh -> Int -> a - outrange sh i = error $ fname' ++ ": out of range (" ++ show i ++ - " in array of shape " ++ show sh ++ ")" |] + 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 \i -> - if i < 0 then outrange sh i else - case (sh, suffixes) of -- unfold first iteration of fromLin to do the range check - ($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" |] + in fromLin0 sh suffixes |] return [SigD fname typesig ,FunD fname [Clause [] (NormalB body) locals]] |