Internal.hs

{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# OPTIONS_GHC -fno-warn-partial-type-signatures #-}
{-# OPTIONS_HADDOCK hide #-}

module Data.Array.Polarized.Pull.Internal where

import qualified Data.Functor.Linear as Data
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import Prelude.Linear
import qualified Unsafe.Linear as Unsafe
import qualified Prelude

-- | A pull array is an array from which it is easy to extract elements, and
-- this can be done in any order. The linear consumption of a pull array means
-- each element is consumed exactly once, but the length can be accessed
-- freely.
data Array a where
  Array :: (Int -> a) -> Int -> Array a
  deriving (Prelude.Semigroup) via NonLinear (Array a)

-- In the linear consumption of a pull array f n, (f i) should be consumed
-- linearly for every 0 <= i < n. The exported functions (from non-internal
-- modules) should enforce this invariant, but the current type of PullArray
-- does not.

instance Data.Functor Array where
  fmap f (Array g n) = fromFunction (\x -> f (g x)) n

-- XXX: This should be well-typed without the unsafe, but it isn't accepted:
-- the pull array type probably isn't the ideal choice (making Array linear in
-- (Int -> a) would mean only one value could be taken out of the Array (which
-- is interesting in and of itself: I think this is like an n-ary With), and
-- changing the other arrows makes no difference)

-- | Produce a pull array of lenght 1 consisting of solely the given element.
singleton :: a %1 -> Array a
singleton = Unsafe.toLinear (\x -> fromFunction (\_ -> x) 1)

-- | @zip [x1, ..., xn] [y1, ..., yn] = [(x1,y1), ..., (xn,yn)]@
-- __Partial:__ `zip [x1,x2,...,xn] [y1,y2,...,yp]` is an error if @n ≠ p@.
zip :: Array a %1 -> Array b %1 -> Array (a, b)
zip (Array g n) (Array h m)
  | n /= m = error "Polarized.zip: size mismatch"
  | otherwise = fromFunction (\k -> (g k, h k)) n

-- | Concatenate two pull arrays.
append :: Array a %1 -> Array a %1 -> Array a
append (Array f m) (Array g n) = Array h (m + n)
  where
    h k =
      if k < m
        then f k
        else g (k - m)

-- | Creates a pull array of given size, filled with the given element.
make :: a -> Int -> Array a
make x n = fromFunction (const x) n

instance Semigroup (Array a) where
  (<>) = append

-- | A right-fold of a pull array.
foldr :: (a %1 -> b %1 -> b) -> b %1 -> Array a %1 -> b
foldr f z (Array g n) = go f z g n
  where
    go :: (_) => (_ %1 -> _ %1 -> _) -> _ %1 -> _ -> _ -> _
    go _ z' _ 0 = z'
    go f' z' g' k = go f' (f' (g' (k - 1)) z') g' (k - 1)

-- go is strict in its last argument

-- | Extract the length of an array, and give back the original array.
findLength :: Array a %1 -> (Int, Array a)
findLength (Array f n) = (n, Array f n)

-- | @'fromFunction' arrIndexer len@ constructs a pull array given a function
-- @arrIndexer@ that goes from an array index to array values and a specified
-- length @len@.
fromFunction :: (Int -> a) -> Int -> Array a
fromFunction f n = Array f' n
  where
    f' k
      | k < 0 = error "Pull.Array: negative index"
      | k >= n = error "Pull.Array: index too large"
      | otherwise = f k

-- XXX: this is used internally to ensure out of bounds errors occur, but
-- is unnecessary if the input function can be assumed to already have bounded
-- domain, for instance in `append`.

-- | This is a convenience function for @alloc . transfer@
toVector :: Array a %1 -> Vector a
toVector (Array f n) = Vector.generate n f

-- TODO: A test to make sure alloc . transfer == toVector

-- | @'split' n v = (vl, vr)@ such that @vl@ has length @n@.
--
-- 'split' is total: if @n@ is larger than the length of @v@,
-- then @vr@ is empty.
split :: Int -> Array a %1 -> (Array a, Array a)
split k (Array f n) = (fromFunction f (min k n), fromFunction (\x -> f (x + k)) (max (n - k) 0))

-- | Reverse a pull array.
reverse :: Array a %1 -> Array a
reverse (Array f n) = Array (\x -> f (n + 1 - x)) n

-- | Index a pull array (without checking bounds)
index :: Array a %1 -> Int -> (a, Array a)
index (Array f n) ix = (f ix, Array f n)