Merge pull request #291 from tweag/dest-arr/no-newtype

Changed destination type to have less unsafe linear casts

Author: Divesh Otwani

src/Data/Array/Destination.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE GADTs #-}
 {-# LANGUAGE LinearTypes #-}
 {-# LANGUAGE NoImplicitPrelude #-}
 {-# LANGUAGE ScopedTypeVariables #-}
@@ -127,41 +128,38 @@ module Data.Array.Destination
   )
   where
 
-import Control.Exception (evaluate)
 import Data.Vector (Vector, (!))
 import qualified Data.Vector as Vector
 import Data.Vector.Mutable (MVector)
 import qualified Data.Vector.Mutable as MVector
 import GHC.Exts (RealWorld)
 import qualified Prelude as Prelude
-import Prelude.Linear hiding (replicate)
-import System.IO.Unsafe
+import System.IO.Unsafe (unsafeDupablePerformIO)
 import GHC.Stack
+import Data.Unrestricted.Linear
+import Prelude.Linear hiding (replicate)
 import qualified Unsafe.Linear as Unsafe
 
 -- | A destination array, or @DArray@, is a write-only array that is filled
 -- by some computation which ultimately returns an array.
-newtype DArray a = DArray (MVector RealWorld a)
+data DArray a where
+  DArray :: MVector RealWorld a -> DArray a
 
 -- XXX: use of Vector in types is temporary. I will probably move away from
 -- vectors and implement most stuff in terms of Array# and MutableArray#
 -- eventually, anyway. This would allow to move the MutableArray logic to
 -- linear IO, possibly, and segregate the unsafe casts to the Linear IO
 -- module.  @`alloc` n k@ must be called with a non-negative value of @n@.
 alloc :: Int -> (DArray a %1-> ()) %1-> Vector a
-alloc n = Unsafe.toLinear unsafeAlloc
-  where
-    unsafeAlloc :: (DArray a %1-> ()) -> Vector a
-    unsafeAlloc build = unsafeDupablePerformIO Prelude.$ do
-      dest <- MVector.unsafeNew n
-      evaluate (build (DArray dest))
-      Vector.unsafeFreeze dest
+alloc n writer = (\(Ur dest, vec) -> writer (DArray dest) `lseq` vec) $
+  unsafeDupablePerformIO Prelude.$ do
+    destArray <- MVector.unsafeNew n
+    vec <- Vector.unsafeFreeze destArray
+    Prelude.return (Ur destArray, vec)
 
 -- | Get the size of a destination array.
 size :: DArray a %1-> (Ur Int, DArray a)
-size (DArray vec) = Unsafe.toLinear go vec
- where
-  go vec' = (Ur (MVector.length vec'), DArray vec')
+size (DArray mvec) = (Ur (MVector.length mvec), DArray mvec)
 
 -- | Fill a destination array with a constant
 replicate :: a -> DArray a %1-> ()
@@ -170,34 +168,25 @@ replicate a = fromFunction (const a)
 -- | @fill a dest@ fills a singleton destination array.
 -- Caution, @'fill' a dest@ will fail is @dest@ isn't of length exactly one.
 fill :: HasCallStack => a %1-> DArray a %1-> ()
-fill = Unsafe.toLinear2 unsafeFill
-    -- XXX: we will probably be able to spare this unsafe cast given a
-    -- (linear) length function on destination.
-  where
-    unsafeFill a (DArray ds) =
-      if MVector.length ds /= 1 then
-        error "Destination.fill: requires a destination of size 1"
-      else
-        unsafeDupablePerformIO Prelude.$ MVector.write ds 0 a
+fill a (DArray mvec) =
+  if MVector.length mvec /= 1
+  then error "Destination.fill: requires a destination of size 1" $ a
+  else a &
+    Unsafe.toLinear (\x -> unsafeDupablePerformIO (MVector.write mvec 0 x))
 
 -- | @dropEmpty dest@ consumes and empty array and fails otherwise.
 dropEmpty :: HasCallStack => DArray a %1-> ()
-dropEmpty = Unsafe.toLinear unsafeDrop where
-  unsafeDrop :: DArray a -> ()
-  unsafeDrop (DArray ds)
-    | MVector.length ds > 0 = error "Destination.dropEmpty on non-empty array."
-    | otherwise = ds `seq` ()
+dropEmpty (DArray mvec)
+  | MVector.length mvec > 0 = error "Destination.dropEmpty on non-empty array."
+  | otherwise = mvec `seq` ()
 
 -- | @'split' n dest = (destl, destr)@ such as @destl@ has length @n@.
 --
 -- 'split' is total: if @n@ is larger than the length of @dest@, then
 -- @destr@ is empty.
 split :: Int -> DArray a %1-> (DArray a, DArray a)
-split n = Unsafe.toLinear unsafeSplit
-  where
-    unsafeSplit (DArray ds) =
-      let (dsl, dsr) = MVector.splitAt n ds in
-        (DArray dsl, DArray dsr)
+split n (DArray mvec) | (ml, mr) <- MVector.splitAt n mvec =
+  (DArray ml, DArray mr)
 
 -- | Fills the destination array with the contents of given vector.
 --
@@ -211,10 +200,10 @@ mirror v f arr =
 
 -- | Fill a destination array using the given index-to-value function.
 fromFunction :: (Int -> b) -> DArray b %1-> ()
-fromFunction f = Unsafe.toLinear unsafeFromFunction
-  where unsafeFromFunction (DArray ds) = unsafeDupablePerformIO Prelude.$ do
-          let n = MVector.length ds
-          Prelude.sequence_ [MVector.unsafeWrite ds m (f m) | m <- [0..n-1]]
--- The unsafe cast here is actually safe, since getting the length does not
+fromFunction f (DArray mvec) = unsafeDupablePerformIO Prelude.$ do
+  let n = MVector.length mvec
+  Prelude.sequence_ [MVector.unsafeWrite mvec m (f m) | m <- [0..n-1]]
+-- The use of the mutable array is linear, since getting the length does not
 -- touch any elements, and each write fills in exactly one slot, so
 -- each slot of the destination array is filled.
+