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

-- |
-- 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 (lseq, dup2)
import Prelude (Int)
import qualified Prelude as Prelude
import qualified Unsafe.Linear as Unsafe
import qualified GHC.Exts as GHC

-- | 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). (a %p -> b) %1 -> a %1 -> 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).
(a %p -> b %q -> c) %1 -> a %1 -> b %1 -> c
Unsafe.toLinear2 (\Array# a
_ b
b -> b
b)

-- | 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 = (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 -> 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 = (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). (a %p -> b) %1 -> a %1 -> 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). (a %p -> b) %1 -> a %1 -> 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). (a %p -> b) %1 -> a %1 -> 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).
(a %p -> b %q -> c) %1 -> a %1 -> b %1 -> 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
arr =
  Array# a %1 -> (# Ur Int, Array# a #)
forall a. Array# a %1 -> (# Ur Int, Array# a #)
size Array# a
arr
    (# Ur Int, Array# a #)
%1 -> ((# Ur Int, Array# a #) %1 -> Array# b) %1 -> Array# b
forall a b c.
(# b, Array# a #) %1 -> ((# b, Array# a #) %1 -> c) %1 -> c
`chain2` \(# Ur Int
s, Array# a
arr' #) -> Int -> Int -> Array# a %1 -> Array# b
go Int
0 Int
s Array# a
arr'
 where
  -- When we're mapping an array, we first insert `b`'s
  -- inside an `Array# a` by unsafeCoerce'ing, and then we
  -- unsafeCoerce the result to an `Array# b`.
  go :: Int -> Int -> Array# a %1-> Array# b
  go :: Int -> Int -> Array# a %1 -> Array# b
go Int
i Int
s Array# a
arr'
    | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
Prelude.== Int
s =
        (Array# a -> Array# b) %1 -> Array# a %1 -> Array# b
forall a b (p :: Multiplicity). (a %p -> b) %1 -> a %1 -> b
Unsafe.toLinear Array# a -> Array# b
GHC.unsafeCoerce# Array# a
arr'
    | Bool
Prelude.otherwise =
        Int -> Array# a %1 -> (# Ur a, Array# a #)
forall a. Int -> Array# a %1 -> (# Ur a, Array# a #)
get Int
i Array# a
arr'
          (# Ur a, Array# a #)
%1 -> ((# Ur a, Array# a #) %1 -> Array# b) %1 -> Array# b
forall a b c.
(# b, Array# a #) %1 -> ((# b, Array# a #) %1 -> c) %1 -> c
`chain2` \(# Ur a
a, Array# a
arr'' #) -> Int -> a -> Array# a %1 -> Array# a
forall a. Int -> a -> Array# a %1 -> Array# a
set Int
i (b %1 -> a
forall a b. a %1 -> b
Unsafe.coerce (a -> b
f a
a)) Array# a
arr''
          Array# a %1 -> (Array# a %1 -> Array# b) %1 -> Array# b
forall a b. Array# a %1 -> (Array# a %1 -> b) %1 -> b
`chain` \Array# a
arr''' -> Int -> Int -> Array# a %1 -> Array# b
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
Prelude.+ Int
1) Int
s Array# a
arr'''
{-# 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). (a %p -> b) %1 -> a %1 -> 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 #-}

-- * Internal library

-- Below two are variants of (&) specialized for taking commonly used
-- unlifted values and returning a levity-polymorphic result.
--
-- They are not polymorphic on their first parameter since levity-polymorphism
-- disallows binding to levity-polymorphic values.

chain :: forall (r :: GHC.RuntimeRep) a (b :: GHC.TYPE r).
        Array# a %1-> (Array# a %1-> b) %1-> b
chain :: forall a b. Array# a %1 -> (Array# a %1 -> b) %1 -> b
chain Array# a
a Array# a %1 -> b
f = Array# a %1 -> b
f Array# a
a

chain2 :: forall (r :: GHC.RuntimeRep) a b (c :: GHC.TYPE r).
        (# b, Array# a #) %1-> ((# b, Array# a #) %1-> c) %1-> c
chain2 :: forall a b c.
(# b, Array# a #) %1 -> ((# b, Array# a #) %1 -> c) %1 -> c
chain2 (# b, Array# a #)
a (# b, Array# a #) %1 -> c
f = (# b, Array# a #) %1 -> c
f (# b, Array# a #)
a
infixl 1 `chain`, `chain2`