Introduce EltRepr (no more Primitive/Coercible in API)

1 files changed, 59 insertions, 24 deletions

diff --git a/src/Data/Array/Nested/Internal.hs b/src/Data/Array/Nested/Internal.hs
index 7a5add7..29592ac 100644
--- a/src/Data/Array/Nested/Internal.hs
+++ b/src/Data/Array/Nested/Internal.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DefaultSignatures #-}
 {-# LANGUAGE DeriveFunctor #-}
 {-# LANGUAGE DerivingVia #-}
 {-# LANGUAGE FlexibleContexts #-}
@@ -132,6 +133,23 @@ shAppSplit p (_ :!$? ssh) (i :$? idx) = first (i :$?) (shAppSplit p ssh idx)
 -- @'Primitive' T@ instead.
 newtype Primitive a = Primitive a
 
+-- | Element types that are primitive; arrays of these types are just a newtype
+-- wrapper over an array.
+class PrimElt a where
+  fromPrimitive :: Mixed sh (Primitive a) -> Mixed sh a
+  toPrimitive :: Mixed sh a -> Mixed sh (Primitive a)
+
+  default fromPrimitive :: Coercible (Mixed sh a) (Mixed sh (Primitive a)) => Mixed sh (Primitive a) -> Mixed sh a
+  fromPrimitive = coerce
+
+  default toPrimitive :: Coercible (Mixed sh (Primitive a)) (Mixed sh a) => Mixed sh a -> Mixed sh (Primitive a)
+  toPrimitive = coerce
+
+-- [PRIMITIVE ELEMENT TYPES LIST]
+instance PrimElt Int
+instance PrimElt Double
+instance PrimElt ()
+
 
 -- | Mixed arrays: some dimensions are size-typed, some are not. Distributes
 -- over product-typed elements using a data family so that the full array is
@@ -488,8 +506,11 @@ mappend = mlift2 go
            => Proxy sh' -> XArray (n : sh ++ sh') b -> XArray (m : sh ++ sh') b -> XArray (X.AddMaybe n m : sh ++ sh') b
         go Proxy | Dict <- X.lemAppKnownShapeX (knownShapeX @sh) (knownShapeX @sh') = X.append
 
-mfromVector :: forall sh a. (KnownShapeX sh, Storable a) => IShX sh -> VS.Vector a -> Mixed sh (Primitive a)
-mfromVector sh v = M_Primitive (X.fromVector sh v)
+mfromVectorP :: forall sh a. (KnownShapeX sh, Storable a) => IShX sh -> VS.Vector a -> Mixed sh (Primitive a)
+mfromVectorP sh v = M_Primitive (X.fromVector sh v)
+
+mfromVector :: forall sh a. (KnownShapeX sh, Storable a, PrimElt a) => IShX sh -> VS.Vector a -> Mixed sh a
+mfromVector sh v = fromPrimitive (mfromVectorP sh v)
 
 mfromList1 :: (KnownShapeX '[n], Elt a) => NonEmpty a -> Mixed '[n] a
 mfromList1 = mfromList . fmap mscalar
@@ -503,11 +524,9 @@ munScalar arr = mindex arr ZIX
 mconstantP :: forall sh a. (KnownShapeX sh, Storable a) => IShX sh -> a -> Mixed sh (Primitive a)
 mconstantP sh x = M_Primitive (X.constant sh x)
 
--- | This 'Coercible' constraint holds for the scalar types for which 'Mixed'
--- is defined.
-mconstant :: forall sh a. (KnownShapeX sh, Storable a, Coercible (Mixed sh (Primitive a)) (Mixed sh a))
+mconstant :: forall sh a. (KnownShapeX sh, Storable a, PrimElt a)
           => IShX sh -> a -> Mixed sh a
-mconstant sh x = coerce (mconstantP sh x)
+mconstant sh x = fromPrimitive (mconstantP sh x)
 
 mslice :: (KnownShapeX sh, Elt a) => [(Int, Int)] -> Mixed sh a -> Mixed sh a
 mslice ivs = mlift $ \_ -> X.slice ivs
@@ -981,10 +1000,10 @@ rlift f (Ranked arr)
   | Dict <- lemKnownReplicate (Proxy @n2)
   = Ranked (mlift f arr)
 
-rsumOuter1 :: forall n a.
-              (Storable a, Num a, KnownINat n)
-           => Ranked (S n) (Primitive a) -> Ranked n (Primitive a)
-rsumOuter1 (Ranked arr)
+rsumOuter1P :: forall n a.
+               (Storable a, Num a, KnownINat n)
+            => Ranked (S n) (Primitive a) -> Ranked n (Primitive a)
+rsumOuter1P (Ranked arr)
   | Dict <- lemKnownReplicate (Proxy @n)
   = Ranked
     . coerce @(XArray (Replicate n 'Nothing) a) @(Mixed (Replicate n 'Nothing) (Primitive a))
@@ -992,6 +1011,11 @@ rsumOuter1 (Ranked arr)
     . coerce @(Mixed (Replicate (S n) Nothing) (Primitive a)) @(XArray (Replicate (S n) Nothing) a)
     $ arr
 
+rsumOuter1 :: forall n a.
+              (Storable a, Num a, PrimElt a, KnownINat n)
+           => Ranked (S n) a -> Ranked n a
+rsumOuter1 = coerce fromPrimitive . rsumOuter1P @n @a . coerce toPrimitive
+
 rtranspose :: forall n a. (KnownINat n, Elt a) => [Int] -> Ranked n a -> Ranked n a
 rtranspose perm (Ranked arr)
   | Dict <- lemKnownReplicate (Proxy @n)
@@ -1004,10 +1028,13 @@ rappend | Dict <- lemKnownReplicate (Proxy @n) = coerce mappend
 rscalar :: Elt a => a -> Ranked I0 a
 rscalar x = Ranked (mscalar x)
 
-rfromVector :: forall n a. (KnownINat n, Storable a) => IShR n -> VS.Vector a -> Ranked n (Primitive a)
-rfromVector sh v
+rfromVectorP :: forall n a. (KnownINat n, Storable a) => IShR n -> VS.Vector a -> Ranked n (Primitive a)
+rfromVectorP sh v
   | Dict <- lemKnownReplicate (Proxy @n)
-  = Ranked (mfromVector (shCvtRX sh) v)
+  = Ranked (mfromVectorP (shCvtRX sh) v)
+
+rfromVector :: forall n a. (KnownINat n, Storable a, PrimElt a) => IShR n -> VS.Vector a -> Ranked n a
+rfromVector sh v = coerce fromPrimitive (rfromVectorP sh v)
 
 rfromList :: forall n a. (KnownINat n, Elt a) => NonEmpty (Ranked n a) -> Ranked (S n) a
 rfromList l
@@ -1031,9 +1058,9 @@ rconstantP sh x
   | Dict <- lemKnownReplicate (Proxy @n)
   = Ranked (mconstantP (shCvtRX sh) x)
 
-rconstant :: forall n a. (KnownINat n, Storable a, Coercible (Mixed (Replicate n Nothing) (Primitive a)) (Mixed (Replicate n Nothing) a))
+rconstant :: forall n a. (KnownINat n, Storable a, PrimElt a)
           => IShR n -> a -> Ranked n a
-rconstant sh x = coerce (rconstantP sh x)
+rconstant sh x = coerce fromPrimitive (rconstantP sh x)
 
 rslice :: (KnownINat n, Elt a) => [(Int, Int)] -> Ranked n a -> Ranked n a
 rslice ivs = rlift $ \_ -> X.slice ivs
@@ -1173,10 +1200,10 @@ slift f (Shaped arr)
   | Dict <- lemKnownMapJust (Proxy @sh2)
   = Shaped (mlift f arr)
 
-ssumOuter1 :: forall sh n a.
-              (Storable a, Num a, KnownNat n, KnownShape sh)
-           => Shaped (n : sh) (Primitive a) -> Shaped sh (Primitive a)
-ssumOuter1 (Shaped arr)
+ssumOuter1P :: forall sh n a.
+               (Storable a, Num a, KnownNat n, KnownShape sh)
+            => Shaped (n : sh) (Primitive a) -> Shaped sh (Primitive a)
+ssumOuter1P (Shaped arr)
   | Dict <- lemKnownMapJust (Proxy @sh)
   = Shaped
     . coerce @(XArray (MapJust sh) a) @(Mixed (MapJust sh) (Primitive a))
@@ -1184,6 +1211,11 @@ ssumOuter1 (Shaped arr)
     . coerce @(Mixed (Just n : MapJust sh) (Primitive a)) @(XArray (Just n : MapJust sh) a)
     $ arr
 
+ssumOuter1 :: forall sh n a.
+              (Storable a, Num a, PrimElt a, KnownNat n, KnownShape sh)
+           => Shaped (n : sh) a -> Shaped sh a
+ssumOuter1 = coerce fromPrimitive . ssumOuter1P @sh @n @a . coerce toPrimitive
+
 stranspose :: forall sh a. (KnownShape sh, Elt a) => [Int] -> Shaped sh a -> Shaped sh a
 stranspose perm (Shaped arr)
   | Dict <- lemKnownMapJust (Proxy @sh)
@@ -1196,10 +1228,13 @@ sappend | Dict <- lemKnownMapJust (Proxy @sh) = coerce mappend
 sscalar :: Elt a => a -> Shaped '[] a
 sscalar x = Shaped (mscalar x)
 
-sfromVector :: forall sh a. (KnownShape sh, Storable a) => VS.Vector a -> Shaped sh (Primitive a)
-sfromVector v
+sfromVectorP :: forall sh a. (KnownShape sh, Storable a) => VS.Vector a -> Shaped sh (Primitive a)
+sfromVectorP v
   | Dict <- lemKnownMapJust (Proxy @sh)
-  = Shaped (mfromVector (shCvtSX (knownShape @sh)) v)
+  = Shaped (mfromVectorP (shCvtSX (knownShape @sh)) v)
+
+sfromVector :: forall sh a. (KnownShape sh, Storable a, PrimElt a) => VS.Vector a -> Shaped sh a
+sfromVector v = coerce fromPrimitive (sfromVectorP @sh @a v)
 
 sfromList :: forall n sh a. (KnownNat n, KnownShape sh, Elt a)
           => NonEmpty (Shaped sh a) -> Shaped (n : sh) a
@@ -1224,9 +1259,9 @@ sconstantP x
   | Dict <- lemKnownMapJust (Proxy @sh)
   = Shaped (mconstantP (shCvtSX (knownShape @sh)) x)
 
-sconstant :: forall sh a. (KnownShape sh, Storable a, Coercible (Mixed (MapJust sh) (Primitive a)) (Mixed (MapJust sh) a))
+sconstant :: forall sh a. (KnownShape sh, Storable a, PrimElt a)
           => a -> Shaped sh a
-sconstant x = coerce (sconstantP @sh x)
+sconstant x = coerce fromPrimitive (sconstantP @sh x)
 
 sslice :: (KnownShape sh, Elt a) => [(Int, Int)] -> Shaped sh a -> Shaped sh a
 sslice ivs = slift $ \_ -> X.slice ivs