Stored roots with caching

3 files changed, 33 insertions, 5 deletions

diff --git a/src/Storage.hs b/src/Storage.hs
index 84df213..cc2476f 100644
--- a/src/Storage.hs
+++ b/src/Storage.hs
@@ -115,13 +115,25 @@ openStorage path = do
     createDirectoryIfMissing True $ path ++ "/heads"
     watchers <- newMVar ([], WatchList 1 [])
     refgen <- newMVar =<< HT.new
-    return $ Storage { stBacking = StorageDir path watchers, stParent = Nothing, stRefGeneration = refgen }
+    refroots <- newMVar =<< HT.new
+    return $ Storage
+        { stBacking = StorageDir path watchers
+        , stParent = Nothing
+        , stRefGeneration = refgen
+        , stRefRoots = refroots
+        }
 
 memoryStorage' :: IO (Storage' c')
 memoryStorage' = do
     backing <- StorageMemory <$> newMVar [] <*> newMVar M.empty <*> newMVar M.empty <*> newMVar (WatchList 1 [])
     refgen <- newMVar =<< HT.new
-    return $ Storage { stBacking = backing, stParent = Nothing, stRefGeneration = refgen }
+    refroots <- newMVar =<< HT.new
+    return $ Storage
+        { stBacking = backing
+        , stParent = Nothing
+        , stRefGeneration = refgen
+        , stRefRoots = refroots
+        }
 
 memoryStorage :: IO Storage
 memoryStorage = memoryStorage'
diff --git a/src/Storage/Internal.hs b/src/Storage/Internal.hs
index 3a53567..1f093b0 100644
--- a/src/Storage/Internal.hs
+++ b/src/Storage/Internal.hs
@@ -45,6 +45,7 @@ data Storage' c = Storage
     { stBacking :: StorageBacking c
     , stParent :: Maybe (Storage' Identity)
     , stRefGeneration :: MVar (HT.BasicHashTable RefDigest Generation)
+    , stRefRoots :: MVar (HT.BasicHashTable RefDigest [RefDigest])
     }
 
 instance Eq (Storage' c) where
diff --git a/src/Storage/Merge.hs b/src/Storage/Merge.hs
index 6353dad..cedf56a 100644
--- a/src/Storage/Merge.hs
+++ b/src/Storage/Merge.hs
@@ -10,19 +10,20 @@ module Storage.Merge (
     ancestors,
     precedes,
     filterAncestors,
+    storedRoots,
 
     findProperty,
 ) where
 
 import Control.Concurrent.MVar
 
-import qualified Data.ByteString.Char8 as BC
-import qualified Data.HashTable.IO as HT
+import Data.ByteString.Char8 qualified as BC
+import Data.HashTable.IO qualified as HT
 import Data.Kind
 import Data.List
 import Data.Maybe
 import Data.Set (Set)
-import qualified Data.Set as S
+import Data.Set qualified as S
 
 import System.IO.Unsafe (unsafePerformIO)
 
@@ -107,6 +108,20 @@ filterAncestors xs = let xs' = uniq $ sort xs
                                                 in helper remains' $ uniq $ sort (px ++ filter (/=x) walk)
                                      Nothing -> remains
 
+storedRoots :: Storable a => Stored a -> [Stored a]
+storedRoots x = do
+    let st = refStorage $ storedRef x
+    unsafePerformIO $ withMVar (stRefRoots st) $ \ht -> do
+        let doLookup y = HT.lookup ht (refDigest $ storedRef y) >>= \case
+                Just roots -> return roots
+                Nothing -> do
+                    roots <- case previous y of
+                        [] -> return [refDigest $ storedRef y]
+                        ps -> map (refDigest . storedRef) . filterAncestors . map (wrappedLoad @Object . Ref st) . concat <$> mapM doLookup ps
+                    HT.insert ht (refDigest $ storedRef y) roots
+                    return roots
+        map (wrappedLoad . Ref st) <$> doLookup x
+
 findProperty :: forall a b. Storable a => (a -> Maybe b) -> [Stored a] -> [b]
 findProperty sel = map (fromJust . sel . fromStored) . filterAncestors . (findPropHeads =<<)
     where findPropHeads :: Stored a -> [Stored a]