{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE UnboxedTuples #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}
{-# OPTIONS_HADDOCK hide #-}

module Data.Array.Mutable.Linear.Internal
  ( -- * Mutable Linear Arrays
    Array (..),

    -- * Performing Computations with Arrays
    alloc,
    allocBeside,
    fromList,

    -- * Modifications
    set,
    unsafeSet,
    resize,
    map,

    -- * Accessors
    get,
    unsafeGet,
    size,
    slice,
    toList,
    freeze,

    -- * Mutable-style interface
    read,
    unsafeRead,
    write,
    unsafeWrite,
  )
where

import Data.Array.Mutable.Unlifted.Linear (Array#)
import qualified Data.Array.Mutable.Unlifted.Linear as Unlifted
import qualified Data.Functor.Linear as Data
import qualified Data.Primitive.Array as Prim
import Data.Unrestricted.Linear
import qualified Data.Vector as Vector
import GHC.Stack
import Prelude.Linear (forget, (&))
import Prelude hiding (map, read)

-- # Data types
-------------------------------------------------------------------------------

data Array a = Array (Array# a)

-- # Creation
-------------------------------------------------------------------------------

-- | Allocate a constant array given a size and an initial value
-- The size must be non-negative, otherwise this errors.
alloc ::
  HasCallStack =>
  Int ->
  a ->
  (Array a %1 -> Ur b) %1 ->
  Ur b
alloc s x f
  | s < 0 =
      (error ("Array.alloc: negative size: " ++ show s) :: x %1 -> x)
        (f undefined)
  | otherwise = Unlifted.alloc s x (\arr -> f (Array arr))

-- | Allocate a constant array given a size and an initial value,
-- using another array as a uniqueness proof.
allocBeside :: Int -> a -> Array b %1 -> (Array a, Array b)
allocBeside s x (Array orig)
  | s < 0 =
      Unlifted.lseq
        orig
        (error ("Array.allocBeside: negative size: " ++ show s))
  | otherwise =
      wrap (Unlifted.allocBeside s x orig)
  where
    wrap :: (# Array# a, Array# b #) %1 -> (Array a, Array b)
    wrap (# orig, new #) = (Array orig, Array new)

-- | Allocate an array from a list
fromList ::
  HasCallStack =>
  [a] ->
  (Array a %1 -> Ur b) %1 ->
  Ur b
fromList list (f :: Array a %1 -> Ur b) =
  alloc
    (Prelude.length list)
    (error "invariant violation: unintialized array position")
    (\arr -> f (insert arr))
  where
    insert :: Array a %1 -> Array a
    insert = doWrites (zip list [0 ..])

    doWrites :: [(a, Int)] -> Array a %1 -> Array a
    doWrites [] arr = arr
    doWrites ((a, ix) : xs) arr = doWrites xs (unsafeSet ix a arr)

-- # Mutations and Reads
-------------------------------------------------------------------------------

size :: Array a %1 -> (Ur Int, Array a)
size (Array arr) = f (Unlifted.size arr)
  where
    f :: (# Ur Int, Array# a #) %1 -> (Ur Int, Array a)
    f (# s, arr #) = (s, Array arr)

-- | Sets the value of an index. The index should be less than the arrays
-- size, otherwise this errors.
set :: HasCallStack => Int -> a -> Array a %1 -> Array a
set i x arr = unsafeSet i x (assertIndexInRange i arr)

-- | Same as 'set, but does not do bounds-checking. The behaviour is undefined
-- if an out-of-bounds index is provided.
unsafeSet :: Int -> a -> Array a %1 -> Array a
unsafeSet ix val (Array arr) =
  Array (Unlifted.set ix val arr)

-- | Get the value of an index. The index should be less than the arrays 'size',
-- otherwise this errors.
get :: HasCallStack => Int -> Array a %1 -> (Ur a, Array a)
get i arr = unsafeGet i (assertIndexInRange i arr)

-- | Same as 'get', but does not do bounds-checking. The behaviour is undefined
-- if an out-of-bounds index is provided.
unsafeGet :: Int -> Array a %1 -> (Ur a, Array a)
unsafeGet ix (Array arr) = wrap (Unlifted.get ix arr)
  where
    wrap :: (# Ur a, Array# a #) %1 -> (Ur a, Array a)
    wrap (# ret, arr #) = (ret, Array arr)

-- | Resize an array. That is, given an array, a target size, and a seed
-- value; resize the array to the given size using the seed value to fill
-- in the new cells when necessary and copying over all the unchanged cells.
--
-- Target size should be non-negative.
--
-- @
-- let b = resize n x a,
--   then size b = n,
--   and b[i] = a[i] for i < size a,
--   and b[i] = x for size a <= i < n.
-- @
resize :: HasCallStack => Int -> a -> Array a %1 -> Array a
resize newSize seed (Array arr :: Array a)
  | newSize < 0 =
      Unlifted.lseq
        arr
        (error "Trying to resize to a negative size.")
  | otherwise =
      doCopy (Unlifted.allocBeside newSize seed arr)
  where
    doCopy :: (# Array# a, Array# a #) %1 -> Array a
    doCopy (# new, old #) = wrap (Unlifted.copyInto 0 old new)

    wrap :: (# Array# a, Array# a #) %1 -> Array a
    wrap (# src, dst #) = src `Unlifted.lseq` Array dst

-- | Return the array elements as a lazy list.
toList :: Array a %1 -> Ur [a]
toList (Array arr) = Unlifted.toList arr

-- | Copy a slice of the array, starting from given offset and copying given
-- number of elements. Returns the pair (oldArray, slice).
--
-- Start offset + target size should be within the input array, and both should
-- be non-negative.
--
-- @
-- let b = slice i n a,
--   then size b = n,
--   and b[j] = a[i+j] for 0 <= j < n
-- @
slice ::
  HasCallStack =>
  -- | Start offset
  Int ->
  -- | Target size
  Int ->
  Array a %1 ->
  (Array a, Array a)
slice from targetSize arr =
  size arr & \case
    (Ur s, Array old)
      | s < from + targetSize ->
          Unlifted.lseq
            old
            (error "Slice index out of bounds.")
      | otherwise ->
          doCopy
            ( Unlifted.allocBeside
                targetSize
                (error "invariant violation: uninitialized array index")
                old
            )
  where
    doCopy :: (# Array# a, Array# a #) %1 -> (Array a, Array a)
    doCopy (# new, old #) = wrap (Unlifted.copyInto from old new)

    wrap :: (# Array# a, Array# a #) %1 -> (Array a, Array a)
    wrap (# old, new #) = (Array old, Array new)

-- | /O(1)/ Convert an 'Array' to an immutable 'Vector.Vector' (from
-- 'vector' package).
freeze :: Array a %1 -> Ur (Vector.Vector a)
freeze (Array arr) =
  Unlifted.freeze (\a -> Vector.fromArray (Prim.Array a)) arr

map :: (a -> b) -> Array a %1 -> Array b
map f (Array arr) = Array (Unlifted.map f arr)

-- # Mutation-style API
-------------------------------------------------------------------------------

-- | Same as 'set', but takes the 'Array' as the first parameter.
write :: HasCallStack => Array a %1 -> Int -> a -> Array a
write arr i a = set i a arr

-- | Same as 'unsafeSet', but takes the 'Array' as the first parameter.
unsafeWrite :: Array a %1 -> Int -> a -> Array a
unsafeWrite arr i a = unsafeSet i a arr

-- | Same as 'get', but takes the 'Array' as the first parameter.
read :: HasCallStack => Array a %1 -> Int -> (Ur a, Array a)
read arr i = get i arr

-- | Same as 'unsafeGet', but takes the 'Array' as the first parameter.
unsafeRead :: Array a %1 -> Int -> (Ur a, Array a)
unsafeRead arr i = unsafeGet i arr

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

instance Consumable (Array a) where
  consume :: Array a %1 -> ()
  consume (Array arr) = arr `Unlifted.lseq` ()

instance Dupable (Array a) where
  dup2 :: Array a %1 -> (Array a, Array a)
  dup2 (Array arr) = wrap (Unlifted.dup2 arr)
    where
      wrap :: (# Array# a, Array# a #) %1 -> (Array a, Array a)
      wrap (# a1, a2 #) = (Array a1, Array a2)

instance Data.Functor Array where
  fmap f arr = map (forget f) arr

-- # Internal library
-------------------------------------------------------------------------------

-- | Check if given index is within the Array, otherwise panic.
assertIndexInRange :: HasCallStack => Int -> Array a %1 -> Array a
assertIndexInRange i arr =
  size arr & \(Ur s, arr') ->
    if 0 <= i && i < s
      then arr'
      else arr' `lseq` error "Array: index out of bounds"