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

module Data.Ref.Linear (
  Ref,
  new,
  free,
  unsafeReadRef,
  unsafeWriteRef,
  atomicModify,
  atomicModify_,
) where

import Control.Functor.Linear qualified as Control
import Control.Monad.Borrow.Pure.Affine
import Control.Monad.Borrow.Pure.Affine.Unsafe (unsafeAff)
import Control.Monad.Borrow.Pure.BO
import Control.Monad.Borrow.Pure.BO.Unsafe (Alias (..))
import Control.Monad.Borrow.Pure.Clone
import Control.Monad.Borrow.Pure.Copyable
import Control.Monad.Borrow.Pure.Lifetime.Token.Internal (
  LinearOnly (..),
  LinearOnlyWitness (..),
 )
import Data.Ref.Linear.Unlifted
import GHC.TypeError
import Prelude.Linear (Consumable (..), Dupable (..))
import Prelude.Linear qualified as PL
import Unsafe.Linear qualified as Unsafe

-- | Linearly owned mutable reference.
data Ref a = Ref (Ref# a)

type role Ref nominal

new :: a %1 -> Linearly %1 -> Ref a
{-# INLINE new #-}
new :: forall a. a %1 -> Linearly %1 -> Ref a
new a
a Linearly
lin = Ref# a -> Ref a
forall a. Ref# a -> Ref a
Ref (a %1 -> Linearly %1 -> Ref# a
forall a. a %1 -> Linearly %1 -> Ref# a
newRef# a
a Linearly
lin)

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

instance (Consumable a) => Consumable (Ref a) where
  consume :: Ref a %1 -> ()
consume = a %1 -> ()
forall a. Consumable a => a %1 -> ()
consume (a %1 -> ()) -> (Ref a %1 -> a) -> Ref a %1 -> ()
forall b c a (q :: Multiplicity) (m :: Multiplicity)
       (n :: Multiplicity).
(b %1 -> c) %q -> (a %1 -> b) %m -> a %n -> c
PL.. Ref a %1 -> a
forall a. Ref a %1 -> a
free
  {-# INLINE consume #-}

instance (PL.Dupable a) => PL.Dupable (Ref a) where
  dup2 :: Ref a %1 -> (Ref a, Ref a)
dup2 = (Ref a -> (Ref a, Ref a)) %1 -> Ref a %1 -> (Ref a, Ref a)
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear \ !Ref a
v ->
    Ref a %1 -> (Linearly, Ref a)
forall a. LinearOnly a => a %1 -> (Linearly, a)
withLinearly Ref a
v (Linearly, Ref a)
-> ((Linearly, Ref a) -> (Ref a, Ref a)) -> (Ref a, Ref a)
forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
PL.& \(Linearly
l, !Ref a
v) ->
      let !v2 :: a
v2 = ((a, a) -> a) %1 -> (a, a) -> a
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear (\(!a
_, !a
v) -> a
v) ((a, a) -> a) -> (a, a) -> a
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
PL.$ a %1 -> (a, a)
forall a. Dupable a => a %1 -> (a, a)
dup2 (a %1 -> (a, a)) -> a %1 -> (a, a)
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
PL.$ Ref a %1 -> a
forall a. Ref a %1 -> a
free Ref a
v
       in (Ref a
v, a %1 -> Linearly %1 -> Ref a
forall a. a %1 -> Linearly %1 -> Ref a
new a
v2 Linearly
l)
  {-# INLINE dup2 #-}

instance Affine (Ref a) where
  aff :: Ref a %1 -> Aff (Ref a)
aff = Ref a %1 -> Aff (Ref a)
forall a. a %1 -> Aff a
unsafeAff

atomicModify_ :: (a %1 -> a) %1 -> Ref a %1 -> Ref a
{-# INLINE atomicModify_ #-}
atomicModify_ :: forall a. (a %1 -> a) %1 -> Ref a %1 -> Ref a
atomicModify_ a %1 -> a
f (Ref Ref# a
v) = Ref# a -> Ref a
forall a. Ref# a -> Ref a
Ref ((a %1 -> a) %1 -> Ref# a %1 -> Ref# a
forall a. (a %1 -> a) %1 -> Ref# a %1 -> Ref# a
atomicModify_# a %1 -> a
f Ref# a
v)

atomicModify :: (a %1 -> (b, a)) %1 -> Ref a %1 -> (b, Ref a)
{-# INLINE atomicModify #-}
atomicModify :: forall a b. (a %1 -> (b, a)) %1 -> Ref a %1 -> (b, Ref a)
atomicModify a %1 -> (b, a)
f (Ref Ref# a
v) = case (a %1 -> (b, a)) %1 -> Ref# a %1 -> (# b, Ref# a #)
forall a b. (a %1 -> (b, a)) %1 -> Ref# a %1 -> (# b, Ref# a #)
atomicModify# a %1 -> (b, a)
f Ref# a
v of
  (# b
b, Ref# a
v' #) -> (b
b, Ref# a -> Ref a
forall a. Ref# a -> Ref a
Ref Ref# a
v')

free :: Ref a %1 -> a
{-# INLINE free #-}
free :: forall a. Ref a %1 -> a
free (Ref Ref# a
v) = Ref# a %1 -> a
forall a. Ref# a %1 -> a
freeRef# Ref# a
v

unsafeReadRef :: Ref a %1 -> (a, Ref a)
{-# INLINE unsafeReadRef #-}
unsafeReadRef :: forall a. Ref a %1 -> (a, Ref a)
unsafeReadRef (Ref Ref# a
v) = case Ref# a %1 -> (# a, Ref# a #)
forall a. Ref# a %1 -> (# a, Ref# a #)
unsafeReadRef# Ref# a
v of
  (# a
a, Ref# a
v' #) -> (a
a, Ref# a -> Ref a
forall a. Ref# a -> Ref a
Ref Ref# a
v')

unsafeWriteRef :: Ref a %1 -> a %1 -> Ref a
{-# INLINE unsafeWriteRef #-}
unsafeWriteRef :: forall a. Ref a %1 -> a %1 -> Ref a
unsafeWriteRef (Ref Ref# a
v) a
a = Ref# a -> Ref a
forall a. Ref# a -> Ref a
Ref (Ref# a %1 -> a %1 -> Ref# a
forall a. Ref# a %1 -> a %1 -> Ref# a
unsafeWriteRef# Ref# a
v a
a)

instance
  (Unsatisfiable (ShowType (Ref a) :<>: Text " cannot be copied!")) =>
  Copyable (Ref a)
  where
  copy :: forall (bk :: BorrowKind) (α :: Lifetime).
Borrow bk α (Ref a) %1 -> Ref a
copy = Borrow bk α (Ref a) %1 -> Ref a
forall (msg :: ErrorMessage) a. Unsatisfiable msg => a
unsatisfiable

instance (Dupable a) => Clone (Ref a) where
  clone :: forall (α :: Lifetime). Share α (Ref a) %1 -> BO α (Ref a)
clone = (Share α (Ref a) %1 -> BO α (Ref a))
%1 -> Share α (Ref a) %1 -> BO α (Ref a)
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear \(UnsafeAlias Ref a
ref) -> Control.do
    !a <- a %1 -> BO α a
forall a. a %1 -> BO α a
forall (f :: * -> *) a. Applicative f => a %1 -> f a
Control.pure (a %1 -> BO α a) -> a %1 -> BO α a
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
PL.$ Ref a %1 -> a
forall a. Ref a %1 -> a
free Ref a
ref
    !a' <- Unsafe.toLinear (\(!a
_, !a
a') -> a %1 -> BO α a
forall a. a %1 -> BO α a
forall (f :: * -> *) a. Applicative f => a %1 -> f a
Control.pure a
a') PL.$ PL.dup a
    new a' Control.<$> askLinearly
  {-# INLINE clone #-}