module Unison.Codebase.SqliteCodebase.Branch.Cache where

import Data.Map qualified as Map
import System.Mem.Weak
import U.Codebase.HashTags qualified as V2
import Unison.Codebase.Branch qualified as V1.Branch
import Unison.Prelude
import Unison.Sqlite qualified as Sqlite
import UnliftIO.STM

-- | A cache of 'V1.Branch.Branch' by 'V2.CausalHash'es.
data BranchCache m = BranchCache
  { forall (m :: * -> *).
BranchCache m -> CausalHash -> m (Maybe (Branch m))
lookupCachedBranch :: V2.CausalHash -> m (Maybe (V1.Branch.Branch m)),
    forall (m :: * -> *).
BranchCache m -> CausalHash -> Branch m -> m ()
insertCachedBranch :: V2.CausalHash -> V1.Branch.Branch m -> m ()
  }

-- | Creates a 'BranchCache' which uses weak references to only keep branches in the cache for
-- as long as they're reachable by something else in the app.
--
-- This means you don't need to worry about a Branch not being GC'd because it's in the cache.
newBranchCache :: forall m. (MonadIO m) => m (BranchCache Sqlite.Transaction)
newBranchCache :: forall (m :: * -> *). MonadIO m => m (BranchCache Transaction)
newBranchCache = do
  TVar (Map CausalHash (Weak (Branch Transaction)))
var <- Map CausalHash (Weak (Branch Transaction))
-> m (TVar (Map CausalHash (Weak (Branch Transaction))))
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Map CausalHash (Weak (Branch Transaction))
forall a. Monoid a => a
mempty
  pure $
    BranchCache
      { $sel:lookupCachedBranch:BranchCache :: CausalHash -> Transaction (Maybe (Branch Transaction))
lookupCachedBranch = TVar (Map CausalHash (Weak (Branch Transaction)))
-> CausalHash -> Transaction (Maybe (Branch Transaction))
lookupCachedBranch' TVar (Map CausalHash (Weak (Branch Transaction)))
var,
        $sel:insertCachedBranch:BranchCache :: CausalHash -> Branch Transaction -> Transaction ()
insertCachedBranch = TVar (Map CausalHash (Weak (Branch Transaction)))
-> CausalHash -> Branch Transaction -> Transaction ()
insertCachedBranch' TVar (Map CausalHash (Weak (Branch Transaction)))
var
      }
  where
    lookupCachedBranch' :: TVar (Map V2.CausalHash (Weak (V1.Branch.Branch Sqlite.Transaction))) -> V2.CausalHash -> Sqlite.Transaction (Maybe (V1.Branch.Branch Sqlite.Transaction))
    lookupCachedBranch' :: TVar (Map CausalHash (Weak (Branch Transaction)))
-> CausalHash -> Transaction (Maybe (Branch Transaction))
lookupCachedBranch' TVar (Map CausalHash (Weak (Branch Transaction)))
var CausalHash
ch = IO (Maybe (Branch Transaction))
-> Transaction (Maybe (Branch Transaction))
forall a. HasCallStack => IO a -> Transaction a
Sqlite.unsafeIO do
      Map CausalHash (Weak (Branch Transaction))
cache <- TVar (Map CausalHash (Weak (Branch Transaction)))
-> IO (Map CausalHash (Weak (Branch Transaction)))
forall (m :: * -> *) a. MonadIO m => TVar a -> m a
readTVarIO TVar (Map CausalHash (Weak (Branch Transaction)))
var
      case CausalHash
-> Map CausalHash (Weak (Branch Transaction))
-> Maybe (Weak (Branch Transaction))
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup CausalHash
ch Map CausalHash (Weak (Branch Transaction))
cache of
        Maybe (Weak (Branch Transaction))
Nothing -> Maybe (Branch Transaction) -> IO (Maybe (Branch Transaction))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Branch Transaction)
forall a. Maybe a
Nothing
        Just Weak (Branch Transaction)
weakRef -> Weak (Branch Transaction) -> IO (Maybe (Branch Transaction))
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak (Branch Transaction)
weakRef

    insertCachedBranch' :: TVar (Map V2.CausalHash (Weak (V1.Branch.Branch Sqlite.Transaction))) -> V2.CausalHash -> (V1.Branch.Branch Sqlite.Transaction) -> Sqlite.Transaction ()
    insertCachedBranch' :: TVar (Map CausalHash (Weak (Branch Transaction)))
-> CausalHash -> Branch Transaction -> Transaction ()
insertCachedBranch' TVar (Map CausalHash (Weak (Branch Transaction)))
var CausalHash
ch Branch Transaction
b = IO () -> Transaction ()
forall a. HasCallStack => IO a -> Transaction a
Sqlite.unsafeIO do
      -- It's worth reading the semantics of these operations.
      -- We may in the future wish to instead keep the branch object alive for as long as the
      -- CausalHash is alive, this is easy to do with 'mkWeak', but we'll start with only
      -- keeping the branch alive as long as it's directly referenced.
      Weak (Branch Transaction)
wk <- Branch Transaction
-> Maybe (IO ()) -> IO (Weak (Branch Transaction))
forall k. k -> Maybe (IO ()) -> IO (Weak k)
mkWeakPtr Branch Transaction
b (IO () -> Maybe (IO ())
forall a. a -> Maybe a
Just (IO () -> Maybe (IO ())) -> IO () -> Maybe (IO ())
forall a b. (a -> b) -> a -> b
$ TVar (Map CausalHash (Weak (Branch Transaction)))
-> CausalHash -> IO ()
removeDeadVal TVar (Map CausalHash (Weak (Branch Transaction)))
var CausalHash
ch)
      STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Map CausalHash (Weak (Branch Transaction)))
-> (Map CausalHash (Weak (Branch Transaction))
    -> Map CausalHash (Weak (Branch Transaction)))
-> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map CausalHash (Weak (Branch Transaction)))
var (CausalHash
-> Weak (Branch Transaction)
-> Map CausalHash (Weak (Branch Transaction))
-> Map CausalHash (Weak (Branch Transaction))
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert CausalHash
ch Weak (Branch Transaction)
wk)

    -- Use this as a finalizer to remove the key from the map when its value gets GC'd
    removeDeadVal :: TVar (Map V2.CausalHash (Weak (V1.Branch.Branch Sqlite.Transaction))) -> V2.CausalHash -> IO ()
    removeDeadVal :: TVar (Map CausalHash (Weak (Branch Transaction)))
-> CausalHash -> IO ()
removeDeadVal TVar (Map CausalHash (Weak (Branch Transaction)))
var CausalHash
ch = IO () -> IO ()
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
      STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Map CausalHash (Weak (Branch Transaction)))
-> (Map CausalHash (Weak (Branch Transaction))
    -> Map CausalHash (Weak (Branch Transaction)))
-> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map CausalHash (Weak (Branch Transaction)))
var (CausalHash
-> Map CausalHash (Weak (Branch Transaction))
-> Map CausalHash (Weak (Branch Transaction))
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete CausalHash
ch)