{-# language BangPatterns #-}
{-# language GADTSyntax #-}
{-# language MagicHash #-}
{-# language RankNTypes #-}
{-# language ScopedTypeVariables #-}
{-# language StandaloneKindSignatures #-}
{-# language TypeApplications #-}
{-# language TypeFamilies #-}
{-# language TypeInType #-}
{-# language UnboxedTuples #-}
{-# language UnliftedNewtypes #-}

module Unlifted
  ( R
  , A#
  , ArrayRep
  , M#
  , empty#
  , index#
  , write#
  , read#
  , size#
  , unsafeFreeze#
  , set#
  , unsafeShrinkFreeze#
  , thaw#
  , initialized#
  , freeze#
  , copy#
  ) where

import GHC.Exts
import Data.Kind (Type)
import Unsafe.Coerce (unsafeCoerceUnlifted)

import qualified GHC.Exts as Exts

type ArrayRep = 'BoxedRep 'Unlifted
type R = 'BoxedRep 'Unlifted

type A# :: TYPE ('BoxedRep 'Unlifted) -> TYPE ('BoxedRep 'Unlifted)
type A# = Array#

type M# :: Type -> TYPE ('BoxedRep 'Unlifted) -> TYPE ('BoxedRep 'Unlifted)
type M# = MutableArray#

size# :: forall (a :: TYPE R). A# a -> Int#
size# :: forall (a :: TYPE R). A# a -> Int#
size# = Array# a -> Int#
forall a. Array# a -> Int#
sizeofArray#

index# :: forall (a :: TYPE R). A# a -> Int# -> a
index# :: forall (a :: TYPE R). A# a -> Int# -> a
index# A# a
a Int#
i = case A# a -> Int# -> (# a #)
forall a. Array# a -> Int# -> (# a #)
indexArray# A# a
a Int#
i of
  (# a
r #) -> a
r

write# :: forall (s :: Type) (a :: TYPE R).
  M# s a -> Int# -> a -> State# s -> State# s
write# :: forall s (a :: TYPE R). M# s a -> Int# -> a -> State# s -> State# s
write# = MutableArray# s a -> Int# -> a -> State# s -> State# s
forall d a. MutableArray# d a -> Int# -> a -> State# d -> State# d
writeArray#

read# :: forall (s :: Type) (a :: TYPE R).
  M# s a -> Int# -> State# s -> (# State# s, a #)
read# :: forall s (a :: TYPE R).
M# s a -> Int# -> State# s -> (# State# s, a #)
read# = MutableArray# s a -> Int# -> State# s -> (# State# s, a #)
forall d a.
MutableArray# d a -> Int# -> State# d -> (# State# d, a #)
readArray#

unsafeFreeze# :: forall (s :: Type) (a :: TYPE R).
     M# s a
  -> State# s
  -> (# State# s, A# a #)
unsafeFreeze# :: forall s (a :: TYPE R). M# s a -> State# s -> (# State# s, A# a #)
unsafeFreeze# = MutableArray# s a -> State# s -> (# State# s, Array# a #)
forall d a.
MutableArray# d a -> State# d -> (# State# d, Array# a #)
unsafeFreezeArray#

initialized# :: forall (s :: Type) (a :: TYPE R).
     Int#
  -> a
  -> State# s
  -> (# State# s, M# s a #)
initialized# :: forall s (a :: TYPE R).
Int# -> a -> State# s -> (# State# s, M# s a #)
initialized# Int#
i a
a State# s
s = Int# -> a -> State# s -> (# State# s, MutableArray# s a #)
forall a d.
Int# -> a -> State# d -> (# State# d, MutableArray# d a #)
newArray# Int#
i a
a State# s
s

-- This implementation is ridiculous, but GHC does not currently give
-- us a way to allocate an empty array of unlifted elements without
-- supplying an element.
empty# :: forall (a :: TYPE R). (# #) -> A# a
empty# :: forall (a :: TYPE R). (# #) -> A# a
empty# (# #)
_ = 
  let !(# State# RealWorld
_, Array# a
z #) = (State# RealWorld -> (# State# RealWorld, Array# a #))
-> (# State# RealWorld, Array# a #)
forall o. (State# RealWorld -> o) -> o
Exts.runRW#
        (\State# RealWorld
s0 -> case Int#
-> State# RealWorld
-> (# State# RealWorld, MutableByteArray# RealWorld #)
forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
Exts.newByteArray# Int#
0# State# RealWorld
s0 of
          (# State# RealWorld
s1, MutableByteArray# RealWorld
placeholder #) -> case Int#
-> a
-> State# RealWorld
-> (# State# RealWorld, MutableArray# RealWorld a #)
forall a d.
Int# -> a -> State# d -> (# State# d, MutableArray# d a #)
Exts.newArray# Int#
0# (forall (a :: TYPE R) (b :: TYPE R). a -> b
unsafeCoerceUnlifted @_ @a MutableByteArray# RealWorld
placeholder) State# RealWorld
s1 of
            (# State# RealWorld
s2, MutableArray# RealWorld a
x #) -> MutableArray# RealWorld a
-> State# RealWorld -> (# State# RealWorld, Array# a #)
forall d a.
MutableArray# d a -> State# d -> (# State# d, Array# a #)
Exts.unsafeFreezeArray# MutableArray# RealWorld a
x State# RealWorld
s2
        )
   in Array# a
z

set# :: forall (s :: Type) (a :: TYPE R).
     M# s a
  -> Int#
  -> Int#
  -> a
  -> State# s
  -> State# s
set# :: forall s (a :: TYPE R).
M# s a -> Int# -> Int# -> a -> State# s -> State# s
set# M# s a
m Int#
off0 Int#
len0 a
a State# s
s0 =
  let go :: Int# -> Int# -> State# s -> State# s
go Int#
off Int#
len State# s
s = case Int#
len of
        Int#
0# -> State# s
s
        Int#
_ -> Int# -> Int# -> State# s -> State# s
go (Int#
off Int# -> Int# -> Int#
+# Int#
1#) (Int#
len Int# -> Int# -> Int#
-# Int#
1#) (M# s a -> Int# -> a -> State# s -> State# s
forall s (a :: TYPE R). M# s a -> Int# -> a -> State# s -> State# s
write# M# s a
m Int#
off a
a State# s
s)
   in Int# -> Int# -> State# s -> State# s
go Int#
off0 Int#
len0 State# s
s0

-- shrink and freeze, all at once
unsafeShrinkFreeze# ::
     M# s a
  -> Int#
  -> State# s
  -> (# State# s, A# a #)
unsafeShrinkFreeze# :: forall s (a :: TYPE R).
M# s a -> Int# -> State# s -> (# State# s, A# a #)
unsafeShrinkFreeze# M# s a
m Int#
i State# s
s0 = case M# s a -> Int#
forall d a. MutableArray# d a -> Int#
sizeofMutableArray# M# s a
m Int# -> Int# -> Int#
==# Int#
i of
  Int#
1# -> M# s a -> State# s -> (# State# s, A# a #)
forall d a.
MutableArray# d a -> State# d -> (# State# d, Array# a #)
Exts.unsafeFreezeArray# M# s a
m State# s
s0
  Int#
_ -> M# s a -> Int# -> Int# -> State# s -> (# State# s, A# a #)
forall d a.
MutableArray# d a
-> Int# -> Int# -> State# d -> (# State# d, Array# a #)
Exts.freezeArray# M# s a
m Int#
0# Int#
i State# s
s0

thaw# :: forall (s :: Type) (a :: TYPE R).
     A# a
  -> Int#
  -> Int#
  -> State# s
  -> (# State# s, M# s a #)
thaw# :: forall s (a :: TYPE R).
A# a -> Int# -> Int# -> State# s -> (# State# s, M# s a #)
thaw# = Array# a
-> Int# -> Int# -> State# s -> (# State# s, MutableArray# s a #)
forall a d.
Array# a
-> Int# -> Int# -> State# d -> (# State# d, MutableArray# d a #)
Exts.thawArray#

freeze# :: forall (s :: Type) (a :: TYPE R).
     M# s a
  -> Int#
  -> Int#
  -> State# s
  -> (# State# s, A# a #)
freeze# :: forall s (a :: TYPE R).
M# s a -> Int# -> Int# -> State# s -> (# State# s, A# a #)
freeze# = MutableArray# s a
-> Int# -> Int# -> State# s -> (# State# s, Array# a #)
forall d a.
MutableArray# d a
-> Int# -> Int# -> State# d -> (# State# d, Array# a #)
Exts.freezeArray#

copy# :: forall (s :: Type) (a :: TYPE R).
     M# s a
  -> Int#
  -> A# a
  -> Int#
  -> Int#
  -> State# s
  -> State# s
copy# :: forall s (a :: TYPE R).
M# s a -> Int# -> A# a -> Int# -> Int# -> State# s -> State# s
copy# M# s a
m Int#
doff A# a
v Int#
soff Int#
len State# s
s0 =
  A# a -> Int# -> M# s a -> Int# -> Int# -> State# s -> State# s
forall a d.
Array# a
-> Int#
-> MutableArray# d a
-> Int#
-> Int#
-> State# d
-> State# d
Exts.copyArray# A# a
v Int#
soff M# s a
m Int#
doff Int#
len State# s
s0