1 files changed, 46 insertions, 9 deletions

diff --git a/src/Data/Array/Nested/Convert.hs b/src/Data/Array/Nested/Convert.hs
index 861bf20..fd59ba6 100644
--- a/src/Data/Array/Nested/Convert.hs
+++ b/src/Data/Array/Nested/Convert.hs
@@ -7,11 +7,18 @@
 {-# LANGUAGE TypeAbstractions #-}
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeOperators #-}
+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
 module Data.Array.Nested.Convert (
-  -- * Shape/index/list casting functions
+  -- * Shape\/index\/list casting functions
+  -- ** To ranked
   ixrFromIxS, ixrFromIxX, shrFromShS, shrFromShX, shrFromShX2,
-  ixsFromIxX, shsFromShX,
+  listrCast, ixrCast, shrCast,
+  -- ** To shaped
+  ixsFromIxR, ixsFromIxR', ixsFromIxX, ixsFromIxX', shsFromShX, shsFromSSX,
+  ixsCast,
+  -- ** To mixed
   ixxFromIxR, ixxFromIxS, shxFromShR, shxFromShS,
+  ixxCast, shxCast, shxCast',
 
   -- * Array conversions
   convert,
@@ -57,24 +64,50 @@ shrFromShS ZSS = ZSR
 shrFromShS (n :$$ sh) = fromSNat' n :$: shrFromShS sh
 
 -- shrFromShX re-exported
-
 -- shrFromShX2 re-exported
+-- listrCast re-exported
+-- ixrCast re-exported
+-- shrCast re-exported
 
 -- * To shaped
 
--- ixsFromIxR :: IIxR (Rank sh) -> IIxS sh
--- ixsFromIxR = \ix -> go ix _
---   where
---     go :: IIxR n -> (forall sh. KnownShS sh => IIxS sh -> r) -> r
---     go ZIR k = k ZIS
---     go (i :.: ix) k = go ix (i :.$)
+-- TODO: these take a ShS because there are KnownNats inside IxS.
+
+ixsFromIxR :: ShS sh -> IxR (Rank sh) i -> IxS sh i
+ixsFromIxR ZSS ZIR = ZIS
+ixsFromIxR (_ :$$ sh) (n :.: idx) = n :.$ ixsFromIxR sh idx
+ixsFromIxR _ _ = error "unreachable"
 
+-- | Performs a runtime check that @n@ matches @Rank sh@. Equivalent to the
+-- following, but more efficient:
+--
+-- > ixsFromIxR' sh idx = ixsFromIxR sh (ixrCast (shsRank sh) idx)
+ixsFromIxR' :: ShS sh -> IxR n i -> IxS sh i
+ixsFromIxR' ZSS ZIR = ZIS
+ixsFromIxR' (_ :$$ sh) (n :.: idx) = n :.$ ixsFromIxR' sh idx
+ixsFromIxR' _ _ = error "ixsFromIxR': index rank does not match shape rank"
+
+-- TODO: this takes a ShS because there are KnownNats inside IxS.
 ixsFromIxX :: ShS sh -> IxX (MapJust sh) i -> IxS sh i
 ixsFromIxX ZSS ZIX = ZIS
 ixsFromIxX (_ :$$ sh) (n :.% idx) = n :.$ ixsFromIxX sh idx
 
+-- | Performs a runtime check that @Rank sh'@ match @Rank sh@. Equivalent to
+-- the following, but more efficient:
+--
+-- > ixsFromIxX' sh idx = ixsFromIxX sh (ixxCast (shxFromShS sh) idx)
+ixsFromIxX' :: ShS sh -> IxX sh' i -> IxS sh i
+ixsFromIxX' ZSS ZIX = ZIS
+ixsFromIxX' (_ :$$ sh) (n :.% idx) = n :.$ ixsFromIxX' sh idx
+ixsFromIxX' _ _ = error "ixsFromIxX': index rank does not match shape rank"
+
 -- shsFromShX re-exported
 
+shsFromSSX :: StaticShX (MapJust sh) -> ShS sh
+shsFromSSX = shsFromShX Prelude.. shxFromSSX
+
+-- ixsCast re-exported
+
 -- * To mixed
 
 ixxFromIxR :: IxR n i -> IxX (Replicate n Nothing) i
@@ -97,6 +130,10 @@ shxFromShS :: ShS sh -> IShX (MapJust sh)
 shxFromShS ZSS = ZSX
 shxFromShS (n :$$ sh) = SKnown n :$% shxFromShS sh
 
+-- ixxCast re-exported
+-- shxCast re-exported
+-- shxCast' re-exported
+
 
 -- * Array conversions