{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE QualifiedDo #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

{- |
A reference cell. To mutate, use as @'Mut' α ('Ref' a)@.
This module is inteted to be imported qualified.
-}
module Data.Ref.Linear.Borrow (
  Ref (),
  update,
  modify,
  swap,
  readShare,
  copyRef,
) where

import Control.Functor.Linear qualified as Control
import Control.Monad.Borrow.Pure.BO
import Control.Monad.Borrow.Pure.BO.Unsafe
import Control.Monad.Borrow.Pure.Copyable
import Control.Syntax.DataFlow qualified as DataFlow
import Data.Ref.Linear (Ref)
import Data.Ref.Linear qualified as Ref
import Prelude.Linear
import Unsafe.Linear qualified as Unsafe
import Prelude qualified as NonLinear

update :: (α >= β) => (a %1 -> BO β (b, a)) %1 -> Mut α (Ref a) %1 -> BO β (b, Mut α (Ref a))
{-# INLINE update #-}
update :: forall (α :: Lifetime) (β :: Lifetime) a b.
(α >= β) =>
(a %1 -> BO β (b, a))
%1 -> Mut α (Ref a) %1 -> BO β (b, Mut α (Ref a))
update a %1 -> BO β (b, a)
f (UnsafeAlias Ref a
mv) = DataFlow.do
  -- NOTE: as there is only one reference to @'Ref' a@, we can just use read/write
  -- instead of 'MutVar.atomicModify' (which requires pure function) while retaining atomicity.
  (!a, !mv) <- Ref a %1 -> (a, Ref a)
forall a. Ref a %1 -> (a, Ref a)
Ref.unsafeReadRef Ref a
mv
  f a Control.<&> \(!b
b, !a
a) -> DataFlow.do
    !mv <- Ref a %1 -> a %1 -> Ref a
forall a. Ref a %1 -> a %1 -> Ref a
Ref.unsafeWriteRef Ref a
mv a
a
    (b, UnsafeAlias mv)

modify :: (α >= β) => (a %1 -> a) %1 -> Mut α (Ref a) %1 -> BO β (Mut α (Ref a))
modify :: forall (α :: Lifetime) (β :: Lifetime) a.
(α >= β) =>
(a %1 -> a) %1 -> Mut α (Ref a) %1 -> BO β (Mut α (Ref a))
modify a %1 -> a
f Mut α (Ref a)
ma = Control.do
  ((), ma) <- (a %1 -> BO β ((), a))
%1 -> Mut α (Ref a) %1 -> BO β ((), Mut α (Ref a))
forall (α :: Lifetime) (β :: Lifetime) a b.
(α >= β) =>
(a %1 -> BO β (b, a))
%1 -> Mut α (Ref a) %1 -> BO β (b, Mut α (Ref a))
update (((), 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 -> ((), a)) %1 -> a %1 -> BO β ((), a)
forall b c a (q :: Multiplicity) (m :: Multiplicity)
       (n :: Multiplicity).
(b %1 -> c) %q -> (a %1 -> b) %m -> a %n -> c
. ((),) (a %1 -> ((), a)) -> (a %1 -> a) %1 -> a %1 -> ((), a)
forall b c a (q :: Multiplicity) (m :: Multiplicity)
       (n :: Multiplicity).
(b %1 -> c) %q -> (a %1 -> b) %m -> a %n -> c
. a %1 -> a
f) Mut α (Ref a)
ma
  Control.pure ma

swap :: (α >= β) => Mut α (Ref a) %1 -> Mut α (Ref a) %1 -> BO β (Mut α (Ref a), Mut α (Ref a))
{-# INLINE swap #-}
swap :: forall (α :: Lifetime) (β :: Lifetime) a.
(α >= β) =>
Mut α (Ref a)
%1 -> Mut α (Ref a) %1 -> BO β (Mut α (Ref a), Mut α (Ref a))
swap Mut α (Ref a)
ma Mut α (Ref a)
ma' =
  ((a %1 -> BO β (Mut α (Ref a), a))
 %1 -> Mut α (Ref a) %1 -> BO β (Mut α (Ref a), Mut α (Ref a)))
-> Mut α (Ref a)
%1 -> (a %1 -> BO β (Mut α (Ref a), a))
%1 -> BO β (Mut α (Ref a), Mut α (Ref a))
forall a (p :: Multiplicity) b (q :: Multiplicity) c
       (r :: Multiplicity).
(a %p -> b %q -> c) %r -> b %q -> a %p -> c
flip (a %1 -> BO β (Mut α (Ref a), a))
%1 -> Mut α (Ref a) %1 -> BO β (Mut α (Ref a), Mut α (Ref a))
forall (α :: Lifetime) (β :: Lifetime) a b.
(α >= β) =>
(a %1 -> BO β (b, a))
%1 -> Mut α (Ref a) %1 -> BO β (b, Mut α (Ref a))
update Mut α (Ref a)
ma' \ !a
a' -> Control.do
    (a, ma) <- (a %1 -> BO β (a, a))
%1 -> Mut α (Ref a) %1 -> BO β (a, Mut α (Ref a))
forall (α :: Lifetime) (β :: Lifetime) a b.
(α >= β) =>
(a %1 -> BO β (b, a))
%1 -> Mut α (Ref a) %1 -> BO β (b, Mut α (Ref a))
update (\ !a
a -> (a, a) %1 -> BO β (a, a)
forall a. a %1 -> BO β a
forall (f :: * -> *) a. Applicative f => a %1 -> f a
Control.pure (a
a, a
a')) Mut α (Ref a)
ma
    Control.pure (ma, a)

readShare :: (α >= β) => Share α (Ref a) %1 -> BO β (Ur (Share α a))
{-# INLINE readShare #-}
readShare :: forall (α :: Lifetime) (β :: Lifetime) a.
(α >= β) =>
Share α (Ref a) %1 -> BO β (Ur (Share α a))
readShare = (Share α (Ref a) -> BO β (Ur (Share α a)))
%1 -> Share α (Ref a) %1 -> BO β (Ur (Share α a))
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear \(UnsafeAlias Ref a
mv) ->
  Ur (Share α a) %1 -> BO β (Ur (Share α a))
forall a. a %1 -> BO β a
forall (f :: * -> *) a. Applicative f => a %1 -> f a
Control.pure (Ur (Share α a) %1 -> BO β (Ur (Share α a)))
-> Ur (Share α a) %1 -> BO β (Ur (Share α a))
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ Share α a -> Ur (Share α a)
forall a. a -> Ur a
Ur (Share α a -> Ur (Share α a)) -> Share α a -> Ur (Share α a)
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$! a -> Share α a
forall (ak :: AliasKind) (α :: Lifetime) a. a -> Alias ak α a
UnsafeAlias (a -> Share α a) -> a -> Share α a
forall a b. (a -> b) -> a -> b
NonLinear.$! (a, Ref a) -> a
forall a b. (a, b) -> a
NonLinear.fst ((a, Ref a) -> a) -> (a, Ref a) -> a
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$! Ref a %1 -> (a, Ref a)
forall a. Ref a %1 -> (a, Ref a)
Ref.unsafeReadRef Ref a
mv

copyRef :: (Copyable a, α >= β) => Borrow k α (Ref a) %1 -> BO β a
{-# INLINE copyRef #-}
copyRef :: forall a (α :: Lifetime) (β :: Lifetime) (k :: BorrowKind).
(Copyable a, α >= β) =>
Borrow k α (Ref a) %1 -> BO β a
copyRef Borrow k α (Ref a)
bor =
  Borrow k α (Ref a) %1 -> Ur (Share α (Ref a))
forall (k :: BorrowKind) (α :: Lifetime) a.
Borrow k α a %1 -> Ur (Share α a)
share Borrow k α (Ref a)
bor Ur (Share α (Ref a))
%1 -> (Ur (Share α (Ref a)) %1 -> BO β a) -> BO β a
forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(Ur Share α (Ref a)
bor) -> Control.do
    Ur !shr <- Share α (Ref a) %1 -> BO β (Ur (Share α a))
forall (α :: Lifetime) (β :: Lifetime) a.
(α >= β) =>
Share α (Ref a) %1 -> BO β (Ur (Share α a))
readShare Share α (Ref a)
bor
    Control.pure $! copy shr