Weak graph is working

main.hs

@@ -23,6 +23,10 @@ import Data.Foldable
 main :: IO ()
 main = do
   testDualWeak
+  testWeakChain
+  testGraphChain
+  testGraphX
+  testGraphO
 
 -- Demonstrate that a Weak's key keeps its value alive, even if the Weak is dead
 testWeakChain :: IO ()
@@ -30,9 +34,9 @@ testWeakChain = do
   a <- newIORef ()
   b <- newIORef ()
   _ <- mkWeakWithIORefKey a b
-  w <- mkWeakWithIORefKey b "Still Alive"
+  w <- mkWeakWithIORefKey b "b is still alive"
   performGCUntilFinalizersQuiesce
-  print =<< deRefWeak w
+  Just "b is still alive" <- deRefWeak w
   touch a
 
 testGraphChain :: IO ()
@@ -40,18 +44,18 @@ testGraphChain = do
   do
     (nai, neo) <- buildGraphChain
     performGCUntilFinalizersQuiesce
-    print =<< findForward nai
-    print =<< findBackward neo
+    ["a", "b", "c", "d", "e"] <- toList <$> findForwardVals nai
+    ["a", "b", "c", "d", "e"] <- toList <$> findBackwardVals neo
     touch (nai, neo)
   do
     (nai, _) <- buildGraphChain
     performGCUntilFinalizersQuiesce
-    print =<< findForward nai
+    [] <- toList <$> findForwardVals nai
     touch nai
   do
     (_, neo) <- buildGraphChain
     performGCUntilFinalizersQuiesce
-    print =<< findBackward neo
+    [] <- toList <$> findBackwardVals neo
     touch neo
 
 buildGraphChain :: IO (NodeInput String, NodeOutput String)
@@ -67,8 +71,8 @@ buildGraphChain = do
   link ndo nei
   pure (nai, neo)
 
-testGraphX :: IO ()
-testGraphX = do
+buildGraphX :: IO (NodeInput String, NodeInput String, NodeOutput String, NodeOutput String)
+buildGraphX = do
   (nai, nao) <- newNode "a"
   (nbi, nbo) <- newNode "b"
   (nci, nco) <- newNode "c"
@@ -78,50 +82,72 @@ testGraphX = do
   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
+  pure (nai, nbi, ndo, neo)
+
+testGraphX :: IO ()
+testGraphX = do
+  do
+    (nai, nbi, ndo, neo) <- buildGraphX
+    performGCUntilFinalizersQuiesce
+    ["a", "c", "d", "e"] <- toList <$> findForwardVals nai
+    ["b", "c", "d", "e"] <- toList <$> findForwardVals nbi
+    ["a", "b", "c", "d"] <- toList <$> findBackwardVals ndo
+    ["a", "b", "c", "e"] <- toList <$> findBackwardVals neo
+    touch (nai, nbi, ndo, neo)
+  do
+    (nai, nbi, ndo, neo) <- buildGraphX
+    performGCUntilFinalizersQuiesce
+    ["a", "c", "d"] <- toList <$> findForwardVals nai
+    ["b", "c", "d"] <- toList <$> findForwardVals nbi
+    touch (nai, nbi, ndo)
+  do
+    (nai, nbi, ndo, neo) <- buildGraphX
+    performGCUntilFinalizersQuiesce
+    [] <- toList <$> findForwardVals nai
+    [] <- toList <$> findForwardVals nbi
+    touch (nai, nbi)
 
 testGraphO :: IO ()
 testGraphO = do
+  do
+    (nai, nao) <- buildGraphO
+    performGCUntilFinalizersQuiesce
+    ["a", "b", "c"] <- toList <$> findBackward nao
+    ["a", "b", "c"] <- toList <$> findForward nai
+    touch (nai, nao)
+  do
+    (_, nao) <- buildGraphO
+    performGCUntilFinalizersQuiesce
+    [] <- toList <$> findBackward nao
+    touch nao
+  do
+    (nai, _) <- buildGraphO
+    performGCUntilFinalizersQuiesce
+    [] <- toList <$> findForward nai
+    touch nai
+
+buildGraphO :: IO (NodeInput String, NodeOutput String)
+buildGraphO = do
   (nai, nao) <- newNode "a"
   (nbi, nbo) <- newNode "b"
   (nci, nco) <- newNode "c"
   link nao nbi
   link nbo nci
   link nco nai
