{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- | This module provides push arrays.
--
-- These are part of a larger framework for controlling when memory is
-- allocated for an array. See @Data.Array.Polarized@.
--
-- This module is designed to be imported qualified as @Push@.
module Data.Array.Polarized.Push
  ( -- * Construction
    Array (..),
    make,
    singleton,
    cons,
    snoc,

    -- * Operations
    alloc,
    foldMap,
    unzip,
  )
where

import Data.Array.Destination (DArray)
import qualified Data.Array.Destination as DArray
import qualified Data.Functor.Linear as Data
import Data.Vector (Vector)
import GHC.Stack
import Prelude.Linear hiding (foldMap, unzip)
import qualified Prelude

-- 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. The @m@ is something that
-- represents a writing of some elements to an array, a delayed write.
--
-- 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@
--
-- Developer notes:
--
-- This is the linear monoid @m@ that we instantiate the above array with
-- in order to allocate. An @ArrayWriter a@ is something that holds the
-- ingredients to write some number of elements to an array, without
-- holding the space to do so.

-- API
-------------------------------------------------------------------------------

-- | 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 :: forall a. Array a %1 -> Vector a
alloc (Array forall m. Monoid m => (a -> m) -> m
k) = forall a. ArrayWriter a %1 -> Vector a
allocArrayWriter forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall m. Monoid m => (a -> m) -> m
k forall a. a -> ArrayWriter a
singletonWriter
  where
    singletonWriter :: a -> ArrayWriter a
    singletonWriter :: forall a. a -> ArrayWriter a
singletonWriter a
a = forall a. (DArray a %1 -> ()) -> Int -> ArrayWriter a
ArrayWriter (forall a. HasCallStack => a %1 -> DArray a %1 -> ()
DArray.fill a
a) Int
1

    allocArrayWriter :: ArrayWriter a %1 -> Vector a
    allocArrayWriter :: forall a. ArrayWriter a %1 -> Vector a
allocArrayWriter (ArrayWriter DArray a %1 -> ()
writer Int
len) = forall a. Int -> (DArray a %1 -> ()) %1 -> Vector a
DArray.alloc Int
len DArray a %1 -> ()
writer

-- | @`make` x n@ creates a constant push array of length @n@ in which every
-- element is @x@.
make :: (HasCallStack) => a -> Int -> Array a
make :: forall a. HasCallStack => a -> Int -> Array a
make a
x Int
n
  | Int
n forall a. Ord a => a %1 -> a %1 -> Bool
< Int
0 = forall a. HasCallStack => [Char] -> a
error [Char]
"Making a negative length push array"
  | Bool
otherwise = forall a. (forall m. Monoid m => (a -> m) -> m) -> Array a
Array (\a -> m
makeA -> forall a. Monoid a => [a] %1 -> a
mconcat forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall a. Int -> a -> [a]
Prelude.replicate Int
n (a -> m
makeA a
x))

singleton :: a -> Array a
singleton :: forall a. a -> Array a
singleton a
x = forall a. (forall m. Monoid m => (a -> m) -> m) -> Array a
Array (\a -> m
writeA -> a -> m
writeA a
x)

snoc :: a -> Array a %1 -> Array a
snoc :: forall a. a -> Array a %1 -> Array a
snoc a
x (Array forall m. Monoid m => (a -> m) -> m
k) = forall a. (forall m. Monoid m => (a -> m) -> m) -> Array a
Array (\a -> m
writeA -> (forall m. Monoid m => (a -> m) -> m
k a -> m
writeA) forall a. Semigroup a => a %1 -> a %1 -> a
<> (a -> m
writeA a
x))

cons :: a -> Array a %1 -> Array a
cons :: forall a. a -> Array a %1 -> Array a
cons a
x (Array forall m. Monoid m => (a -> m) -> m
k) = forall a. (forall m. Monoid m => (a -> m) -> m) -> Array a
Array (\a -> m
writeA -> (a -> m
writeA a
x) forall a. Semigroup a => a %1 -> a %1 -> a
<> (forall m. Monoid m => (a -> m) -> m
k a -> m
writeA))

foldMap :: (Monoid b) => (a -> b) -> Array a %1 -> b
foldMap :: forall b a. Monoid b => (a -> b) -> Array a %1 -> b
foldMap a -> b
f (Array forall m. Monoid m => (a -> m) -> m
k) = forall m. Monoid m => (a -> m) -> m
k a -> b
f

unzip :: Array (a, b) %1 -> (Array a, Array b)
unzip :: forall a b. Array (a, b) %1 -> (Array a, Array b)
unzip (Array forall m. Monoid m => ((a, b) -> m) -> m
k) = forall m. Monoid m => ((a, b) -> m) -> m
k (\(a
a, b
b) -> (forall a. a -> Array a
singleton a
a, forall a. a -> Array a
singleton b
b))

-- # Instances
-------------------------------------------------------------------------------

instance Data.Functor Array where
  fmap :: forall a b. (a %1 -> b) -> Array a %1 -> Array b
fmap a %1 -> b
f (Array forall m. Monoid m => (a -> m) -> m
k) = forall a. (forall m. Monoid m => (a -> m) -> m) -> Array a
Array (\b -> m
g -> forall m. Monoid m => (a -> m) -> m
k (\a
x -> (b -> m
g (a %1 -> b
f a
x))))

instance Prelude.Semigroup (Array a) where
  <> :: Array a -> Array a -> Array a
(<>) Array a
x Array a
y = forall a. Array a %1 -> Array a %1 -> Array a
append Array a
x Array a
y

instance Semigroup (Array a) where
  <> :: Array a %1 -> Array a %1 -> Array a
(<>) = forall a. Array a %1 -> Array a %1 -> Array a
append

instance Prelude.Monoid (Array a) where
  mempty :: Array a
mempty = forall a. Array a
empty

instance Monoid (Array a) where
  mempty :: Array a
mempty = forall a. Array a
empty

empty :: Array a
empty :: forall a. Array a
empty = forall a. (forall m. Monoid m => (a -> m) -> m) -> Array a
Array (\a -> m
_ -> forall a. Monoid a => a
mempty)

append :: Array a %1 -> Array a %1 -> Array a
append :: forall a. Array a %1 -> Array a %1 -> Array a
append (Array forall m. Monoid m => (a -> m) -> m
k1) (Array forall m. Monoid m => (a -> m) -> m
k2) = forall a. (forall m. Monoid m => (a -> m) -> m) -> Array a
Array (\a -> m
writeA -> forall m. Monoid m => (a -> m) -> m
k1 a -> m
writeA forall a. Semigroup a => a %1 -> a %1 -> a
<> forall m. Monoid m => (a -> m) -> m
k2 a -> m
writeA)

instance Prelude.Semigroup (ArrayWriter a) where
  <> :: ArrayWriter a -> ArrayWriter a -> ArrayWriter a
(<>) ArrayWriter a
x ArrayWriter a
y = forall a. ArrayWriter a %1 -> ArrayWriter a %1 -> ArrayWriter a
addWriters ArrayWriter a
x ArrayWriter a
y

instance Prelude.Monoid (ArrayWriter a) where
  mempty :: ArrayWriter a
mempty = forall a. ArrayWriter a
emptyWriter

instance Semigroup (ArrayWriter a) where
  <> :: ArrayWriter a %1 -> ArrayWriter a %1 -> ArrayWriter a
(<>) = forall a. ArrayWriter a %1 -> ArrayWriter a %1 -> ArrayWriter a
addWriters

instance Monoid (ArrayWriter a) where
  mempty :: ArrayWriter a
mempty = forall a. ArrayWriter a
emptyWriter

addWriters :: ArrayWriter a %1 -> ArrayWriter a %1 -> ArrayWriter a
addWriters :: forall a. ArrayWriter a %1 -> ArrayWriter a %1 -> ArrayWriter a
addWriters (ArrayWriter DArray a %1 -> ()
k1 Int
l1) (ArrayWriter DArray a %1 -> ()
k2 Int
l2) =
  forall a. (DArray a %1 -> ()) -> Int -> ArrayWriter a
ArrayWriter
    ( \DArray a
darr ->
        (forall a. Int -> DArray a %1 -> (DArray a, DArray a)
DArray.split Int
l1 DArray a
darr) forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(DArray a
darr1, DArray a
darr2) -> forall a. Consumable a => a %1 -> ()
consume (DArray a %1 -> ()
k1 DArray a
darr1, DArray a %1 -> ()
k2 DArray a
darr2)
    )
    (Int
l1 forall a. Additive a => a %1 -> a %1 -> a
+ Int
l2)

emptyWriter :: ArrayWriter a
emptyWriter :: forall a. ArrayWriter a
emptyWriter = forall a. (DArray a %1 -> ()) -> Int -> ArrayWriter a
ArrayWriter forall a. HasCallStack => DArray a %1 -> ()
DArray.dropEmpty Int
0

-- Remark. @emptyWriter@ assumes we can split a destination array at 0.