Explore using weak pointers to clean up unused versions

[noughtmare]

I've been experimenting with using weak pointers to remove unused versions from the history chain. If a sequence versions somewhere in the middle of the chain is unused, their diffs could be combined into an IntMap (see also #13) which has logarithmic time lookup, so this could significantly improve the time it takes to read from old versions of the array.

I tried to start with a simple experiment: implement a list whose elements disappear when they are no longer directly referenced. I think that nicely models the behaviour we want. Unfortunately, it does not work yet and I don't have any more time this weekend. Here's my attempt:

{-# LANGUAGE LambdaCase #-}
import System.Mem.Weak
import System.Mem
import System.IO
import Control.Concurrent
import Data.IORef
import Control.Monad
import Data.Coerce

newtype WeakList a = WL (IORef (WeakListData a))
data WeakListData a = Nil | Cons a !(Weak (WeakList a))

nil :: IO (WeakList a)
nil = WL <$> newIORef Nil

cons :: Show a => a -> WeakList a -> IO (WeakList a)
cons x xs@(WL ref) = do
  ref' <- newIORef Nil
  w <- readIORef ref >>= \case
    Nil -> mkWeakIORef ref (pure ())
    Cons y ys -> mkWeakIORef ref $ print y >> writeIORef ref' (Cons x ys) 
  writeIORef ref' (Cons x (coerce w))
  pure (WL ref')

toList :: WeakList a -> IO [a]
toList (WL ref) = readIORef ref >>= \case
  Nil -> pure []
  Cons x w -> deRefWeak w >>= \case
    Nothing -> pure [x]
    Just xs -> (x :) <$> toList xs

main = do
  hSetBuffering stdout LineBuffering
  xs <- cons 3 =<< cons 2 =<< cons (1 :: Int) =<< nil
  print =<< toList xs
  ys <- cons 6 =<< cons 5 =<< cons 4 xs
  print =<< toList ys
  performGC
  threadDelay 100000
  print =<< toList xs
  print =<< toList ys

It currently prints:

[3,2,1]
[6,5,4,3,2,1]
2
1
[3]
[6,5,4,3]

I wanted it to print:

[3,2,1]
[6,5,4,3,2,1]
5
4
2
1
[3]
[6,3]

Of course this probably would have way too much overhead, but this could show that it is possible to get the GC to help us in theory.

[noughtmare]

I'm starting to think this is not possible using just weak pointers.

[noughtmare]

This does seem to work properly:

{-# LANGUAGE LambdaCase #-}

import System.Mem.Weak
import Data.IORef
import Data.Coerce
import System.Mem
import Control.Concurrent

newtype List a = List (IORef (ListData a))
data ListData a = Nil | Cons a (Weak (List a)) (IORef (ListData a))

nil :: IO (List a)
nil = List <$> newIORef Nil

cons :: a -> List a -> IO (List a)
cons x (List xs) = do
  next <- newIORef =<< readIORef xs
  self <- newIORef Nil
  w <- mkWeakIORef self (pure ())
  writeIORef self (Cons x (coerce w) next)
  pure (List self)

gc :: List a -> IO ()
gc (List r) = do
  readIORef r >>= \case
    Nil -> pure ()
    Cons _ _ r' -> go r' r'
  where
    go prev curr =
      readIORef curr >>= \case
        Nil -> writeIORef prev Nil
        x@(Cons _ w next) ->
          deRefWeak w >>= \case
            Nothing -> go prev next
            Just _ -> do
              writeIORef prev x
              go next next

toList :: List a -> IO [a]
toList (List r) =
  readIORef r >>= \case
    Nil -> pure []
    Cons x w r' -> do
      (x :) <$> toList (List r')

main = do
  xs <- cons 3 =<< cons 2 =<< cons (1 :: Int) =<< nil
  ys <- cons 6 =<< cons 5 =<< cons 4 xs
  print =<< toList xs
  print =<< toList ys
  performGC
  threadDelay 100000
  gc ys
  print =<< toList xs
  print =<< toList ys

[noughtmare]

I've made a more elaborate example which actually keeps track of the values that are evicted from the history chain:

{-# LANGUAGE LambdaCase #-}

import System.Mem.Weak
import Data.IORef
import System.Mem
import Control.Concurrent

newtype Sum = Sum (IORef (IORef SumData))
data SumData = Nil | Cons !Int (Weak (IORef (IORef SumData))) (IORef SumData)

zero :: IO Sum
zero = Sum <$> (newIORef =<< newIORef Nil)

add :: Int -> Sum -> IO Sum
add x (Sum canary) = do
  spine <- readIORef canary
  spine' <- newIORef Nil
  canary' <- newIORef spine'
  w <- mkWeakIORef canary' (pure ())
  writeIORef spine' (Cons x w spine)
  pure (Sum canary')

getSum :: Sum -> IO Int
getSum (Sum canary) = do
  spine <- readIORef canary
  go spine
  where
    go spine = readIORef spine >>= \case
      Nil -> pure 0
      Cons x _ spine' -> (x +) <$> go spine'

gc :: Sum -> IO ()
gc (Sum canary) = do
  spine <- readIORef canary
  readIORef spine >>= \case
    Nil -> pure ()
    Cons x w spine' -> go x w spine spine'
  where
    go :: Int -> Weak (IORef (IORef SumData)) -> IORef SumData -> IORef SumData -> IO ()
    go x w prev curr =
      readIORef curr >>= \case
        Nil -> writeIORef prev (Cons x w curr)
        Cons x' w' next ->
          deRefWeak w' >>= \case
            Nothing -> go (x + x') w prev next
            Just _ -> do
              writeIORef prev (Cons x w curr)
              go x' w' curr next

internalLength :: Sum -> IO Int
internalLength (Sum canary) = do
  spine <- readIORef canary
  go 0 spine
  where
    go i spine =
      readIORef spine >>= \case
        Nil -> pure i
        Cons _ _ spine' -> go (i + 1) spine'

main = do
  xs <- add 3 =<< add 2 =<< add (1 :: Int) =<< zero
  ys <- add 6 =<< add 5 =<< add 4 xs
  putStr "getSum xs: "
  print =<< getSum xs
  putStr "getSum ys: "
  print =<< getSum ys
  putStr "internalLength ys: "
  print =<< internalLength ys
  performGC
  threadDelay 100000
  -- after base 4.20 (GHC 9.10) you could use performBlockingMajorGC
  gc ys
  putStr "getSum xs: "
  print =<< getSum xs
  putStr "getSum ys: "
  print =<< getSum ys
  putStr "internalLength ys: "
  print =<< internalLength ys