Check component size when loading ondemand objectHEADmaster

1 files changed, 38 insertions, 0 deletions

diff --git a/src/Erebos/Object.hs b/src/Erebos/Object.hs
index f00b63d..955b0d3 100644
--- a/src/Erebos/Object.hs
+++ b/src/Erebos/Object.hs
@@ -18,6 +18,44 @@ module Erebos.Object (
     readRefDigest, showRefDigest,
     refDigestFromByteString, hashToRefDigest,
     copyRef, partialRef, partialRefFromDigest,
+
+    componentSize,
+    partialComponentSize,
 ) where
 
+import Data.ByteString.Lazy qualified as BL
+import Data.Maybe
+import Data.Set qualified as S
+import Data.Word
+
 import Erebos.Object.Internal
+
+
+componentSize :: Ref -> Word64
+componentSize ref = go S.empty [ ref ]
+  where
+    go seen (r : rs)
+        | refDigest r `S.member` seen = go seen rs
+        | otherwise = objectSize r + go (S.insert (refDigest r) seen) (referredFrom r ++ rs)
+    go _ [] = 0
+
+    objectSize = fromIntegral . BL.length . lazyLoadBytes
+    referredFrom r = case load r of
+        Rec items -> mapMaybe ((\case RecRef r' -> Just r'; _ -> Nothing) . snd) items
+        _ -> []
+
+partialComponentSize :: PartialRef -> IO Word64
+partialComponentSize ref = go S.empty [ ref ]
+  where
+    go seen (r : rs)
+        | refDigest r `S.member` seen = go seen rs
+        | otherwise = do
+            size <- objectSize r
+            referred <- referredFrom r
+            (size +) <$> go (S.insert (refDigest r) seen) (referred ++ rs)
+    go _ [] = return 0
+
+    objectSize r = either (const 0) (fromIntegral . BL.length) <$> ioLoadBytes r
+    referredFrom r = ioLoadObject r >>= return . \case
+        Right (Rec items) -> mapMaybe ((\case RecRef r' -> Just r'; _ -> Nothing) . snd) items
+        _ -> []