NFData instances

1 files changed, 23 insertions, 3 deletions

diff --git a/src/Data/Array/Mixed.hs b/src/Data/Array/Mixed.hs
index d894b85..1e8cee2 100644
--- a/src/Data/Array/Mixed.hs
+++ b/src/Data/Array/Mixed.hs
@@ -2,6 +2,7 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DeriveFoldable #-}
 {-# LANGUAGE DeriveFunctor #-}
+{-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE DerivingStrategies #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTs #-}
@@ -25,6 +26,7 @@
 {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
 module Data.Array.Mixed where
 
+import Control.DeepSeq (NFData(..))
 import qualified Data.Array.RankedS as S
 import qualified Data.Array.Ranked as ORB
 import Data.Bifunctor (first)
@@ -38,6 +40,7 @@ import Data.Type.Bool
 import Data.Type.Equality
 import qualified Data.Vector.Storable as VS
 import Foreign.Storable (Storable)
+import GHC.Generics (Generic)
 import GHC.IsList (IsList)
 import qualified GHC.IsList as IsList
 import GHC.TypeError
@@ -106,6 +109,10 @@ infixr 3 ::%
 instance (forall n. Show (f n)) => Show (ListX sh f) where
   showsPrec _ = showListX shows
 
+instance (forall n. NFData (f n)) => NFData (ListX sh f) where
+  rnf ZX = ()
+  rnf (x ::% l) = rnf x `seq` rnf l
+
 data UnconsListXRes f sh1 =
   forall n sh. (n : sh ~ sh1) => UnconsListXRes (ListX sh f) (f n)
 unconsListX :: ListX sh1 f -> Maybe (UnconsListXRes f sh1)
@@ -142,7 +149,7 @@ listXToList (Const i ::% is) = i : listXToList is
 type role IxX nominal representational
 type IxX :: [Maybe Nat] -> Type -> Type
 newtype IxX sh i = IxX (ListX sh (Const i))
-  deriving (Eq, Ord)
+  deriving (Eq, Ord, Generic)
 
 pattern ZIX :: forall sh i. () => sh ~ '[] => IxX sh i
 pattern ZIX = IxX ZX
@@ -168,6 +175,8 @@ instance Functor (IxX sh) where
 instance Foldable (IxX sh) where
   foldMap f (IxX l) = foldListX (f . getConst) l
 
+instance NFData i => NFData (IxX sh i)
+
 
 data SMayNat i f n where
   SUnknown :: i -> SMayNat i f Nothing
@@ -176,6 +185,10 @@ deriving instance (Show i, forall m. Show (f m)) => Show (SMayNat i f n)
 deriving instance (Eq i, forall m. Eq (f m)) => Eq (SMayNat i f n)
 deriving instance (Ord i, forall m. Ord (f m)) => Ord (SMayNat i f n)
 
+instance (NFData i, forall m. NFData (f m)) => NFData (SMayNat i f n) where
+  rnf (SUnknown i) = rnf i
+  rnf (SKnown x) = rnf x
+
 fromSMayNat :: (n ~ Nothing => i -> r) -> (forall m. n ~ Just m => f m -> r) -> SMayNat i f n -> r
 fromSMayNat f _ (SUnknown i) = f i
 fromSMayNat _ g (SKnown s) = g s
@@ -186,7 +199,7 @@ fromSMayNat' = fromSMayNat id fromSNat'
 type role ShX nominal representational
 type ShX :: [Maybe Nat] -> Type -> Type
 newtype ShX sh i = ShX (ListX sh (SMayNat i SNat))
-  deriving (Eq, Ord)
+  deriving (Eq, Ord, Generic)
 
 pattern ZSX :: forall sh i. () => sh ~ '[] => ShX sh i
 pattern ZSX = ShX ZX
@@ -209,6 +222,11 @@ instance Show i => Show (ShX sh i) where
 instance Functor (ShX sh) where
   fmap f (ShX l) = ShX (fmapListX (fromSMayNat (SUnknown . f) SKnown) l)
 
+instance NFData i => NFData (ShX sh i) where
+  rnf (ShX ZX) = ()
+  rnf (ShX (SUnknown i ::% l)) = rnf i `seq` rnf (ShX l)
+  rnf (ShX (SKnown SNat ::% l)) = rnf (ShX l)
+
 lengthShX :: ShX sh i -> Int
 lengthShX (ShX l) = lengthListX l
 
@@ -294,11 +312,13 @@ type family Rank sh where
 
 type XArray :: [Maybe Nat] -> Type -> Type
 newtype XArray sh a = XArray (S.Array (Rank sh) a)
-  deriving (Show, Eq)
+  deriving (Show, Eq, Generic)
 
 -- | Only on scalars, because lexicographical ordering is strange on multi-dimensional arrays.
 deriving instance (Ord a, Storable a) => Ord (XArray '[] a)
 
+instance NFData a => NFData (XArray sh a)
+
 zeroIxX :: StaticShX sh -> IIxX sh
 zeroIxX ZKX = ZIX
 zeroIxX (_ :!% ssh) = 0 :.% zeroIxX ssh