De-triplicate ix?ToLinear

3 files changed, 10 insertions, 18 deletions

diff --git a/src/Data/Array/Nested/Convert.hs b/src/Data/Array/Nested/Convert.hs
index 91752c4..86534f1 100644
--- a/src/Data/Array/Nested/Convert.hs
+++ b/src/Data/Array/Nested/Convert.hs
@@ -141,11 +141,8 @@ shsFromSSX = shsFromShX Prelude.. shxFromSSX
 
 -- * To mixed
 
-ixxFromIxR :: IxR n i -> IxX (Replicate n Nothing) i
-ixxFromIxR = unsafeCoerce
-
-ixxFromIxS :: IxS sh i -> IxX (MapJust sh) i
-ixxFromIxS = unsafeCoerce
+-- ixxFromIxR re-exported
+-- ixxFromIxS re-exported
 
 shxFromShR :: ShR n i -> ShX (Replicate n Nothing) i
 shxFromShR = coerce
diff --git a/src/Data/Array/Nested/Ranked/Shape.hs b/src/Data/Array/Nested/Ranked/Shape.hs
index 0ac980e..74f4c5c 100644
--- a/src/Data/Array/Nested/Ranked/Shape.hs
+++ b/src/Data/Array/Nested/Ranked/Shape.hs
@@ -290,14 +290,10 @@ ixrPermutePrefix = coerce (listrPermutePrefix @i)
 -- index into the buffer.
 {-# INLINEABLE ixrToLinear #-}
 ixrToLinear :: Num i => IShR m -> IxR m i -> i
-ixrToLinear = \sh i -> go sh i 0
-  where
-    -- Additional argument: index, in the @m - m1@ dimensional array so far,
-    -- of the @m - m1 + n@ dimensional tensor pointed to by the current
-    -- @m - m1@ dimensional index prefix.
-    go :: Num i => IShR m1 -> IxR m1 i -> i -> i
-    go ZSR ZIR a = a
-    go (n :$: sh) (i :.: ix) a = go sh ix (fromIntegral n * a + i)
+ixrToLinear (ShR sh) ix = ixxToLinear sh (ixxFromIxR ix)
+
+ixxFromIxR :: IxR n i -> IxX (Replicate n Nothing) i
+ixxFromIxR = unsafeCoerce
 
 {-# INLINEABLE ixrFromLinear #-}
 ixrFromLinear :: forall i m. Num i => IShR m -> Int -> IxR m i
diff --git a/src/Data/Array/Nested/Shaped/Shape.hs b/src/Data/Array/Nested/Shaped/Shape.hs
index 39be729..634c854 100644
--- a/src/Data/Array/Nested/Shaped/Shape.hs
+++ b/src/Data/Array/Nested/Shaped/Shape.hs
@@ -311,11 +311,10 @@ ixsPermutePrefix = coerce (listsPermutePrefix @(Const i))
 -- index into the buffer.
 {-# INLINEABLE ixsToLinear #-}
 ixsToLinear :: Num i => ShS sh -> IxS sh i -> i
-ixsToLinear = \sh i -> go sh i 0
-  where
-    go :: Num i => ShS sh -> IxS sh i -> i -> i
-    go ZSS ZIS a = a
-    go (n :$$ sh) (i :.$ ix) a = go sh ix (fromIntegral (fromSNat' n) * a + i)
+ixsToLinear (ShS sh) ix = ixxToLinear sh (ixxFromIxS ix)
+
+ixxFromIxS :: IxS sh i -> IxX (MapJust sh) i
+ixxFromIxS = unsafeCoerce
 
 {-# INLINEABLE ixsFromLinear #-}
 ixsFromLinear :: Num i => ShS sh -> Int -> IxS sh i