Fix a few KnownNat in ShS-related TODOs and reword the restmvecsReplicate

3 files changed, 9 insertions, 12 deletions

diff --git a/src/Data/Array/Nested/Convert.hs b/src/Data/Array/Nested/Convert.hs
index 8c88d23..3d0da37 100644
--- a/src/Data/Array/Nested/Convert.hs
+++ b/src/Data/Array/Nested/Convert.hs
@@ -75,12 +75,12 @@ shrFromShS (n :$$ sh) = fromSNat' n :$: shrFromShS sh
 
 -- * To shaped
 
--- TODO: these take a ShS because there are KnownNats inside IxS.
-
+-- TODO: remove the ShS now that no KnownNats is inside IxS.
 ixsFromIxR :: ShS sh -> IxR (Rank sh) i -> IxS sh i
 ixsFromIxR ZSS ZIR = ZIS
 ixsFromIxR (_ :$$ sh) (n :.: idx) = n :.$ ixsFromIxR sh idx
 
+-- TODO: if possible, remove the ShS now that no KnownNats is inside IxS.
 -- | Performs a runtime check that @n@ matches @Rank sh@. Equivalent to the
 -- following, but more efficient:
 --
@@ -90,11 +90,12 @@ 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.
+-- TODO: remove the ShS now that no KnownNats is inside IxS.
 ixsFromIxX :: ShS sh -> IxX (MapJust sh) i -> IxS sh i
 ixsFromIxX ZSS ZIX = ZIS
 ixsFromIxX (_ :$$ sh) (n :.% idx) = n :.$ ixsFromIxX sh idx
 
+-- TODO: if possible, remove the ShS now that no KnownNats is inside IxS.
 -- | Performs a runtime check that @Rank sh'@ match @Rank sh@. Equivalent to
 -- the following, but more efficient:
 --
diff --git a/src/Data/Array/Nested/Shaped.hs b/src/Data/Array/Nested/Shaped.hs
index d23a025..85042f2 100644
--- a/src/Data/Array/Nested/Shaped.hs
+++ b/src/Data/Array/Nested/Shaped.hs
@@ -246,9 +246,7 @@ sreshape :: (Elt a, Product sh ~ Product sh') => ShS sh' -> Shaped sh a -> Shape
 sreshape sh' (Shaped arr) = Shaped (mreshape (shxFromShS sh') arr)
 
 sflatten :: Elt a => Shaped sh a -> Shaped '[Product sh] a
-sflatten arr =
-  case shsProduct (sshape arr) of  -- TODO: simplify when removing the KnownNat stuff
-    n@SNat -> sreshape (n :$$ ZSS) arr
+sflatten arr = sreshape (shsProduct (sshape arr) :$$ ZSS) arr
 
 siota :: (Enum a, PrimElt a) => SNat n -> Shaped '[n] a
 siota sn = Shaped (miota sn)
diff --git a/src/Data/Array/Nested/Shaped/Shape.hs b/src/Data/Array/Nested/Shaped/Shape.hs
index bfc6ad2..18bd2e9 100644
--- a/src/Data/Array/Nested/Shaped/Shape.hs
+++ b/src/Data/Array/Nested/Shaped/Shape.hs
@@ -188,7 +188,7 @@ listsPermute (i `PCons` (is :: Perm is')) (sh :: ListS sh f) =
   case listsIndex (Proxy @is') (Proxy @sh) i sh of
     item -> item ::$ listsPermute is sh
 
--- TODO: remove this SNat when the KnownNat constaint in ListS is removed
+-- TODO: try to remove this SNat now that the KnownNat constraint in ListS is removed
 listsIndex :: forall f i is sh shT. Proxy is -> Proxy shT -> SNat i -> ListS sh f -> f (Index i sh)
 listsIndex _ _ SZ (n ::$ _) = n
 listsIndex p pT (SS (i :: SNat i')) ((_ :: f n) ::$ (sh :: ListS sh' f))
@@ -278,11 +278,9 @@ ixsInit (IxS list) = IxS (listsInit list)
 ixsLast :: IxS (n : sh) i -> i
 ixsLast (IxS list) = getConst (listsLast list)
 
--- TODO: this takes a ShS because there are KnownNats inside IxS.
-ixsCast :: ShS sh' -> IxS sh i -> IxS sh' i
-ixsCast ZSS ZIS = ZIS
-ixsCast (_ :$$ sh) (i :.$ idx) = i :.$ ixsCast sh idx
-ixsCast _ _ = error "ixsCast: ranks don't match"
+ixsCast :: IxS sh i -> IxS sh i
+ixsCast ZIS = ZIS
+ixsCast (i :.$ idx) = i :.$ ixsCast idx
 
 ixsAppend :: forall sh sh' i. IxS sh i -> IxS sh' i -> IxS (sh ++ sh') i
 ixsAppend = coerce (listsAppend @_ @(Const i))