5 files changed, 23 insertions, 37 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/Mixed.hs b/src/Data/Array/Nested/Mixed.hs
index 6b96a15..54b2a9f 100644
--- a/src/Data/Array/Nested/Mixed.hs
+++ b/src/Data/Array/Nested/Mixed.hs
@@ -378,7 +378,9 @@ class Elt a where
   -- | Given the shape of this array, finalise the vectors into 'XArray's.
   mvecsFreeze :: IShX sh -> MixedVecs s sh a -> ST s (Mixed sh a)
 
-  -- | Given the shape of this array, finalise the vectors into 'XArray's.
+  -- | 'mvecsFreeze' but without copying the mutable vectors before freezing
+  -- them. If the mutable vectors are mutated after this function, referential
+  -- transparency is broken.
   mvecsUnsafeFreeze :: IShX sh -> MixedVecs s sh a -> ST s (Mixed sh a)
 
 -- | Element types for which we have evidence of the (static part of the) shape
diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs
index 11ef757..b1b4f81 100644
--- a/src/Data/Array/Nested/Mixed/Shape.hs
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -62,7 +62,7 @@ type role ListX nominal representational
 type ListX :: [Maybe Nat] -> (Maybe Nat -> Type) -> Type
 data ListX sh f where
   ZX :: ListX '[] f
-  (::%) :: f n -> ListX sh f -> ListX (n : sh) f
+  (::%) :: forall n sh {f}. f n -> ListX sh f -> ListX (n : sh) f
 deriving instance (forall n. Eq (f n)) => Eq (ListX sh f)
 deriving instance (forall n. Ord (f n)) => Ord (ListX sh f)
 infixr 3 ::%
diff --git a/src/Data/Array/Nested/Shaped.hs b/src/Data/Array/Nested/Shaped.hs
index 99ad590..85042f2 100644
--- a/src/Data/Array/Nested/Shaped.hs
+++ b/src/Data/Array/Nested/Shaped.hs
@@ -56,7 +56,7 @@ ssize = shsSize . sshape
 sindex :: Elt a => Shaped sh a -> IIxS sh -> a
 sindex (Shaped arr) idx = mindex arr (ixxFromIxS idx)
 
-shsTakeIx :: Proxy sh' -> ShS (sh ++ sh') -> IIxS sh -> ShS sh
+shsTakeIx :: Proxy sh' -> ShS (sh ++ sh') -> IxS sh i -> ShS sh
 shsTakeIx _ _ ZIS = ZSS
 shsTakeIx p sh (_ :.$ idx) = case sh of n :$$ sh' -> n :$$ shsTakeIx p sh' idx
 
@@ -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 f616946..18bd2e9 100644
--- a/src/Data/Array/Nested/Shaped/Shape.hs
+++ b/src/Data/Array/Nested/Shaped/Shape.hs
@@ -52,16 +52,14 @@ import Data.Array.Nested.Types
 
 -- * Shaped lists
 
--- | Note: The 'KnownNat' constraint on '(::$)' is deprecated and should be
--- removed in a future release.
 type role ListS nominal representational
 type ListS :: [Nat] -> (Nat -> Type) -> Type
 data ListS sh f where
   ZS :: ListS '[] f
-  -- TODO: when the KnownNat constraint is removed, restore listsIndex to sanity
-  (::$) :: forall n sh {f}. KnownNat n => f n -> ListS sh f -> ListS (n : sh) f
+  (::$) :: forall n sh {f}. f n -> ListS sh f -> ListS (n : sh) f
 deriving instance (forall n. Eq (f n)) => Eq (ListS sh f)
 deriving instance (forall n. Ord (f n)) => Ord (ListS sh f)
+
 infixr 3 ::$
 
 #ifdef OXAR_DEFAULT_SHOW_INSTANCES
@@ -76,7 +74,7 @@ instance (forall m. NFData (f m)) => NFData (ListS n f) where
   rnf (x ::$ l) = rnf x `seq` rnf l
 
 data UnconsListSRes f sh1 =
