{-# LANGUAGE MagicHash #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE ScopedTypeVariables #-}
import System.Mem
import System.Mem.Weak
import Data.IORef
import GHC.IORef
import GHC.STRef
import GHC.IO
import GHC.Weak
import GHC.Prim
import Control.Monad.Primitive
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import System.IO.Unsafe
import Control.Monad
import Control.Concurrent

main :: IO ()
main = do
  testDualWeak

-- Demonstrate that a Weak's key keeps its value alive, even if the Weak is dead
testWeakChain :: IO ()
testWeakChain = do
  a <- newIORef ()
  b <- newIORef ()
  _ <- mkWeakWithIORefKey a b
  w <- mkWeakWithIORefKey b "Still Alive"
  performGCUntilFinalizersQuiesce
  print =<< deRefWeak w
  touch a

testGraphChain :: IO ()
testGraphChain = do
  do
    (nai, neo) <- buildGraphChain
    performGCUntilFinalizersQuiesce
    print =<< findForward nai
    print =<< findBackward neo
    touch (nai, neo)
  do
    (nai, _) <- buildGraphChain
    performGCUntilFinalizersQuiesce
    print =<< findForward nai
    touch nai
  do
    (_, neo) <- buildGraphChain
    performGCUntilFinalizersQuiesce
    print =<< findBackward neo
    touch neo

buildGraphChain :: IO (NodeInput String, NodeOutput String)
buildGraphChain = do
  (nai, nao) <- newNode "a"
  (nbi, nbo) <- newNode "b"
  (nci, nco) <- newNode "c"
  (ndi, ndo) <- newNode "d"
  (nei, neo) <- newNode "e"
  link nao nbi
  link nbo nci
  link nco ndi
  link ndo nei
  pure (nai, neo)

testGraphX :: IO ()
testGraphX = do
  (nai, nao) <- newNode "a"
  (nbi, nbo) <- newNode "b"
  (nci, nco) <- newNode "c"
  (ndi, ndo) <- newNode "d"
  (nei, neo) <- newNode "e"
  link nao nci
  link nbo nci
  link nco ndi
  link nco nei
  performGCUntilFinalizersQuiesce
  print =<< findForward nai
  print =<< findForward nbi
  print =<< findBackward ndo
  print =<< findBackward neo
  touch (nai, nbi, ndo, neo)
  performGCUntilFinalizersQuiesce
  print =<< findForward nai
  print =<< findForward nbi
  touch ndo
  performGCUntilFinalizersQuiesce
  print =<< findForward nai
  print =<< findForward nbi

data NodeContents a = NodeContents
  { _nodeContents_key :: !Int
  , _nodeContents_value :: a
  , _nodeContents_forwardLinks :: !(IORef (IntMap (NodeInput a)))
  , _nodeContents_backLinks :: !(IORef (IntMap (NodeOutput a)))
  }

data NodeInput a = NodeInput
  { _nodeInput_key :: !Int
  , _nodeInput_contents :: !(DualWeak (NodeContents a))
  , _nodeInput_ticket :: !Ticket
  }

data NodeOutput a = NodeOutput
  { _nodeOutput_key :: !Int
  , _nodeOutput_contents :: !(DualWeak (NodeContents a))
--  , _nodeOutput_backLinks :: !(IORef (IntMap (NodeOutput a)))
  , _nodeOutput_ticket :: !Ticket
  }

{-# NOINLINE globalNodeIdRef #-}
globalNodeIdRef :: IORef Int
globalNodeIdRef = unsafePerformIO $ newIORef 1

newNodeId :: IO Int
newNodeId = atomicModifyIORef' globalNodeIdRef $ \n -> (succ n, n)

newNode :: a -> IO (NodeInput a, NodeOutput a)
newNode v = do
  nodeId <- newNodeId
  backLinks <- newIORef mempty
  forwardLinks <- newIORef mempty
  (w, tBackward, tForward) <- newDualWeak $ NodeContents
    { _nodeContents_key = nodeId
    , _nodeContents_value = v
    , _nodeContents_backLinks = backLinks
    , _nodeContents_forwardLinks = forwardLinks
    }
  pure
    ( NodeInput
      { _nodeInput_key = nodeId
      , _nodeInput_contents = w
      , _nodeInput_ticket = tForward
      }
    , NodeOutput
      { _nodeOutput_key = nodeId
      , _nodeOutput_contents = w
--      , _nodeOutput_backLinks = backLinks
      , _nodeOutput_ticket = tBackward
      }
    )

link :: NodeOutput a -> NodeInput a -> IO ()
link aOut bIn = do
  getDualWeak (_nodeOutput_contents aOut) >>= \case
    Nothing -> pure () -- Linking to a dead node does nothing
    Just a -> getDualWeak (_nodeInput_contents bIn) >>= \case
      Nothing -> pure () -- Linking to a dead node does nothing
      Just b -> do
        atomicModifyIORef' (_nodeContents_forwardLinks a) $ \m -> (IntMap.insert (_nodeContents_key b) bIn m, ())
        atomicModifyIORef' (_nodeContents_backLinks b) $ \m -> (IntMap.insert (_nodeContents_key a) aOut m, ())

findForward :: forall a. NodeInput a -> IO (IntMap a)
findForward i = go mempty (IntMap.singleton (_nodeInput_key i) i)
  where
    go :: IntMap a -> IntMap (NodeInput a) -> IO (IntMap a)
    go found toSearch = case IntMap.minViewWithKey toSearch of
      Nothing -> pure found
      Just ((k, thisIn), restToSearch) -> do
        getDualWeak (_nodeInput_contents thisIn :: DualWeak (NodeContents a)) >>= \case
          Nothing -> go found restToSearch
          Just this -> do
            newLinks <- readIORef (_nodeContents_forwardLinks this)
            go (IntMap.insert k (_nodeContents_value this) found) (IntMap.union restToSearch newLinks)

findBackward :: forall a. NodeOutput a -> IO (IntMap a)
findBackward i = go mempty (IntMap.singleton (_nodeOutput_key i) i)
  where
    go :: IntMap a -> IntMap (NodeOutput a) -> IO (IntMap a)
    go found toSearch = case IntMap.minViewWithKey toSearch of
      Nothing -> pure found
      Just ((k, thisIn), restToSearch) -> do
        getDualWeak (_nodeOutput_contents thisIn :: DualWeak (NodeContents a)) >>= \case
          Nothing -> go found restToSearch
          Just this -> do
            newLinks <- readIORef (_nodeContents_backLinks this)
            go (IntMap.insert k (_nodeContents_value this) found) (IntMap.union restToSearch newLinks)

testDualWeak :: IO ()
testDualWeak = do
  do
    putStrLn "A"
    target <- newIORef ()
    r <- mkWeakWithIORefKey target ()
    (w, a, b) <- newDualWeak target
    performGC
    threadDelay 1000000
    Just () <- deRefWeak r
    Just _ <- getDualWeak w
    touch (a, b)
    performGCUntilFinalizersQuiesce
  do
    putStrLn "B"
    target <- newIORef ()
    r <- mkWeakWithIORefKey target ()
    (w, a, _) <- newDualWeak target
    performGC
    threadDelay 1000000
    Nothing <- deRefWeak r
    Nothing <- getDualWeak w
    touch a
    performGCUntilFinalizersQuiesce
  do
    putStrLn "C"
    target <- newIORef ()
    r <- mkWeakWithIORefKey target ()
    (w, _, b) <- newDualWeak target
    performGC
    threadDelay 1000000
    performGC
    Nothing <- deRefWeak r
    Nothing <- getDualWeak w
    touch b
    performGCUntilFinalizersQuiesce

performGCUntilFinalizersQuiesce :: IO ()
performGCUntilFinalizersQuiesce = do
  writeIORef finalizerDidRunRef False
  performGC
  fence <- createFinalizerFence -- Based on my shaky memory, I think finalizers are processed in a queue, and therefore this fence will run after all the other finalizers in the GC we just finished.  However, it may run at ANY point in the subsequent GC, so we can't rely on it to know that *that* GC had no finalizers run.
  performGC
  waitForFinalizerFence fence
  readIORef finalizerDidRunRef >>= \case
    True -> performGCUntilFinalizersQuiesce
    False -> pure ()

{-# NOINLINE finalizerDidRunRef #-}
finalizerDidRunRef :: IORef Bool
finalizerDidRunRef = unsafePerformIO $ newIORef False

newtype FinalizerFence = FinalizerFence (MVar ())

createFinalizerFence :: IO FinalizerFence
createFinalizerFence = do
  r <- newIORef ()
  v <- newEmptyMVar
  _ <- mkWeakWithIORefKeyWithFinalizer r () $ putMVar v ()
  pure $ FinalizerFence v

waitForFinalizerFence :: FinalizerFence -> IO ()
waitForFinalizerFence (FinalizerFence v) = takeMVar v

newtype DualWeak a = DualWeak (Weak (Weak (IORef (Maybe a))))
newtype Ticket = Ticket (IORef ())

-- | Creates a weak reference to `a` which only remains alive if *both* tickets are alive
newDualWeak :: a -> IO (DualWeak a, Ticket, Ticket)
newDualWeak v = do
  vRef <- newIORef $ Just v
  tInner <- newIORef ()
  wInner <- mkWeakWithIORefKey tInner vRef
  tOuter <- newIORef ()
  wOuter <- mkWeakWithIORefKeyWithFinalizer tOuter wInner $ do
    writeIORef finalizerDidRunRef True
    deRefWeak wInner >>= \case
      Nothing -> pure ()
      Just vRef' -> do
        writeIORef vRef' Nothing
  pure (DualWeak wOuter, Ticket tOuter, Ticket tInner)

getDualWeak :: DualWeak a -> IO (Maybe a)
getDualWeak (DualWeak wOuter) = do
  deRefWeak wOuter >>= \case
    Nothing -> pure Nothing
    Just wInner -> deRefWeak wInner >>= \case
      Nothing -> pure Nothing
      Just vRef -> readIORef vRef

mkWeakWithIORefKey :: IORef a -> b -> IO (Weak b)
mkWeakWithIORefKey (IORef (STRef r#)) v = IO $ \s ->
    case mkWeakNoFinalizer# r# v s of (# s1, w #) -> (# s1, Weak w #)

mkWeakWithIORefKeyWithFinalizer :: IORef a -> b -> IO () -> IO (Weak b)
mkWeakWithIORefKeyWithFinalizer (IORef (STRef r#)) v (IO f) = IO $ \s ->
    case mkWeak# r# v f s of (# s1, w #) -> (# s1, Weak w #)