{-# 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 = (Array# a -> Ur b) %1 -> Array# a %1 -> Ur b
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear (\(Array# MutableArray# RealWorld a
a) -> b -> Ur b
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 = (Array# a -> b -> b) %1 -> Array# a %1 -> b %1 -> b
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 :: (Movable b) => Int -> a -> (Array# a %1 -> b) %1 -> b
alloc :: forall b a. Movable b => Int -> a -> (Array# a %1 -> b) %1 -> b
alloc Int
i a
a Array# a %1 -> b
f = case b %1 -> Ur b
forall a. Movable a => a %1 -> Ur a
move (Int -> a -> (Array# a %1 -> b) %1 -> b
forall a b. Int -> a -> (Array# a %1 -> b) %1 -> b
unsafe_alloc Int
i a
a Array# a %1 -> b
f) of
  Ur b
b -> b
b
{-# INLINEABLE alloc #-}

-- The `alloc` function is split in two. One very unsafe below (it's very
-- unsafe, because `unafe_alloc 57 0 id` returns an unrestricted _mutable_
-- `Array#` breaking the module's invariants). Because `unsafe_alloc` calls
-- `runRW#`, it's marked as `NOINLINE`.
--
-- It's made safe by the wrapping function `alloc`, which restricts `b` to be
-- `Movable` (`Array#` is crucially not `Movable`, therefore `alloc 57 0 id`
-- doesn't type). Furthermore, `alloc` cases on `move` to make sure that all the
-- effects have been run by the time we evaluate the result of an `alloc`. It's
-- fine that `alloc` is inlined: its semantics is preserved by program
-- transformations. It's useful that `alloc` be inlined, because in most
-- instance `case move … of` will trigger a case-of-known-constructor avoiding
-- an extra allocation. This is in particular the case for the common case where
-- `b = Ur x`.
unsafe_alloc :: Int -> a -> (Array# a %1 -> b) %1 -> b
unsafe_alloc :: forall a b. Int -> a -> (Array# a %1 -> b) %1 -> b
unsafe_alloc (GHC.I# Int#
s) a
a Array# a %1 -> b
f =
  let new :: Array# a
new = (State# RealWorld -> Array# a) -> Array# a
forall o. (State# RealWorld -> o) -> o
GHC.runRW# ((State# RealWorld -> Array# a) -> Array# a)
-> (State# RealWorld -> Array# a) -> Array# a
forall a b. (a -> b) -> a -> b
Prelude.$ \State# RealWorld
st ->
        case Int#
-> a
-> State# RealWorld
-> (# State# RealWorld, MutableArray# RealWorld a #)
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 #) -> MutableArray# RealWorld a -> Array# a
forall a. MutableArray# RealWorld a -> Array# a
Array# MutableArray# RealWorld a
arr
   in Array# a %1 -> b
f Array# a
new
{-# NOINLINE unsafe_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 = (State# RealWorld -> Array# a) -> Array# a
forall o. (State# RealWorld -> o) -> o
GHC.runRW# ((State# RealWorld -> Array# a) -> Array# a)
-> (State# RealWorld -> Array# a) -> Array# a
forall a b. (a -> b) -> a -> b
Prelude.$ \State# RealWorld
st ->
        case Int#
-> a
-> State# RealWorld
-> (# State# RealWorld, MutableArray# RealWorld a #)
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 #) -> MutableArray# RealWorld a -> Array# a
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 = (Array# a -> (# Ur Int, Array# a #))
%1 -> Array# a %1 -> (# Ur Int, Array# a #)
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear Array# a -> (# Ur Int, Array# a #)
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 = MutableArray# RealWorld a -> Int#
forall d a. MutableArray# d a -> Int#
GHC.sizeofMutableArray# MutableArray# RealWorld a
arr
       in (# Int -> Ur Int
forall a. a -> Ur a
Ur (Int# -> Int
GHC.I# Int#
s), MutableArray# RealWorld a -> Array# a
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) = (Array# a -> (# Ur a, Array# a #))
%1 -> Array# a %1 -> (# Ur a, Array# a #)
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear Array# a -> (# Ur a, Array# a #)
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 (State# RealWorld -> (# State# RealWorld, a #))
-> (# State# RealWorld, a #)
forall o. (State# RealWorld -> o) -> o
GHC.runRW# (MutableArray# RealWorld a
-> Int# -> State# RealWorld -> (# State# RealWorld, a #)
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 -> Ur a
forall a. a -> Ur a
Ur a
ret, MutableArray# RealWorld a -> Array# a
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) = (Array# a -> Array# a) %1 -> Array# a %1 -> Array# 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 (State# RealWorld -> State# RealWorld) -> State# RealWorld
forall o. (State# RealWorld -> o) -> o
GHC.runRW# (MutableArray# RealWorld a
-> Int# -> a -> State# RealWorld -> State# RealWorld
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
_ -> MutableArray# RealWorld a -> Array# a
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#) = (Array# a -> Array# a -> (# Array# a, Array# a #))
%1 -> Array# a %1 -> Array# a %1 -> (# Array# a, Array# a #)
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 -> Array# a -> (# Array# a, Array# a #)
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#) =
            Int -> Int -> Int
forall a. Ord a => a -> a -> a
Prelude.min
              (Int# -> Int
GHC.I# (MutableArray# RealWorld a -> Int#
forall d a. MutableArray# d a -> Int#
GHC.sizeofMutableArray# MutableArray# RealWorld a
src) Int -> Int -> Int
forall a. Num a => a -> a -> a
Prelude.- Int
start)
              (Int# -> Int
GHC.I# (MutableArray# RealWorld a -> Int#
forall d a. MutableArray# d a -> Int#
GHC.sizeofMutableArray# MutableArray# RealWorld a
dst))
       in case (State# RealWorld -> State# RealWorld) -> State# RealWorld
forall o. (State# RealWorld -> o) -> o
GHC.runRW# (MutableArray# RealWorld a
-> Int#
-> MutableArray# RealWorld a
-> Int#
-> Int#
-> State# RealWorld
-> State# RealWorld
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
_ -> (# MutableArray# RealWorld a -> Array# a
forall a. MutableArray# RealWorld a -> Array# a
Array# MutableArray# RealWorld a
src, MutableArray# RealWorld a -> Array# a
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) =
  (Array# a -> Array# b) %1 -> Array# a %1 -> Array# 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 = MutableArray# RealWorld a -> MutableArray# RealWorld b
forall a b. a -> b
GHC.unsafeCoerce# MutableArray# RealWorld a
as
            len :: GHC.Int#
            len :: Int#
len = MutableArray# RealWorld a -> Int#
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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
Prelude.== Int# -> Int
GHC.I# Int#
len = ()
              | Bool
Prelude.otherwise =
                  case MutableArray# RealWorld a
-> Int# -> State# RealWorld -> (# State# RealWorld, a #)
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 MutableArray# RealWorld b
-> Int# -> b -> State# RealWorld -> State# RealWorld
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 (State# RealWorld -> ()) -> ()
forall o. (State# RealWorld -> o) -> o
GHC.runRW# (Int# -> State# RealWorld -> ()
go Int#
0#) () -> Array# b -> Array# b
forall a b. a -> b -> b
`GHC.seq` MutableArray# RealWorld b -> Array# b
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 = (MutableArray# RealWorld a -> [a]) -> Array# a %1 -> Ur [a]
forall a b. (MutableArray# RealWorld a -> b) -> Array# a %1 -> Ur b
unArray# ((MutableArray# RealWorld a -> [a]) -> Array# a %1 -> Ur [a])
-> (MutableArray# RealWorld a -> [a]) -> Array# a %1 -> Ur [a]
forall a b. (a -> b) -> a -> b
Prelude.$ \MutableArray# RealWorld a
arr ->
  Int -> Int -> MutableArray# RealWorld a -> [a]
forall {a}. Int -> Int -> MutableArray# RealWorld a -> [a]
go
    Int
0
    (Int# -> Int
GHC.I# (MutableArray# RealWorld a -> Int#
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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
Prelude.== Int
len = []
      | GHC.I# Int#
i# <- Int
i =
          case (State# RealWorld -> (# State# RealWorld, a #))
-> (# State# RealWorld, a #)
forall o. (State# RealWorld -> o) -> o
GHC.runRW# (MutableArray# RealWorld a
-> Int# -> State# RealWorld -> (# State# RealWorld, a #)
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 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Int -> Int -> MutableArray# RealWorld a -> [a]
go (Int
i Int -> Int -> Int
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 = (MutableArray# RealWorld a -> b) -> Array# a %1 -> Ur b
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 (State# RealWorld -> (# State# RealWorld, Array# a #))
-> (# State# RealWorld, Array# a #)
forall o. (State# RealWorld -> o) -> o
GHC.runRW# (MutableArray# RealWorld a
-> State# RealWorld -> (# State# RealWorld, Array# a #)
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 = (Array# a -> (# Array# a, Array# a #))
%1 -> Array# a %1 -> (# Array# a, Array# a #)
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear Array# a -> (# Array# a, Array# a #)
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 (State# RealWorld
 -> (# State# RealWorld, MutableArray# RealWorld a #))
-> (# State# RealWorld, MutableArray# RealWorld a #)
forall o. (State# RealWorld -> o) -> o
GHC.runRW#
        (MutableArray# RealWorld a
-> Int#
-> Int#
-> State# RealWorld
-> (# State# RealWorld, MutableArray# RealWorld a #)
forall d a.
MutableArray# d a
-> Int# -> Int# -> State# d -> (# State# d, MutableArray# d a #)
GHC.cloneMutableArray# MutableArray# RealWorld a
arr Int#
0# (MutableArray# RealWorld a -> Int#
forall d a. MutableArray# d a -> Int#
GHC.sizeofMutableArray# MutableArray# RealWorld a
arr)) of
        (# State# RealWorld
_, MutableArray# RealWorld a
new #) -> (# MutableArray# RealWorld a -> Array# a
forall a. MutableArray# RealWorld a -> Array# a
Array# MutableArray# RealWorld a
arr, MutableArray# RealWorld a -> Array# a
forall a. MutableArray# RealWorld a -> Array# a
Array# MutableArray# RealWorld a
new #)
{-# NOINLINE dup2 #-}