{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UnliftedNewtypes #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- |
-- This module provides an unlifted mutable array with a pure
-- interface. Though the array itself is unlifted, it's elements are
-- lifted types. This is made possible by using linear types to make
-- sure array references are single threaded through reads and writes.
--
-- Accessing out-of-bounds indices causes undefined behaviour.
--
-- This module is meant to be imported qualified.
module Data.Array.Mutable.Unlifted.Linear
  ( Array#,
    unArray#,
    alloc,
    allocBeside,
    lseq,
    size,
    get,
    set,
    copyInto,
    map,
    toList,
    freeze,
    dup2,
  )
where

import Data.Unrestricted.Linear hiding (dup2, lseq)
import qualified GHC.Exts as GHC
import qualified Unsafe.Linear as Unsafe
import Prelude (Int)
import qualified Prelude as Prelude

-- | A mutable array holding @a@s
newtype Array# a = Array# (GHC.MutableArray# GHC.RealWorld a)

-- | Extract the underlying 'GHC.MutableArray#', consuming the 'Array#'
-- in process.
unArray# :: (GHC.MutableArray# GHC.RealWorld a -> b) -> Array# a %1 -> Ur b
unArray# :: forall a b. (MutableArray# RealWorld a -> b) -> Array# a %1 -> Ur b
unArray# MutableArray# RealWorld a -> b
f = forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear (\(Array# MutableArray# RealWorld a
a) -> forall a. a -> Ur a
Ur (MutableArray# RealWorld a -> b
f MutableArray# RealWorld a
a))

-- | Consume an 'Array#'.
--
-- Note that we can not implement a 'Consumable' instance because 'Array#'
-- is unlifted.
lseq :: Array# a %1 -> b %1 -> b
lseq :: forall a b. Array# a %1 -> b %1 -> b
lseq = forall a b c (p :: Multiplicity) (q :: Multiplicity)
       (x :: Multiplicity) (y :: Multiplicity).
(a %p -> b %q -> c) %1 -> a %x -> b %y -> c
Unsafe.toLinear2 (\Array# a
_ b
b -> b
b)

infixr 0 `lseq` -- same fixity as base.seq

-- | Allocate a mutable array of given size using a default value.
--
-- The size should be non-negative.
alloc :: Int -> a -> (Array# a %1 -> Ur b) %1 -> Ur b
alloc :: forall a b. Int -> a -> (Array# a %1 -> Ur b) %1 -> Ur b
alloc (GHC.I# Int#
s) a
a Array# a %1 -> Ur b
f =
  let new :: Array# a
new = forall o. (State# RealWorld -> o) -> o
GHC.runRW# forall a b. (a -> b) -> a -> b
Prelude.$ \State# RealWorld
st ->
        case forall a d.
Int# -> a -> State# d -> (# State# d, MutableArray# d a #)
GHC.newArray# Int#
s a
a State# RealWorld
st of
          (# State# RealWorld
_, MutableArray# RealWorld a
arr #) -> forall a. MutableArray# RealWorld a -> Array# a
Array# MutableArray# RealWorld a
arr
   in Array# a %1 -> Ur b
f Array# a
new
{-# NOINLINE alloc #-} -- prevents runRW# from floating outwards

-- For the reasoning behind these NOINLINE pragmas, see the discussion at:
-- https://github.com/tweag/linear-base/pull/187#pullrequestreview-489183531

-- | Allocate a mutable array of given size using a default value,
-- using another 'Array#' as a uniqueness proof.
--
-- The size should be non-negative.
allocBeside :: Int -> a -> Array# b %1 -> (# Array# a, Array# b #)
allocBeside :: forall a b. Int -> a -> Array# b %1 -> (# Array# a, Array# b #)
allocBeside (GHC.I# Int#
s) a
a Array# b
orig =
  let new :: Array# a
new = forall o. (State# RealWorld -> o) -> o
GHC.runRW# forall a b. (a -> b) -> a -> b
Prelude.$ \State# RealWorld
st ->
        case forall a d.
Int# -> a -> State# d -> (# State# d, MutableArray# d a #)
GHC.newArray# Int#
s a
a State# RealWorld
st of
          (# State# RealWorld
_, MutableArray# RealWorld a
arr #) -> forall a. MutableArray# RealWorld a -> Array# a
Array# MutableArray# RealWorld a
arr
   in (# Array# a
new, Array# b
orig #)
{-# NOINLINE allocBeside #-} -- prevents runRW# from floating outwards

size :: Array# a %1 -> (# Ur Int, Array# a #)
size :: forall a. Array# a %1 -> (# Ur Int, Array# a #)
size = forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear forall a. Array# a -> (# Ur Int, Array# a #)
go
  where
    go :: Array# a -> (# Ur Int, Array# a #)
    go :: forall a. Array# a -> (# Ur Int, Array# a #)
go (Array# MutableArray# RealWorld a
arr) =
      let !s :: Int#
s = forall d a. MutableArray# d a -> Int#
GHC.sizeofMutableArray# MutableArray# RealWorld a
arr
       in (# forall a. a -> Ur a
Ur (Int# -> Int
GHC.I# Int#
s), forall a. MutableArray# RealWorld a -> Array# a
Array# MutableArray# RealWorld a
arr #)

get :: Int -> Array# a %1 -> (# Ur a, Array# a #)
get :: forall a. Int -> Array# a %1 -> (# Ur a, Array# a #)
get (GHC.I# Int#
i) = forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear forall a. Array# a -> (# Ur a, Array# a #)
go
  where
    go :: Array# a -> (# Ur a, Array# a #)
    go :: forall a. Array# a -> (# Ur a, Array# a #)
go (Array# MutableArray# RealWorld a
arr) =
      case forall o. (State# RealWorld -> o) -> o
GHC.runRW# (forall d a.
MutableArray# d a -> Int# -> State# d -> (# State# d, a #)
GHC.readArray# MutableArray# RealWorld a
arr Int#
i) of
        (# State# RealWorld
_, a
ret #) -> (# forall a. a -> Ur a
Ur a
ret, forall a. MutableArray# RealWorld a -> Array# a
Array# MutableArray# RealWorld a
arr #)
{-# NOINLINE get #-} -- prevents the runRW# effect from being reordered

set :: Int -> a -> Array# a %1 -> Array# a
set :: forall a. Int -> a -> Array# a %1 -> Array# a
set (GHC.I# Int#
i) (a
a :: a) = forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear Array# a -> Array# a
go
  where
    go :: Array# a -> Array# a
    go :: Array# a -> Array# a
go (Array# MutableArray# RealWorld a
arr) =
      case forall o. (State# RealWorld -> o) -> o
GHC.runRW# (forall d a. MutableArray# d a -> Int# -> a -> State# d -> State# d
GHC.writeArray# MutableArray# RealWorld a
arr Int#
i a
a) of
        State# RealWorld
_ -> forall a. MutableArray# RealWorld a -> Array# a
Array# MutableArray# RealWorld a
arr
{-# NOINLINE set #-} -- prevents the runRW# effect from being reordered

-- | Copy the first mutable array into the second mutable array, starting
-- from the given index of the source array.
--
-- It copies fewer elements if the second array is smaller than the
-- first. 'n' should be within [0..size src).
--
-- @
--  copyInto n src dest:
--   dest[i] = src[n+i] for i < size dest, i < size src + n
-- @
copyInto :: Int -> Array# a %1 -> Array# a %1 -> (# Array# a, Array# a #)
copyInto :: forall a.
Int -> Array# a %1 -> Array# a %1 -> (# Array# a, Array# a #)
copyInto start :: Int
start@(GHC.I# Int#
start#) = forall a b c (p :: Multiplicity) (q :: Multiplicity)
       (x :: Multiplicity) (y :: Multiplicity).
(a %p -> b %q -> c) %1 -> a %x -> b %y -> c
Unsafe.toLinear2 forall a. Array# a -> Array# a -> (# Array# a, Array# a #)
go
  where
    go :: Array# a -> Array# a -> (# Array# a, Array# a #)
    go :: forall a. Array# a -> Array# a -> (# Array# a, Array# a #)
go (Array# MutableArray# RealWorld a
src) (Array# MutableArray# RealWorld a
dst) =
      let !(GHC.I# Int#
len#) =
            forall a. Ord a => a -> a -> a
Prelude.min
              (Int# -> Int
GHC.I# (forall d a. MutableArray# d a -> Int#
GHC.sizeofMutableArray# MutableArray# RealWorld a
src) forall a. Num a => a -> a -> a
Prelude.- Int
start)
              (Int# -> Int
GHC.I# (forall d a. MutableArray# d a -> Int#
GHC.sizeofMutableArray# MutableArray# RealWorld a
dst))
       in case forall o. (State# RealWorld -> o) -> o
GHC.runRW# (forall d a.
MutableArray# d a
-> Int#
-> MutableArray# d a
-> Int#
-> Int#
-> State# d
-> State# d
GHC.copyMutableArray# MutableArray# RealWorld a
src Int#
start# MutableArray# RealWorld a
dst Int#
0# Int#
len#) of
            State# RealWorld
_ -> (# forall a. MutableArray# RealWorld a -> Array# a
Array# MutableArray# RealWorld a
src, forall a. MutableArray# RealWorld a -> Array# a
Array# MutableArray# RealWorld a
dst #)
{-# NOINLINE copyInto #-} -- prevents the runRW# effect from being reordered

map :: (a -> b) -> Array# a %1 -> Array# b
map :: forall a b. (a -> b) -> Array# a %1 -> Array# b
map (a -> b
f :: a -> b) =
  forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear
    ( \(Array# MutableArray# RealWorld a
as) ->
        let -- We alias the input array to write the resulting -- 'b's to,
            -- just to make the typechecker happy. Care must be taken to
            -- only read indices from 'as' that is not yet written to 'bs'.
            bs :: GHC.MutableArray# GHC.RealWorld b
            bs :: MutableArray# RealWorld b
bs = unsafeCoerce# :: forall a b. a -> b
GHC.unsafeCoerce# MutableArray# RealWorld a
as
            len :: GHC.Int#
            len :: Int#
len = forall d a. MutableArray# d a -> Int#
GHC.sizeofMutableArray# MutableArray# RealWorld a
as

            -- For each index ([0..len]), we read the element on 'as', pass
            -- it through 'f' and write to the same location on 'bs'.
            go :: GHC.Int# -> GHC.State# GHC.RealWorld -> ()
            go :: Int# -> State# RealWorld -> ()
go Int#
i State# RealWorld
st
              | Int# -> Int
GHC.I# Int#
i forall a. Eq a => a -> a -> Bool
Prelude.== Int# -> Int
GHC.I# Int#
len = ()
              | Bool
Prelude.otherwise =
                  case forall d a.
MutableArray# d a -> Int# -> State# d -> (# State# d, a #)
GHC.readArray# MutableArray# RealWorld a
as Int#
i State# RealWorld
st of
                    (# State# RealWorld
st', a
a #) ->
                      case forall d a. MutableArray# d a -> Int# -> a -> State# d -> State# d
GHC.writeArray# MutableArray# RealWorld b
bs Int#
i (a -> b
f a
a) State# RealWorld
st' of
                        !State# RealWorld
st'' -> Int# -> State# RealWorld -> ()
go (Int#
i Int# -> Int# -> Int#
GHC.+# Int#
1#) State# RealWorld
st''
         in forall o. (State# RealWorld -> o) -> o
GHC.runRW# (Int# -> State# RealWorld -> ()
go Int#
0#) seq :: forall a b. a -> b -> b
`GHC.seq` forall a. MutableArray# RealWorld a -> Array# a
Array# MutableArray# RealWorld b
bs
    )
{-# NOINLINE map #-}

-- | Return the array elements as a lazy list.
toList :: Array# a %1 -> Ur [a]
toList :: forall a. Array# a %1 -> Ur [a]
toList = forall a b. (MutableArray# RealWorld a -> b) -> Array# a %1 -> Ur b
unArray# forall a b. (a -> b) -> a -> b
Prelude.$ \MutableArray# RealWorld a
arr ->
  forall {a}. Int -> Int -> MutableArray# RealWorld a -> [a]
go
    Int
0
    (Int# -> Int
GHC.I# (forall d a. MutableArray# d a -> Int#
GHC.sizeofMutableArray# MutableArray# RealWorld a
arr))
    MutableArray# RealWorld a
arr
  where
    go :: Int -> Int -> MutableArray# RealWorld a -> [a]
go Int
i Int
len MutableArray# RealWorld a
arr
      | Int
i forall a. Eq a => a -> a -> Bool
Prelude.== Int
len = []
      | GHC.I# Int#
i# <- Int
i =
          case forall o. (State# RealWorld -> o) -> o
GHC.runRW# (forall d a.
MutableArray# d a -> Int# -> State# d -> (# State# d, a #)
GHC.readArray# MutableArray# RealWorld a
arr Int#
i#) of
            (# State# RealWorld
_, a
ret #) -> a
ret forall a. a -> [a] -> [a]
: Int -> Int -> MutableArray# RealWorld a -> [a]
go (Int
i forall a. Num a => a -> a -> a
Prelude.+ Int
1) Int
len MutableArray# RealWorld a
arr

-- | /O(1)/ Convert an 'Array#' to an immutable 'GHC.Array#'.
freeze :: (GHC.Array# a -> b) -> Array# a %1 -> Ur b
freeze :: forall a b. (Array# a -> b) -> Array# a %1 -> Ur b
freeze Array# a -> b
f = forall a b. (MutableArray# RealWorld a -> b) -> Array# a %1 -> Ur b
unArray# MutableArray# RealWorld a -> b
go
  where
    go :: MutableArray# RealWorld a -> b
go MutableArray# RealWorld a
mut =
      case forall o. (State# RealWorld -> o) -> o
GHC.runRW# (forall d a.
MutableArray# d a -> State# d -> (# State# d, Array# a #)
GHC.unsafeFreezeArray# MutableArray# RealWorld a
mut) of
        (# State# RealWorld
_, Array# a
ret #) -> Array# a -> b
f Array# a
ret

-- | Clone an array.
dup2 :: Array# a %1 -> (# Array# a, Array# a #)
dup2 :: forall a. Array# a %1 -> (# Array# a, Array# a #)
dup2 = forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear forall a. Array# a -> (# Array# a, Array# a #)
go
  where
    go :: Array# a -> (# Array# a, Array# a #)
    go :: forall a. Array# a -> (# Array# a, Array# a #)
go (Array# MutableArray# RealWorld a
arr) =
      case forall o. (State# RealWorld -> o) -> o
GHC.runRW#
        (forall d a.
MutableArray# d a
-> Int# -> Int# -> State# d -> (# State# d, MutableArray# d a #)
GHC.cloneMutableArray# MutableArray# RealWorld a
arr Int#
0# (forall d a. MutableArray# d a -> Int#
GHC.sizeofMutableArray# MutableArray# RealWorld a
arr)) of
        (# State# RealWorld
_, MutableArray# RealWorld a
new #) -> (# forall a. MutableArray# RealWorld a -> Array# a
Array# MutableArray# RealWorld a
arr, forall a. MutableArray# RealWorld a -> Array# a
Array# MutableArray# RealWorld a
new #)
{-# NOINLINE dup2 #-}