content: {-# LANGUAGE GADTs #-}
  {-# LANGUAGE LinearTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE NoImplicitPrelude #-}
  {-# OPTIONS_HADDOCK hide #-}
  
  module Data.Array.Destination.Internal where
  
  import Data.Unrestricted.Linear
  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 GHC.Stack
  import Prelude.Linear hiding (replicate)
  import System.IO.Unsafe (unsafeDupablePerformIO)
  import qualified Unsafe.Linear as Unsafe
  import qualified Prelude as Prelude
  
  -- | A destination array, or @DArray@, is a write-only array that is filled
  -- by some computation which ultimately returns an array.
  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 writer = (\(Ur dest, vec) -> writer (DArray dest) `lseq` vec) $
    unsafeDupablePerformIO $ 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 mvec) = (Ur (MVector.length mvec), DArray mvec)
  
  -- | Fill a destination array with a constant
  replicate :: a -> DArray a %1 -> ()
  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 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 (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 (DArray mvec)
    | (ml, mr) <- MVector.splitAt n mvec =
        (DArray ml, DArray mr)
  
  -- | Fills the destination array with the contents of given vector.
  --
  -- Errors if the given vector is smaller than the destination array.
  mirror :: (HasCallStack) => Vector a -> (a %1 -> b) -> DArray b %1 -> ()
  mirror v f arr =
    size arr & \(Ur sz, arr') ->
      if Vector.length v < sz
        then error "Destination.mirror: argument smaller than DArray" $ arr'
        else fromFunction (\t -> f (v ! t)) arr'
  
  -- | Fill a destination array using the given index-to-value function.
  fromFunction :: (Int -> b) -> DArray b %1 -> ()
  fromFunction f (DArray mvec) = unsafeDupablePerformIO $ 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.
  