Define ix?FromLinear without TH

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]]