aboutsummaryrefslogtreecommitdiff
path: root/src/Data/Array/Nested/Mixed/Shape.hs
diff options
context:
space:
mode:
Diffstat (limited to 'src/Data/Array/Nested/Mixed/Shape.hs')
-rw-r--r--src/Data/Array/Nested/Mixed/Shape.hs34
1 files changed, 7 insertions, 27 deletions
diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs
index 8aa5a77..5a45a09 100644
--- a/src/Data/Array/Nested/Mixed/Shape.hs
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -16,6 +16,7 @@
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE StrictData #-}
+{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
@@ -44,6 +45,7 @@ import GHC.TypeLits
import GHC.TypeLits.Orphans ()
#endif
+import Data.Array.Nested.Mixed.Shape.Internal
import Data.Array.Nested.Types
@@ -276,33 +278,6 @@ 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
-{-# 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 \i ->
- if i < 0 then outrange sh i else
- case (sh, suffixes) of -- unfold first iteration of fromLin to do the range check
- (ZSX, _) | i > 0 -> outrange sh i
- | otherwise -> ZIX
- (n :$% sh', suff : suffs) ->
- let (q, r) = i `quotRem` suff
- in if q >= fromSMayNat' n then outrange sh i else
- fromIntegral q :.% fromLin sh' suffs r
- _ -> error "impossible"
- where
- 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 ++ ")"
-
ixxToLinear :: IShX sh -> IIxX sh -> Int
ixxToLinear = \sh i -> fst (go sh i)
where
@@ -684,3 +659,8 @@ instance KnownShX sh => IsList (ShX sh Int) where
type Item (ShX sh Int) = 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 #-}