1 files changed, 20 insertions, 13 deletions

diff --git a/src/Data/Array/Nested/Internal/Mixed.hs b/src/Data/Array/Nested/Internal/Mixed.hs
index b155da5..d3e8088 100644
--- a/src/Data/Array/Nested/Internal/Mixed.hs
+++ b/src/Data/Array/Nested/Internal/Mixed.hs
@@ -20,7 +20,7 @@ module Data.Array.Nested.Internal.Mixed where
 
 import Prelude hiding (mconcat)
 
-import Control.DeepSeq (NFData)
+import Control.DeepSeq (NFData(..))
 import Control.Monad (forM_, when)
 import Control.Monad.ST
 import Data.Array.RankedS qualified as S
@@ -143,8 +143,6 @@ data instance Mixed sh (Primitive a) = M_Primitive !(IShX sh) !(XArray sh a)
 -- | Only on scalars, because lexicographical ordering is strange on multi-dimensional arrays.
 deriving instance (Ord a, Storable a) => Ord (Mixed sh (Primitive a))
 
-instance NFData a => NFData (Mixed sh (Primitive a))
-
 -- [PRIMITIVE ELEMENT TYPES LIST]
 newtype instance Mixed sh Bool = M_Bool (Mixed sh (Primitive Bool)) deriving (Eq, Generic)
 newtype instance Mixed sh Int = M_Int (Mixed sh (Primitive Int)) deriving (Eq, Generic)
@@ -157,21 +155,19 @@ newtype instance Mixed sh () = M_Nil (Mixed sh (Primitive ())) deriving (Eq, Gen
 -- etc.
 
 -- [PRIMITIVE ELEMENT TYPES LIST]
-deriving instance Ord (Mixed sh Bool) ; instance NFData (Mixed sh Bool)
-deriving instance Ord (Mixed sh Int) ; instance NFData (Mixed sh Int)
-deriving instance Ord (Mixed sh Int64) ; instance NFData (Mixed sh Int64)
-deriving instance Ord (Mixed sh Int32) ; instance NFData (Mixed sh Int32)
-deriving instance Ord (Mixed sh CInt) ; instance NFData (Mixed sh CInt)
-deriving instance Ord (Mixed sh Float) ; instance NFData (Mixed sh Float)
-deriving instance Ord (Mixed sh Double) ; instance NFData (Mixed sh Double)
-deriving instance Ord (Mixed sh ()) ; instance NFData (Mixed sh ())
+deriving instance Ord (Mixed sh Bool)
+deriving instance Ord (Mixed sh Int)
+deriving instance Ord (Mixed sh Int64)
+deriving instance Ord (Mixed sh Int32)
+deriving instance Ord (Mixed sh CInt)
+deriving instance Ord (Mixed sh Float)
+deriving instance Ord (Mixed sh Double)
+deriving instance Ord (Mixed sh ())
 
 data instance Mixed sh (a, b) = M_Tup2 !(Mixed sh a) !(Mixed sh b) deriving (Generic)
-instance (NFData (Mixed sh a), NFData (Mixed sh b)) => NFData (Mixed sh (a, b))
 -- etc., larger tuples (perhaps use generics to allow arbitrary product types)
 
 data instance Mixed sh1 (Mixed sh2 a) = M_Nest !(IShX sh1) !(Mixed (sh1 ++ sh2) a) deriving (Generic)
-instance NFData (Mixed (sh1 ++ sh2) a) => NFData (Mixed sh1 (Mixed sh2 a))
 
 
 -- | Internal helper data family mirroring 'Mixed' that consists of mutable
@@ -218,6 +214,9 @@ instance (Show a, Storable a) => Show (ShowViaPrimitive sh a) where
 
 deriving via (ShowViaToListLinear sh a) instance (Show a, Elt a) => Show (Mixed sh a)
 
+instance Elt a => NFData (Mixed sh a) where
+  rnf = mrnf
+
 
 mliftNumElt1 :: (PrimElt a, PrimElt b)
              => (SNat (Rank sh) -> S.Array (Rank sh) a -> S.Array (Rank sh) b)
@@ -330,6 +329,8 @@ class Elt a where
   -- inside their elements.
   mconcat :: NonEmpty (Mixed (Nothing : sh) a) -> Mixed (Nothing : sh) a
 
+  mrnf :: Mixed sh a -> ()
+
   -- ====== PRIVATE METHODS ====== --
 
   -- | Tree giving the shape of every array component.
@@ -432,6 +433,8 @@ instance Storable a => Elt (Primitive a) where
     let result = X.concat (ssxFromShape sh) (fmap (\(M_Primitive _ arr) -> arr) l)
     in M_Primitive (X.shape (SUnknown () :!% ssxFromShape sh) result) result
 
+  mrnf (M_Primitive sh a) = rnf sh `seq` rnf a
+
   type ShapeTree (Primitive a) = ()
   mshapeTree _ = ()
   mshapeTreeEq _ () () = True
@@ -503,6 +506,8 @@ instance (Elt a, Elt b) => Elt (a, b) where
         unzipT2 (M_Tup2 a b :| l) = let (l1, l2) = unzipT2l l in (a :| l1, b :| l2)
     in uncurry M_Tup2 . bimap mconcat mconcat . unzipT2
 
+  mrnf (M_Tup2 a b) = mrnf a `seq` mrnf b
+
   type ShapeTree (a, b) = (ShapeTree a, ShapeTree b)
   mshapeTree (x, y) = (mshapeTree x, mshapeTree y)
   mshapeTreeEq _ (t1, t2) (t1', t2') = mshapeTreeEq (Proxy @a) t1 t1' && mshapeTreeEq (Proxy @b) t2 t2'
@@ -627,6 +632,8 @@ instance Elt a => Elt (Mixed sh' a) where
     let result = mconcat (fmap (\(M_Nest _ arr) -> arr) l)
     in M_Nest (fst (shxSplitApp (Proxy @sh') (ssxFromShape sh1) (mshape result))) result
 
+  mrnf (M_Nest sh arr) = rnf sh `seq` mrnf arr
+
   type ShapeTree (Mixed sh' a) = (IShX sh', ShapeTree a)
 
   mshapeTree :: Mixed sh' a -> ShapeTree (Mixed sh' a)