{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UnliftedNewtypes #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}

module Data.Ref.Linear.Unlifted (
  Ref#,
  newRef#,
  freeRef#,
  unsafeReadRef#,
  unsafeWriteRef#,
  atomicModify_#,
  atomicModify#,
) where

import Control.Monad.Borrow.Pure.Lifetime.Token
import Control.Monad.Borrow.Pure.Lifetime.Token.Unsafe (LinearOnly (..), LinearOnlyWitness (..))
import Control.Monad.Borrow.Pure.Utils (lseq#)
import GHC.Exts
import GHC.Exts qualified as GHC
import Prelude.Linear
import Unsafe.Linear qualified as Unsafe

newtype Ref# a = Ref# (MutVar# RealWorld a)

type role Ref# nominal

newRef# :: a %1 -> Linearly %1 -> Ref# a
{-# NOINLINE newRef# #-}
newRef# :: forall a. a %1 -> Linearly %1 -> Ref# a
newRef# = (a %1 -> Linearly %1 -> Ref# a) -> a %1 -> Linearly %1 -> Ref# a
forall a. a -> a
GHC.noinline ((a %1 -> Linearly %1 -> Ref# a) -> a %1 -> Linearly %1 -> Ref# a)
-> (a %1 -> Linearly %1 -> Ref# a) -> a %1 -> Linearly %1 -> Ref# a
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ (a -> Linearly %1 -> Ref# a) %1 -> a %1 -> Linearly %1 -> Ref# a
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear ((a -> Linearly %1 -> Ref# a) %1 -> a %1 -> Linearly %1 -> Ref# a)
-> (a -> Linearly %1 -> Ref# a) %1 -> a %1 -> Linearly %1 -> Ref# a
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ \a
a Linearly
lin ->
  Linearly
lin
    Linearly %1 -> Ref# a %1 -> Ref# a
forall a (s :: UnliftedType). Consumable a => a %1 -> s %1 -> s
`lseq#` (State# RealWorld -> Ref# a) -> Ref# a
forall o. (State# RealWorld -> o) -> o
GHC.runRW# \State# RealWorld
s ->
      case a
-> State# RealWorld -> (# State# RealWorld, MutVar# RealWorld a #)
forall a d. a -> State# d -> (# State# d, MutVar# d a #)
GHC.newMutVar# a
a State# RealWorld
s of
        (# !State# RealWorld
_, !MutVar# RealWorld a
v #) -> MutVar# RealWorld a -> Ref# a
forall a. MutVar# RealWorld a -> Ref# a
Ref# MutVar# RealWorld a
v

-- | This is unsafe, because the ownership of 'a' is duplicated.
unsafeReadRef# :: Ref# a %1 -> (# a, Ref# a #)
unsafeReadRef# :: forall a. Ref# a %1 -> (# a, Ref# a #)
unsafeReadRef# = (Ref# a %1 -> (# a, Ref# a #)) -> Ref# a %1 -> (# a, Ref# a #)
forall a. a -> a
GHC.noinline ((Ref# a %1 -> (# a, Ref# a #)) -> Ref# a %1 -> (# a, Ref# a #))
-> (Ref# a %1 -> (# a, Ref# a #)) -> Ref# a %1 -> (# a, Ref# a #)
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ (Ref# a -> (# a, Ref# a #)) %1 -> Ref# a %1 -> (# a, Ref# a #)
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear \(Ref# !MutVar# RealWorld a
mv) ->
  (State# RealWorld -> (# a, Ref# a #)) -> (# a, Ref# a #)
forall o. (State# RealWorld -> o) -> o
runRW# \State# RealWorld
s ->
    case MutVar# RealWorld a
-> State# RealWorld -> (# State# RealWorld, a #)
forall d a. MutVar# d a -> State# d -> (# State# d, a #)
GHC.readMutVar# MutVar# RealWorld a
mv State# RealWorld
s of
      (# !State# RealWorld
_, !a
a #) -> (# a
a, MutVar# RealWorld a -> Ref# a
forall a. MutVar# RealWorld a -> Ref# a
Ref# MutVar# RealWorld a
mv #)

-- | This is unsafe, because the ownership of original 'a' is dropped.
unsafeWriteRef# :: Ref# a %1 -> a %1 -> Ref# a
{-# NOINLINE unsafeWriteRef# #-}
unsafeWriteRef# :: forall a. Ref# a %1 -> a %1 -> Ref# a
unsafeWriteRef# = (Ref# a %1 -> a %1 -> Ref# a) -> Ref# a %1 -> a %1 -> Ref# a
forall a. a -> a
GHC.noinline ((Ref# a %1 -> a %1 -> Ref# a) -> Ref# a %1 -> a %1 -> Ref# a)
-> (Ref# a %1 -> a %1 -> Ref# a) -> Ref# a %1 -> a %1 -> Ref# a
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ (Ref# a -> a -> Ref# a) %1 -> Ref# a %1 -> a %1 -> Ref# 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 \(Ref# MutVar# RealWorld a
mv) !a
a ->
  (State# RealWorld -> Ref# a) -> Ref# a
forall o. (State# RealWorld -> o) -> o
runRW# \State# RealWorld
s ->
    case MutVar# RealWorld a -> a -> State# RealWorld -> State# RealWorld
forall d a. MutVar# d a -> a -> State# d -> State# d
GHC.writeMutVar# MutVar# RealWorld a
mv a
a State# RealWorld
s of
      State# RealWorld
_ -> MutVar# RealWorld a -> Ref# a
forall a. MutVar# RealWorld a -> Ref# a
Ref# MutVar# RealWorld a
mv

freeRef# :: Ref# a %1 -> a
{-# NOINLINE freeRef# #-}
freeRef# :: forall a. Ref# a %1 -> a
freeRef# = (Ref# a -> a) %1 -> Ref# a %1 -> a
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear \(Ref# MutVar# RealWorld a
a) ->
  (State# RealWorld -> a) -> a
forall o. (State# RealWorld -> o) -> o
runRW# \State# RealWorld
s ->
    case MutVar# RealWorld a
-> State# RealWorld -> (# State# RealWorld, a #)
forall d a. MutVar# d a -> State# d -> (# State# d, a #)
GHC.readMutVar# MutVar# RealWorld a
a State# RealWorld
s of
      (# State# RealWorld
_, !a
a #) -> a
a

instance LinearOnly (Ref# a) where
  linearOnly :: LinearOnlyWitness (Ref# a)
linearOnly = LinearOnlyWitness (Ref# a)
forall {k} (a :: k). LinearOnlyWitness a
UnsafeLinearOnly

atomicModify_# :: (a %1 -> a) %1 -> Ref# a %1 -> Ref# a
{-# NOINLINE atomicModify_# #-}
atomicModify_# :: forall a. (a %1 -> a) %1 -> Ref# a %1 -> Ref# a
atomicModify_# = ((a %1 -> a) %1 -> Ref# a %1 -> Ref# a)
-> (a %1 -> a) %1 -> Ref# a %1 -> Ref# a
forall a. a -> a
GHC.noinline (((a %1 -> a) %1 -> Ref# a %1 -> Ref# a)
 -> (a %1 -> a) %1 -> Ref# a %1 -> Ref# a)
-> ((a %1 -> a) %1 -> Ref# a %1 -> Ref# a)
-> (a %1 -> a)
%1 -> Ref# a
%1 -> Ref# a
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ ((a %1 -> a) -> Ref# a -> Ref# a)
%1 -> (a %1 -> a) %1 -> Ref# a %1 -> Ref# 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 \a %1 -> a
f (Ref# MutVar# RealWorld a
mv) ->
  (State# RealWorld -> Ref# a) -> Ref# a
forall o. (State# RealWorld -> o) -> o
runRW# \State# RealWorld
s ->
    case MutVar# RealWorld a
-> (a -> a) -> State# RealWorld -> (# State# RealWorld, a, a #)
forall d a c.
MutVar# d a -> (a -> c) -> State# d -> (# State# d, a, c #)
GHC.atomicModifyMutVar2# MutVar# RealWorld a
mv ((a %1 -> a) %1 -> a -> a
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear a %1 -> a
f) State# RealWorld
s of
      (# State# RealWorld
_, !a
_, !a
_ #) -> MutVar# RealWorld a -> Ref# a
forall a. MutVar# RealWorld a -> Ref# a
Ref# MutVar# RealWorld a
mv

atomicModify# :: (a %1 -> (b, a)) %1 -> Ref# a %1 -> (# b, Ref# a #)
{-# NOINLINE atomicModify# #-}
atomicModify# :: forall a b. (a %1 -> (b, a)) %1 -> Ref# a %1 -> (# b, Ref# a #)
atomicModify# = ((a %1 -> (b, a)) %1 -> Ref# a %1 -> (# b, Ref# a #))
-> (a %1 -> (b, a)) %1 -> Ref# a %1 -> (# b, Ref# a #)
forall a. a -> a
GHC.noinline (((a %1 -> (b, a)) %1 -> Ref# a %1 -> (# b, Ref# a #))
 -> (a %1 -> (b, a)) %1 -> Ref# a %1 -> (# b, Ref# a #))
-> ((a %1 -> (b, a)) %1 -> Ref# a %1 -> (# b, Ref# a #))
-> (a %1 -> (b, a))
%1 -> Ref# a
%1 -> (# b, Ref# a #)
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ ((a %1 -> (b, a)) -> Ref# a -> (# b, Ref# a #))
%1 -> (a %1 -> (b, a)) %1 -> Ref# a %1 -> (# b, Ref# 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 \a %1 -> (b, a)
f (Ref# MutVar# RealWorld a
mv) ->
  (State# RealWorld -> (# b, Ref# a #)) -> (# b, Ref# a #)
forall o. (State# RealWorld -> o) -> o
runRW# \State# RealWorld
s ->
    case MutVar# RealWorld a
-> (a -> (b, a))
-> State# RealWorld
-> (# State# RealWorld, a, (b, a) #)
forall d a c.
MutVar# d a -> (a -> c) -> State# d -> (# State# d, a, c #)
GHC.atomicModifyMutVar2# MutVar# RealWorld a
mv ((a %1 -> (b, a)) %1 -> a -> (b, a)
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear a %1 -> (b, a)
f) State# RealWorld
s of
      (# State# RealWorld
_, !a
_, (!b
b, !a
_) #) -> (# b
b, MutVar# RealWorld a -> Ref# a
forall a. MutVar# RealWorld a -> Ref# a
Ref# MutVar# RealWorld a
mv #)