{-# LANGUAGE GHC2021 #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE MagicHash #-}
module Symbolize.WeakSymbol (WeakSymbol, new, deref) where

import Data.Array.Byte (ByteArray (ByteArray))
import GHC.Exts (ByteArray#, StableName#, Weak#, deRefWeak#, mkWeak#, makeStableName#)
import GHC.IO (IO (IO))
import Symbolize.Accursed qualified as Accursed

-- Inside the WeakSymbol
-- we keep:
-- - A weak pointer to the underlying ByteArray#.
--   This weak pointer will be invalidated (turn into a 'tombstone')
--   when the final instance of this symbol is GC'd
-- - A `StableName` for the same ByteArray#
--   This ensures we have a stable hash even in the presence of the ByteArray#
--   being moved around by the GC.
--   We never read it after construction,
--   but by keeping it around until the WeakSymbol is removed by the finalizer,
--   we ensure that future calls to `makeStableName` return the same hash.
data WeakSymbol where
  WeakSymbol# :: Weak# ByteArray# -> StableName# ByteArray# -> WeakSymbol

-- | Create a new weak symbol
-- based on the given symbol content ByteArray
-- and finalizer to run when the weak symbol
-- is no longer needed.
new :: ByteArray# -> IO () -> IO WeakSymbol
{-# INLINE new #-}
new :: ByteArray# -> IO () -> IO WeakSymbol
new ByteArray#
ba# (IO State# RealWorld -> (# State# RealWorld, () #)
finalizer#) =
    -- SAFETY: This should even be safe
    -- in the presence of inlining, CSE and full laziness
    --
    -- because the result is outwardly pure
    -- and the finalizer we use is idempotent
  (State# RealWorld -> (# State# RealWorld, WeakSymbol #))
-> IO WeakSymbol
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, WeakSymbol #))
 -> IO WeakSymbol)
-> (State# RealWorld -> (# State# RealWorld, WeakSymbol #))
-> IO WeakSymbol
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s1 -> case ByteArray#
-> ByteArray#
-> (State# RealWorld -> (# State# RealWorld, () #))
-> State# RealWorld
-> (# State# RealWorld, Weak# ByteArray# #)
forall a b c.
a
-> b
-> (State# RealWorld -> (# State# RealWorld, c #))
-> State# RealWorld
-> (# State# RealWorld, Weak# b #)
mkWeak# ByteArray#
ba# ByteArray#
ba# State# RealWorld -> (# State# RealWorld, () #)
finalizer# State# RealWorld
s1 of
    (# State# RealWorld
s2, Weak# ByteArray#
weak# #) ->
      case ByteArray#
-> State# RealWorld
-> (# State# RealWorld, StableName# ByteArray# #)
forall a.
a -> State# RealWorld -> (# State# RealWorld, StableName# a #)
makeStableName# ByteArray#
ba# State# RealWorld
s2 of
        (# State# RealWorld
s3, StableName# ByteArray#
sname# #) ->
          (# State# RealWorld
s3, Weak# ByteArray# -> StableName# ByteArray# -> WeakSymbol
WeakSymbol# Weak# ByteArray#
weak# StableName# ByteArray#
sname# #)

-- | Attempt to get back the containing ByteArray#
-- by looking inside this `WeakSymbol`
--
-- Returns `Nothing` if it was GC'd in the meantime
-- (which may be before, after or concurrently with when the finalizer runs)
deref :: WeakSymbol -> Maybe ByteArray
{-# INLINE deref #-}
deref :: WeakSymbol -> Maybe ByteArray
deref (WeakSymbol# Weak# ByteArray#
w StableName# ByteArray#
_sn) =
    -- SAFETY: This should even be safe
    -- in the presence of inlining, CSE and full laziness;
    IO (Maybe ByteArray) -> Maybe ByteArray
forall a. IO a -> a
Accursed.accursedUnutterablePerformIO (IO (Maybe ByteArray) -> Maybe ByteArray)
-> IO (Maybe ByteArray) -> Maybe ByteArray
forall a b. (a -> b) -> a -> b
$ (State# RealWorld -> (# State# RealWorld, Maybe ByteArray #))
-> IO (Maybe ByteArray)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, Maybe ByteArray #))
 -> IO (Maybe ByteArray))
-> (State# RealWorld -> (# State# RealWorld, Maybe ByteArray #))
-> IO (Maybe ByteArray)
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
  case Weak# ByteArray#
-> State# RealWorld -> (# State# RealWorld, Int#, ByteArray# #)
forall a.
Weak# a -> State# RealWorld -> (# State# RealWorld, Int#, a #)
deRefWeak# Weak# ByteArray#
w State# RealWorld
s of
    (# State# RealWorld
s1, Int#
flag, ByteArray#
p #) -> case Int#
flag of
      Int#
0# -> (# State# RealWorld
s1, Maybe ByteArray
forall a. Maybe a
Nothing #)
      Int#
_ -> (# State# RealWorld
s1, ByteArray -> Maybe ByteArray
forall a. a -> Maybe a
Just (ByteArray# -> ByteArray
ByteArray ByteArray#
p) #)