content: {-# 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 (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 :: (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))
  
  singleton :: a -> Array a
  singleton x = Array (\writeA -> writeA x)
  
  snoc :: a -> Array a %1 -> Array a
  snoc x (Array k) = Array (\writeA -> (k writeA) <> (writeA x))
  
  cons :: a -> Array a %1 -> Array a
  cons x (Array k) = Array (\writeA -> (writeA x) <> (k writeA))
  
  foldMap :: (Monoid b) => (a -> b) -> Array a %1 -> b
  foldMap f (Array k) = k f
  
  unzip :: Array (a, b) %1 -> (Array a, Array b)
  unzip (Array k) = k (\(a, b) -> (singleton a, singleton b))
  
  -- # 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)
  
  append :: Array a %1 -> Array a %1 -> Array a
  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.
  