-  performGCUntilFinalizersQuiesce
-  [] <- toList <$> findForward nai
-  pure ()
-
-data NodeContents a = NodeContents
-  { _nodeContents_key :: !Int
-  , _nodeContents_value :: a
-  , _nodeContents_forwardLinks :: !(IORef (IntMap (NodeInput a)))
-  , _nodeContents_backLinks :: !(IORef (IntMap (NodeOutput a)))
-  }
+  pure (nai, nao)
 
 data NodeInput a = NodeInput
   { _nodeInput_key :: !Int
-  , _nodeInput_contents :: !(DualWeak (NodeContents a))
-  , _nodeInput_ticket :: !Ticket
+  , _nodeInput_contents :: !(Weak (IORef (Maybe a)))
+  , _nodeInput_forwardLinks :: !(IORef (IntMap (NodeInput a)))
+  , _nodeInput_backLinks :: !(Weak (IORef (IntMap (NodeOutput a))))
   }
 
 data NodeOutput a = NodeOutput
   { _nodeOutput_key :: !Int
-  , _nodeOutput_contents :: !(DualWeak (NodeContents a))
---  , _nodeOutput_backLinks :: !(IORef (IntMap (NodeOutput a)))
-  , _nodeOutput_ticket :: !Ticket
+  , _nodeOutput_contents :: !(Weak (IORef (Maybe a)))
+  , _nodeOutput_forwardLinks :: !(Weak (IORef (IntMap (NodeInput a))))
+  , _nodeOutput_backLinks :: !(IORef (IntMap (NodeOutput a)))
   }
 
 {-# NOINLINE globalNodeIdRef #-}
@@ -136,35 +162,50 @@ 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
-    }
+  contentsRef <- newIORef $ Just v
+  w <- mkWeakWithIORefKey backLinks contentsRef
+  wBack <- mkWeakWithIORefKey backLinks backLinks
+  wForward <- mkWeakWithIORefKeyWithFinalizer forwardLinks forwardLinks $ do
+    writeIORef finalizerDidRunRef True
+    deRefWeak w >>= \case
+      Nothing -> pure ()
+      Just contentsRef' -> do
+        writeIORef contentsRef' Nothing
+    deRefWeak wBack >>= \case
+      Nothing -> pure ()
+      Just backLinks' -> do
+        writeIORef backLinks' mempty
+        --TODO: Clear out all the nodes forward of us
   pure
     ( NodeInput
       { _nodeInput_key = nodeId
       , _nodeInput_contents = w
-      , _nodeInput_ticket = tForward
+      , _nodeInput_backLinks = wBack
+      , _nodeInput_forwardLinks = forwardLinks
       }
     , NodeOutput
       { _nodeOutput_key = nodeId
       , _nodeOutput_contents = w
---      , _nodeOutput_backLinks = backLinks
-      , _nodeOutput_ticket = tBackward
+      , _nodeOutput_backLinks = backLinks
+      , _nodeOutput_forwardLinks = wForward
       }
     )
 
+getNodeContents :: Weak (IORef (Maybe a)) -> IO (Maybe a)
+getNodeContents w = do
+  deRefWeak w >>= \case
+    Nothing -> pure Nothing
+    Just r -> readIORef r
+
 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, ())
+  deRefWeak (_nodeInput_backLinks bIn) >>= \case
+    Nothing -> pure ()
+    Just backLinks -> deRefWeak (_nodeOutput_forwardLinks aOut) >>= \case
+      Nothing -> pure ()
+      Just forwardLinks -> do
+        atomicModifyIORef' forwardLinks $ \m -> (IntMap.insert (_nodeInput_key bIn) bIn m, ())
+        atomicModifyIORef' backLinks $ \m -> (IntMap.insert (_nodeOutput_key aOut) aOut m, ())
 
 findForwardVals :: Ord a => NodeInput a -> IO (Set a)
 findForwardVals i = Set.fromList . IntMap.elems <$> findForward i
@@ -176,11 +217,14 @@ findForward i = go mempty (IntMap.singleton (_nodeInput_key i) i)
     go found toSearch = case IntMap.minViewWithKey toSearch of
       Nothing -> pure found
       Just ((k, thisIn), restToSearch) -> do
-        getDualWeak (_nodeInput_contents thisIn :: DualWeak (NodeContents a)) >>= \case
+        getNodeContents (_nodeInput_contents thisIn) >>= \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 `IntMap.difference` found))
+            newLinks <- readIORef (_nodeInput_forwardLinks thisIn)
+            go (IntMap.insert k this found) (IntMap.union restToSearch (newLinks `IntMap.difference` found))
+
+findBackwardVals :: Ord a => NodeOutput a -> IO (Set a)
+findBackwardVals i = Set.fromList . IntMap.elems <$> findBackward i
 
 findBackward :: forall a. NodeOutput a -> IO (IntMap a)
 findBackward o = go mempty (IntMap.singleton (_nodeOutput_key o) o)
@@ -188,17 +232,16 @@ findBackward o = go mempty (IntMap.singleton (_nodeOutput_key o) o)
     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
+      Just ((k, thisOut), restToSearch) -> do
+        getNodeContents (_nodeOutput_contents thisOut) >>= \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 `IntMap.difference` found))
+            newLinks <- readIORef (_nodeOutput_backLinks thisOut)
+            go (IntMap.insert k this found) (IntMap.union restToSearch (newLinks `IntMap.difference` found))
 
 testDualWeak :: IO ()
 testDualWeak = do
   do
-    putStrLn "A"
     target <- newIORef ()
     r <- mkWeakWithIORefKey target ()
     (w, a, b) <- newDualWeak target
@@ -209,7 +252,6 @@ testDualWeak = do
     touch (a, b)
     performGCUntilFinalizersQuiesce
   do
-    putStrLn "B"
     target <- newIORef ()
     r <- mkWeakWithIORefKey target ()
     (w, a, _) <- newDualWeak target
@@ -220,7 +262,6 @@ testDualWeak = do
     touch a
     performGCUntilFinalizersQuiesce
   do
-    putStrLn "C"
     target <- newIORef ()
     r <- mkWeakWithIORefKey target ()
     (w, _, b) <- newDualWeak target