{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE ViewPatterns #-}

module ErsatzPointer
  ( -- * Ersatz pointer
    (:=>) ((:=>)),
    establish,
    establish_,
    onDismantle,

    -- * Ersatz pointer reference
    (:=>?),
    dereference,
    dismantle,

    -- * Source
    Source (..),
    PrimitiveIdentity,
  )
where

import Data.Functor
import GHC.Conc
import GHC.Exts
import GHC.IO
import GHC.IORef
import GHC.MVar
import GHC.STRef
import GHC.Weak

-- | An __ersatz pointer__.
data a :=> b = forall (a# :: TYPE 'UnliftedRep).
  ErsatzPointer
  { (a :=> b) -> a
source :: a,
    ()
sourceIdentity# :: a#,
    (a :=> b) -> b
target :: b,
    (a :=> b) -> Maybe (IO ())
maybeFinalizer :: !(Maybe (IO ()))
  }

{-# COMPLETE (:=>) #-}

pattern (:=>) :: Source a => a -> b -> (a :=> b)
pattern source $b:=> :: a -> b -> a :=> b
$m:=> :: forall r a b.
Source a =>
(a :=> b) -> (a -> b -> r) -> (Void# -> r) -> r
:=> target <-
  ErsatzPointer {source, target}
  where
    a
source :=> b
target =
      case a -> PrimitiveIdentity
forall a. Source a => a -> PrimitiveIdentity
primitiveIdentity a
source of
        PrimitiveIdentity# a
sourceIdentity# ->
          ErsatzPointer :: forall a b (a# :: TYPE 'UnliftedRep).
a -> a# -> b -> Maybe (IO ()) -> a :=> b
ErsatzPointer
            { a
source :: a
source :: a
source,
              a
sourceIdentity# :: a
sourceIdentity# :: a
sourceIdentity#,
              b
target :: b
target :: b
target,
              maybeFinalizer :: Maybe (IO ())
maybeFinalizer = Maybe (IO ())
forall a. Maybe a
Nothing
            }

-- | /Establish/ an __ersatz pointer__ @__p__@ from object @__a__@ to object @__b__@.
--
-- When this function is called,
--
-- * A relationship is established between @__a__@ and @__b__@ such that if @__a__@ is still alive, then @__b__@ is too,
--   exactly as if @__a__@ contained an actual pointer to @__b__@, as in (for example) the relationship between a record
--   and one of its fields.
--
-- * An __ersatz pointer reference__ @__r__@ is created, and can be used to determine whether @__p__@ is still
--   /established/, which will be the case until either @__a__@ is garbage-collected, or @__p__@ is /dismantled/
--   explicitly, whichever comes first.
--
-- @
--        ┌ /Memory/ ───┐
--        │ __a__       __b__ │
--        └───────────┘
--
--              ┊
--              ▼
--
-- ┌ /Code/ ────────────────────┐
-- │ __r__ \<- 'establish' (__a__ ':=>' __b__) │
-- └──────────────────────────┘
--
--              ┊
--              ▼
--
--        ┌ /Memory/ ───┐
--        │ __a__ ──__p__─➤ __b__ │
--        │     ⇡     │
--        │     __r__     │
--        └───────────┘
-- @
--
-- Note that it may be the case that
--
-- * @__a__@ already cotains an actual pointer to @__b__@.
-- * @__a__@ and @__b__@ are the same object.
--
-- In either case, /establishing/ an __ersatz pointer__ from @__a__@ to @__b__@ may still be useful, because @__r__@ can
-- then be used to determine whether @__a__@ has been garbage-collected, so long as @__r__@ is not /dismantled/
-- explicitly.
establish :: (a :=> b) -> IO (a :=>? b)
establish :: (a :=> b) -> IO (a :=>? b)
establish pointer :: a :=> b
pointer@ErsatzPointer {a#
sourceIdentity# :: a#
sourceIdentity# :: ()
sourceIdentity#, Maybe (IO ())
maybeFinalizer :: Maybe (IO ())
maybeFinalizer :: forall a b. (a :=> b) -> Maybe (IO ())
maybeFinalizer} =
  IO (Weak (a :=> b)) -> IO (a :=>? b)
coerce (a# -> (a :=> b) -> Maybe (IO ()) -> IO (Weak (a :=> b))
forall (source# :: TYPE 'UnliftedRep) target.
source# -> target -> Maybe (IO ()) -> IO (Weak target)
makeWeakPointer a#
sourceIdentity# a :=> b
pointer Maybe (IO ())
maybeFinalizer)

-- | Like 'establish', but does not return the __ersatz pointer reference__ @__r__@.
--
-- This is not useful if either
--
-- * @__a__@ already cotains an actual pointer to @__b__@.
-- * @__a__@ and @__b__@ are the same object.
establish_ :: (a :=> b) -> IO ()
establish_ :: (a :=> b) -> IO ()
establish_ =
  IO (a :=>? b) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (a :=>? b) -> IO ())
-> ((a :=> b) -> IO (a :=>? b)) -> (a :=> b) -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a :=> b) -> IO (a :=>? b)
forall a b. (a :=> b) -> IO (a :=>? b)
establish

-- | Schedule an @IO@ action to be run when @__p__@ is /dismantled/, which is either when @__a__@ is garbage-collected,
-- or when @__p__@ is /dismantled/ explicitly, whichever comes first.
onDismantle :: (a :=> b) -> IO () -> (a :=> b)
onDismantle :: (a :=> b) -> IO () -> a :=> b
onDismantle a :=> b
pointer IO ()
finalizer =
  a :=> b
pointer {maybeFinalizer :: Maybe (IO ())
maybeFinalizer = (a :=> b) -> Maybe (IO ())
forall a b. (a :=> b) -> Maybe (IO ())
maybeFinalizer a :=> b
pointer Maybe (IO ()) -> Maybe (IO ()) -> Maybe (IO ())
forall a. Semigroup a => a -> a -> a
<> IO () -> Maybe (IO ())
forall a. a -> Maybe a
Just IO ()
finalizer}

-- | An __ersatz pointer reference__ is a reference to an __ersatz pointer__, and is evidence that the pointer was
-- /established/ at some point.
newtype a :=>? b
  = ErsatzPointerReference (Weak (a :=> b))

-- | /Dereference/ an __ersatz pointer reference__ @__r__@ to determine whether the corresponding __ersatz pointer__
-- @__p__@ from @__a__@ to @__b__@ is still /established/.
--
-- In general, if @__a__@ and @__b__@ are different objects, there are three possible cases, only two of which are
-- distinguishable.
--
-- * @__p__@ is still /established/, because @__a__@ (and therefore @__b__@) are still alive.
--
-- @
-- ┌ /Memory/ ───┐
-- │ __a__ ──__p__─➤ __b__ │
-- │     ⇡     │
-- │     __r__     │
-- └───────────┘
-- @
--
-- * @__p__@ was /dismantled/ because @__a__@ was garbage-collected; it is unknown whether @__b__@ is still alive,
--   because @__b__@ may still be referred to by another object.
--
-- @
-- ┌ /Memory/ ───┐
-- │         ? │
-- │     ⇡     │
-- │     __r__     │
-- └───────────┘
-- @
--
-- * @__p__@ was /dismantled/ explicitly; it is unknown whether @__a__@ is still alive, and whether @__b__@ is still
--   alive.
--
-- @
-- ┌ /Memory/ ───┐
-- │ ?       ? │
-- │     ⇡     │
-- │     __r__     │
-- └───────────┘
-- @
dereference :: (a :=>? b) -> IO (Maybe (a :=> b))
dereference :: (a :=>? b) -> IO (Maybe (a :=> b))
dereference (ErsatzPointerReference Weak (a :=> b)
weak) =
  Weak (a :=> b) -> IO (Maybe (a :=> b))
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak (a :=> b)
weak

-- | /Dismantle/ an __ersatz pointer__ @__p__@ from @__a__@ to @__b__@, which
--
-- 1. Undoes the relationship established by 'establish', i.e. makes it no longer the case that if @__a__@ is alive,
--    @__b__@ is too.
-- 2. Causes any registered 'onDismantle' actions to be run immediately.
--
-- This action is a no-op if @__p__@ was alread /dismantled/, either because @__a__@ was already garbage-collected, or
-- because it was /dismantled/ explicitly.
--
-- @
--  ┌ /Memory/ ───┐
--  │ __a__ ──__p__─➤ __b__ │
--  │     ⇡     │
--  │     __r__     │
--  └───────────┘
--
--        ┊
--        ▼
--
-- ┌ /Code/ ───────┐
-- │ 'dismantle' __r__ │
-- └─────────────┘
--
--        ┊
--        ▼
--
--  ┌ /Memory/ ───┐
--  │ __a__       __b__ │
--  │     ⇡     │
--  │     __r__     │
--  └───────────┘
-- @
dismantle :: (a :=>? b) -> IO ()
dismantle :: (a :=>? b) -> IO ()
dismantle (ErsatzPointerReference Weak (a :=> b)
weak) =
  Weak (a :=> b) -> IO ()
forall v. Weak v -> IO ()
finalize Weak (a :=> b)
weak

------------------------------------------------------------------------------------------------------------------------
-- Source

-- | The class of types that can be used as the source of an __ersatz pointer__.
--
-- This includes types whose values have a primitive identity ('IORef', 'MVar', and 'TVar'), but may also include
-- product types that contain such a type via user-defined instances.
--
-- ==== __Example user-defined instance__
--
-- @
-- data MyRecord = MyRecord
--   { ...
--   , ref :: IORef ()
--   , ...
--   }
--
-- instance 'Source' MyRecord where
--   'primitiveIdentity' MyRecord{ref} = 'primitiveIdentity' ref
-- @
class Source a where
  primitiveIdentity :: a -> PrimitiveIdentity

instance Source (IORef a) where
  primitiveIdentity :: IORef a -> PrimitiveIdentity
  primitiveIdentity :: IORef a -> PrimitiveIdentity
primitiveIdentity (IORef (STRef MutVar# RealWorld a
var#)) = MutVar# RealWorld a -> PrimitiveIdentity
forall (a :: TYPE 'UnliftedRep). a -> PrimitiveIdentity
PrimitiveIdentity# MutVar# RealWorld a
var#

instance Source (MVar a) where
  primitiveIdentity :: MVar a -> PrimitiveIdentity
  primitiveIdentity :: MVar a -> PrimitiveIdentity
primitiveIdentity (MVar MVar# RealWorld a
var#) = MVar# RealWorld a -> PrimitiveIdentity
forall (a :: TYPE 'UnliftedRep). a -> PrimitiveIdentity
PrimitiveIdentity# MVar# RealWorld a
var#

instance Source (TVar a) where
  primitiveIdentity :: TVar a -> PrimitiveIdentity
  primitiveIdentity :: TVar a -> PrimitiveIdentity
primitiveIdentity (TVar TVar# RealWorld a
var#) = TVar# RealWorld a -> PrimitiveIdentity
forall (a :: TYPE 'UnliftedRep). a -> PrimitiveIdentity
PrimitiveIdentity# TVar# RealWorld a
var#

-- | The primitive identity of a value.
data PrimitiveIdentity where
  PrimitiveIdentity# :: forall (a :: TYPE 'UnliftedRep). a -> PrimitiveIdentity

-- The type that System.Mem.Weak.mkWeak *should* have - unlifted first argument. (Even that isn't good enough - we
-- really want to know this value has a primitive identity, hence the 'Source' class above).
makeWeakPointer :: forall (source# :: TYPE 'UnliftedRep) target. source# -> target -> Maybe (IO ()) -> IO (Weak target)
makeWeakPointer :: source# -> target -> Maybe (IO ()) -> IO (Weak target)
makeWeakPointer source#
source# target
target = \case
  Maybe (IO ())
Nothing ->
    (State# RealWorld -> (# State# RealWorld, Weak target #))
-> IO (Weak target)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO \State# RealWorld
s0# ->
      case source#
-> target
-> State# RealWorld
-> (# State# RealWorld, Weak# target #)
forall a b.
a -> b -> State# RealWorld -> (# State# RealWorld, Weak# b #)
mkWeakNoFinalizer# source#
source# target
target State# RealWorld
s0# of
        (# State# RealWorld
s1#, Weak# target
weak# #) -> (# State# RealWorld
s1#, Weak# target -> Weak target
forall v. Weak# v -> Weak v
Weak Weak# target
weak# #)
  Just (IO State# RealWorld -> (# State# RealWorld, () #)
finalizer#) ->
    (State# RealWorld -> (# State# RealWorld, Weak target #))
-> IO (Weak target)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO \State# RealWorld
s0# ->
      case source#
-> target
-> (State# RealWorld -> (# State# RealWorld, () #))
-> State# RealWorld
-> (# State# RealWorld, Weak# target #)
forall a b c.
a
-> b
-> (State# RealWorld -> (# State# RealWorld, c #))
-> State# RealWorld
-> (# State# RealWorld, Weak# b #)
mkWeak# source#
source# target
target State# RealWorld -> (# State# RealWorld, () #)
finalizer# State# RealWorld
s0# of
        (# State# RealWorld
s1#, Weak# target
weak# #) -> (# State# RealWorld
s1#, Weak# target -> Weak target
forall v. Weak# v -> Weak v
Weak Weak# target
weak# #)