Merge pull request #473 from tweag/scope-with-movable

In scope-passing style: use a `Movable b` instead of `Ur b`

Author: aspiwack

src/Data/Array/Mutable/Linear/Internal.hs

@@ -62,11 +62,11 @@ data Array a = Array (Array# a)
 -- | Allocate a constant array given a size and an initial value
 -- The size must be non-negative, otherwise this errors.
 alloc ::
-  (HasCallStack) =>
+  (HasCallStack, Movable b) =>
   Int ->
   a ->
-  (Array a %1 -> Ur b) %1 ->
-  Ur b
+  (Array a %1 -> b) %1 ->
+  b
 alloc s x f
   | s < 0 =
       (error ("Array.alloc: negative size: " ++ show s) :: x %1 -> x)
@@ -89,11 +89,11 @@ allocBeside s x (Array orig)
 
 -- | Allocate an array from a list
 fromList ::
-  (HasCallStack) =>
+  (HasCallStack, Movable b) =>
   [a] ->
-  (Array a %1 -> Ur b) %1 ->
-  Ur b
-fromList list (f :: Array a %1 -> Ur b) =
+  (Array a %1 -> b) %1 ->
+  b
+fromList list (f :: Array a %1 -> b) =
   alloc
     (Prelude.length list)
     (error "invariant violation: unintialized array position")

src/Data/Array/Mutable/Unlifted/Linear.hs

@@ -60,13 +60,32 @@ infixr 0 `lseq` -- same fixity as base.seq
 -- | Allocate a mutable array of given size using a default value.
 --
 -- The size should be non-negative.
-alloc :: Int -> a -> (Array# a %1 -> Ur b) %1 -> Ur b
-alloc (GHC.I# s) a f =
+alloc :: (Movable b) => Int -> a -> (Array# a %1 -> b) %1 -> b
+alloc i a f = case move (unsafe_alloc i a f) of
+  Ur b -> b
+{-# INLINEABLE alloc #-}
+
+-- The `alloc` function is split in two. One very unsafe below (it's very
+-- unsafe, because `unafe_alloc 57 0 id` returns an unrestricted _mutable_
+-- `Array#` breaking the module's invariants). Because `unsafe_alloc` calls
+-- `runRW#`, it's marked as `NOINLINE`.
+--
+-- It's made safe by the wrapping function `alloc`, which restricts `b` to be
+-- `Movable` (`Array#` is crucially not `Movable`, therefore `alloc 57 0 id`
+-- doesn't type). Furthermore, `alloc` cases on `move` to make sure that all the
+-- effects have been run by the time we evaluate the result of an `alloc`. It's
+-- fine that `alloc` is inlined: its semantics is preserved by program
+-- transformations. It's useful that `alloc` be inlined, because in most
+-- instance `case move … of` will trigger a case-of-known-constructor avoiding
+-- an extra allocation. This is in particular the case for the common case where
+-- `b = Ur x`.
+unsafe_alloc :: Int -> a -> (Array# a %1 -> b) %1 -> b
+unsafe_alloc (GHC.I# s) a f =
   let new = GHC.runRW# Prelude.$ \st ->
         case GHC.newArray# s a st of
           (# _, arr #) -> Array# arr
    in f new
-{-# NOINLINE alloc #-} -- prevents runRW# from floating outwards
+{-# NOINLINE unsafe_alloc #-} -- prevents runRW# from floating outwards
 
 -- For the reasoning behind these NOINLINE pragmas, see the discussion at:
 -- https://github.com/tweag/linear-base/pull/187#pullrequestreview-489183531

src/Data/HashMap/Mutable/Linear/Internal.hs

@@ -133,10 +133,10 @@ data ProbeResult k v where
 -- | Run a computation with an empty 'HashMap' with given capacity.
 empty ::
   forall k v b.
-  (Keyed k) =>
+  (Keyed k, Movable b) =>
   Int ->
-  (HashMap k v %1 -> Ur b) %1 ->
-  Ur b
+  (HashMap k v %1 -> b) %1 ->
+  b
 empty size scope =
   let cap = max 1 size
    in Array.alloc cap Nothing (\arr -> scope (HashMap 0 cap arr))
@@ -151,10 +151,10 @@ allocBeside size (HashMap s' c' arr) =
 -- | Run a computation with an 'HashMap' containing given key-value pairs.
 fromList ::
   forall k v b.
-  (Keyed k) =>
+  (Keyed k, Movable b) =>
   [(k, v)] ->
-  (HashMap k v %1 -> Ur b) %1 ->
-  Ur b
+  (HashMap k v %1 -> b) %1 ->
+  b
 fromList xs scope =
   let cap =
         max

src/Data/Set/Mutable/Linear/Internal.hs

@@ -29,8 +29,8 @@ type Keyed a = Linear.Keyed a
 -- # Constructors and Mutators
 -------------------------------------------------------------------------------
 
-empty :: (Keyed a) => Int -> (Set a %1 -> Ur b) %1 -> Ur b
-empty s (f :: Set a %1 -> Ur b) =
+empty :: (Keyed a, Movable b) => Int -> (Set a %1 -> b) %1 -> b
+empty s (f :: Set a %1 -> b) =
   Linear.empty s (\hm -> f (Set hm))
 
 toList :: (Keyed a) => Set a %1 -> Ur [a]
@@ -63,7 +63,7 @@ member :: (Keyed a) => a -> Set a %1 -> (Ur Bool, Set a)
 member a (Set hm) =
   Linear.member a hm Linear.& \(b, hm') -> (b, Set hm')
 
-fromList :: (Keyed a) => [a] -> (Set a %1 -> Ur b) %1 -> Ur b
+fromList :: (Keyed a, Movable b) => [a] -> (Set a %1 -> b) %1 -> b
 fromList xs f =
   Linear.fromList (Prelude.map (,()) xs) (\hm -> f (Set hm))
 

src/Data/Vector/Mutable/Linear/Internal.hs

@@ -56,25 +56,25 @@ fromArray arr =
     & \(Ur size', arr') -> Vec size' arr'
 
 -- Allocate an empty vector
-empty :: (Vector a %1 -> Ur b) %1 -> Ur b
+empty :: (Movable b) => (Vector a %1 -> b) %1 -> b
 empty f = Array.fromList [] (f . fromArray)
 
 -- | Allocate a constant vector of a given non-negative size (and error on a
 -- bad size)
 constant ::
-  (HasCallStack) =>
+  (HasCallStack, Movable b) =>
   Int ->
   a ->
-  (Vector a %1 -> Ur b) %1 ->
-  Ur b
+  (Vector a %1 -> b) %1 ->
+  b
 constant size' x f
   | size' < 0 =
       (error ("Trying to construct a vector of size " ++ show size') :: x %1 -> x)
         (f undefined)
   | otherwise = Array.alloc size' x (f . fromArray)
 
 -- | Allocator from a list
-fromList :: (HasCallStack) => [a] -> (Vector a %1 -> Ur b) %1 -> Ur b
+fromList :: (HasCallStack, Movable b) => [a] -> (Vector a %1 -> b) %1 -> b
 fromList xs f = Array.fromList xs (f . fromArray)
 
 -- | Number of elements inside the vector.

src/Foreign/Marshal/Pure/Internal.hs

@@ -29,10 +29,10 @@ import Foreign.Marshal.Utils
 import Foreign.Ptr
 import Foreign.Storable
 import Foreign.Storable.Tuple ()
-import Prelude.Linear hiding (Eq (..), ($))
+import Prelude.Linear hiding (Eq (..))
 import System.IO.Unsafe
 import qualified Unsafe.Linear as Unsafe
-import Prelude (Eq (..), return, ($), (<$>), (<*>), (=<<))
+import Prelude (Eq (..), return, (<$>), (<*>), (=<<))
 
 -- XXX: [2018-02-09] I'm having trouble with the `constraints` package (it seems
 -- that the version of Type.Reflection.Unsafe in the linear ghc compiler is not
@@ -290,20 +290,20 @@ freeAll start end = do
 -- TODO: document individual functions
 
 -- | Given a linear computation that manages memory, run that computation.
-withPool :: (Pool %1 -> Ur b) %1 -> Ur b
-withPool scope = Unsafe.toLinear performScope scope
+withPool :: forall b. (Movable b) => (Pool %1 -> b) %1 -> b
+withPool scope = unur $ Unsafe.toLinear performScope scope
   where
     -- XXX: do ^ without `toLinear` by using linear IO
 
-    performScope :: (Pool %1 -> Ur b) -> Ur b
+    performScope :: (Pool %1 -> b) -> Ur b
     performScope scope' = unsafeDupablePerformIO $ do
       -- Initialise the pool
       backPtr <- malloc
       let end = DLL backPtr nullPtr nullPtr -- always at the end of the list
       start <- DLL nullPtr nullPtr <$> new end -- always at the start of the list
       poke backPtr start
       -- Run the computation
-      evaluate (scope' (Pool start))
+      evaluate (move $ scope' (Pool start))
         `finally`
         -- Clean up remaining variables.
         (freeAll start end)