-  forall n sh. (KnownNat n, n : sh ~ sh1) => UnconsListSRes (ListS sh f) (f n)
+  forall n sh. (n : sh ~ sh1) => UnconsListSRes (ListS sh f) (f n)
 listsUncons :: ListS sh1 f -> Maybe (UnconsListSRes f sh1)
 listsUncons (x ::$ sh') = Just (UnconsListSRes sh' x)
 listsUncons ZS = Nothing
@@ -188,11 +186,11 @@ listsPermute :: forall f is sh. Perm is -> ListS sh f -> ListS (Permute is sh) f
 listsPermute PNil _ = ZS
 listsPermute (i `PCons` (is :: Perm is')) (sh :: ListS sh f) =
   case listsIndex (Proxy @is') (Proxy @sh) i sh of
-    (item, SNat) -> item ::$ listsPermute is sh
+    item -> item ::$ listsPermute is sh
 
--- TODO: remove this SNat when the KnownNat constaint in ListS is removed
-listsIndex :: forall f i is sh shT. Proxy is -> Proxy shT -> SNat i -> ListS sh f -> (f (Index i sh), SNat (Index i sh))
-listsIndex _ _ SZ (n ::$ _) = (n, SNat)
+-- 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))
   | Refl <- lemIndexSucc (Proxy @i') (Proxy @n) (Proxy @sh')
   = listsIndex p pT i sh
@@ -216,7 +214,7 @@ pattern ZIS = IxS ZS
 -- removed in a future release.
 pattern (:.$)
   :: forall {sh1} {i}.
-     forall n sh. (KnownNat n, n : sh ~ sh1)
+     forall n sh. (n : sh ~ sh1)
   => i -> IxS sh i -> IxS sh1 i
 pattern i :.$ shl <- IxS (listsUncons -> Just (UnconsListSRes (IxS -> shl) (getConst -> i)))
   where i :.$ IxS shl = IxS (Const i ::$ shl)
@@ -280,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))
@@ -331,7 +327,7 @@ pattern ZSS = ShS ZS
 
 pattern (:$$)
   :: forall {sh1}.
-     forall n sh. (KnownNat n, n : sh ~ sh1)
+     forall n sh. (n : sh ~ sh1)
   => SNat n -> ShS sh -> ShS sh1
 pattern i :$$ shl <- ShS (listsUncons -> Just (UnconsListSRes (ShS -> shl) i))
   where i :$$ ShS shl = ShS (i ::$ shl)
@@ -414,7 +410,7 @@ shsPermute :: Perm is -> ShS sh -> ShS (Permute is sh)
 shsPermute = coerce (listsPermute @SNat)
 
 shsIndex :: Proxy is -> Proxy shT -> SNat i -> ShS sh -> SNat (Index i sh)
-shsIndex pis pshT i sh = coerce (fst (listsIndex @SNat pis pshT i (coerce sh)))
+shsIndex pis pshT i sh = coerce (listsIndex @SNat pis pshT i (coerce sh))
 
 shsPermutePrefix :: forall is sh. Perm is -> ShS sh -> ShS (PermutePrefix is sh)
 shsPermutePrefix = coerce (listsPermutePrefix @SNat)
@@ -445,17 +441,6 @@ shsOrthotopeShape :: ShS sh -> Dict O.Shape sh
 shsOrthotopeShape ZSS = Dict
 shsOrthotopeShape (SNat :$$ sh) | Dict <- shsOrthotopeShape sh = Dict
 
--- | This function is a hack made possible by the 'KnownNat' inside 'ListS'.
--- This function may be removed in a future release.
-shsFromListS :: ListS sh f -> ShS sh
-shsFromListS ZS = ZSS
-shsFromListS (_ ::$ l) = SNat :$$ shsFromListS l
-
--- | This function is a hack made possible by the 'KnownNat' inside 'IxS'. This
--- function may be removed in a future release.
-shsFromIxS :: IxS sh i -> ShS sh
-shsFromIxS (IxS l) = shsFromListS l
-
 shsEnum :: ShS sh -> [IIxS sh]
 shsEnum = shsEnum'