Merge pull request #287 from tweag/better-push-arrays-2

More generic push arrays 1/2

Author: Divesh Otwani

src/Data/Array/Destination.hs

@@ -123,6 +123,7 @@ module Data.Array.Destination
   , mirror
   , fromFunction
   , fill
+  , dropEmpty
   )
   where
 
@@ -179,6 +180,14 @@ fill = Unsafe.toLinear2 unsafeFill
       else
         unsafeDupablePerformIO Prelude.$ MVector.write ds 0 a
 
+-- | @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` ()
+
 -- | @'split' n dest = (destl, destr)@ such as @destl@ has length @n@.
 --
 -- 'split' is total: if @n@ is larger than the length of @dest@, then

src/Data/Array/Polarized.hs

@@ -99,10 +99,10 @@ module Data.Array.Polarized
   )
   where
 
-import qualified Data.Array.Destination as DArray
 import qualified Data.Array.Polarized.Pull.Internal as Pull
 import qualified Data.Array.Polarized.Pull as Pull
 import qualified Data.Array.Polarized.Push as Push
+import qualified Data.Foldable as NonLinear
 import Prelude.Linear
 import Data.Vector (Vector)
 
@@ -129,7 +129,8 @@ import Data.Vector (Vector)
 -- | Convert a pull array into a push array.
 -- NOTE: this does NOT require allocation and can be in-lined.
 transfer :: Pull.Array a %1-> Push.Array a
-transfer (Pull.Array f n) = Push.Array (\g -> DArray.fromFunction (\i -> g (f i))) n
+transfer (Pull.Array f n) =
+  Push.Array (\k -> NonLinear.foldMap' (\x -> k (f x)) [0..(n-1)])
 
 -- | This is a shortcut convenience function
 -- for @transfer . Pull.fromVector@.

src/Data/Array/Polarized/Push.hs

@@ -1,4 +1,3 @@
-{-# OPTIONS_GHC -fno-warn-partial-type-signatures #-}
 {-# LANGUAGE DerivingVia #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE LinearTypes #-}
@@ -11,55 +10,107 @@
 -- allocated for an array. See @Data.Array.Polarized@.
 --
 -- This module is designed to be imported qualified as @Push@.
-module Data.Array.Polarized.Push where
+module Data.Array.Polarized.Push
+  ( Array(..)
+  , alloc
+  , make
+  )
+where
 
--- XXX: it might be better to hide the data constructor, in case we wish to
--- change the implementation.
-
-import Data.Array.Destination (DArray)
-import qualified Data.Array.Destination as DArray
 import qualified Data.Functor.Linear as Data
+import qualified Data.Array.Destination as DArray
+import Data.Array.Destination (DArray)
 import Data.Vector (Vector)
-import Prelude.Linear
 import qualified Prelude
+import Prelude.Linear
+import GHC.Stack
+
 
--- TODO: the below isn't really true yet since no friendly way of constructing
--- a PushArray directly is given yet (see issue #62), but the spirit holds.
--- TODO: There's also a slight lie in that one might want to consume a
--- PushArray into a commutative monoid, for instance summing all the elements,
--- and this is not yet possible, although it should be.
+-- The Types
+-------------------------------------------------------------------------------
 
 -- | Push arrays are un-allocated finished arrays. These are finished
 -- computations passed along or enlarged until we are ready to allocate.
 data Array a where
+  Array :: (forall m. Monoid m => (a -> m) -> m) %1-> Array a
+  -- Developer notes:
+  --
+  -- Think of @(a -> m)@ as something that writes an @a@ and think of
+  -- @((a -> m) -> m)@ as something that takes a way to write a single element
+  -- and writes and concatenates all elements.
+  --
+  -- Also, note that in this formulation we don't know the length beforehand.
+
+data ArrayWriter a where
+  ArrayWriter :: (DArray a %1-> ()) %1-> !Int -> ArrayWriter a
   -- The second parameter is the length of the @DArray@
-  Array :: (forall b. (a %1-> b) -> DArray b %1-> ()) %1-> Int -> Array a
-  deriving Prelude.Semigroup via NonLinear (Array a)
 
-instance Data.Functor Array where
-  fmap f (Array k n) = Array (\g dest -> k (g . f) dest) n
 
-instance Semigroup (Array a) where
-  (<>) = append
+-- API
+-------------------------------------------------------------------------------
 
--- XXX: the use of Vector in the type of alloc is temporary (see also
--- "Data.Array.Destination")
 -- | Convert a push array into a vector by allocating. This would be a common
 -- end to a computation using push and pull arrays.
 alloc :: Array a %1-> Vector a
-alloc (Array k n) = DArray.alloc n (k id)
+alloc (Array k) = allocArrayWriter $ k singletonWriter where
+  singletonWriter :: a -> ArrayWriter a
+  singletonWriter a = ArrayWriter (DArray.fill a) 1
+
+  allocArrayWriter :: ArrayWriter a %1-> Vector a
+  allocArrayWriter (ArrayWriter writer len) = DArray.alloc len writer
 
 -- | @`make` x n@ creates a constant push array of length @n@ in which every
 -- element is @x@.
-make :: a -> Int -> Array a
-make x n = Array (\k -> DArray.replicate (k x)) n
+make :: HasCallStack => a -> Int -> Array a
+make x n
+  | n < 0 = error "Making a negative length push array"
+  | otherwise = Array (\makeA -> mconcat $ Prelude.replicate n (makeA x))
+
+
+-- # Instances
+-------------------------------------------------------------------------------
+
+instance Data.Functor Array where
+  fmap f (Array k) = Array (\g -> k (\x -> (g (f x))))
+
+instance Prelude.Semigroup (Array a) where
+  (<>) x y = append x y
+
+instance Semigroup (Array a) where
+  (<>) = append
+
+instance Prelude.Monoid (Array a) where
+  mempty = empty
+
+instance Monoid (Array a) where
+  mempty = empty
+
+empty :: Array a
+empty = Array (\_ -> mempty)
 
--- | Concatenate two push arrays.
 append :: Array a %1-> Array a %1-> Array a
-append (Array kl nl) (Array kr nr) =
-    Array
-      (\f dest -> parallelApply f kl kr (DArray.split nl dest))
-      (nl+nr)
-  where
-    parallelApply :: (a %1-> b) -> ((a %1-> b) -> DArray b %1-> ()) %1-> ((a %1-> b) -> DArray b %1-> ()) %1-> (DArray b, DArray b) %1-> ()
-    parallelApply f' kl' kr' (dl, dr) = kl' f' dl <> kr' f' dr
+append (Array k1) (Array k2) = Array (\writeA -> k1 writeA <> k2 writeA)
+
+instance Prelude.Semigroup (ArrayWriter a) where
+  (<>) x y = addWriters x y
+
+instance Prelude.Monoid (ArrayWriter a) where
+  mempty = emptyWriter
+
+instance Semigroup (ArrayWriter a) where
+  (<>) = addWriters
+
+instance Monoid (ArrayWriter a) where
+  mempty = emptyWriter
+
+addWriters :: ArrayWriter a %1-> ArrayWriter a %1-> ArrayWriter a
+addWriters (ArrayWriter k1 l1) (ArrayWriter k2 l2) =
+  ArrayWriter
+    (\darr ->
+      (DArray.split l1 darr) & \(darr1,darr2) -> consume (k1 darr1, k2 darr2))
+    (l1+l2)
+
+emptyWriter :: ArrayWriter a
+emptyWriter = ArrayWriter DArray.dropEmpty 0
+-- Remark. @emptyWriter@ assumes we can split a destination array at 0.
+