2 files changed, 67 insertions, 2 deletions
diff --git a/src/Storage/Internal.hs b/src/Storage/Internal.hs
index 5f5055a..c70e8ae 100644
--- a/src/Storage/Internal.hs
+++ b/src/Storage/Internal.hs
@@ -2,28 +2,37 @@ module Storage.Internal where
 
 import Codec.Compression.Zlib
 
+import Control.Arrow
 import Control.Concurrent
+import Control.DeepSeq
 import Control.Exception
 import Control.Monad
 import Control.Monad.Identity
 
 import Crypto.Hash
 
+import Data.Bits
 import Data.ByteArray (ByteArrayAccess, ScrubbedBytes)
 import qualified Data.ByteArray as BA
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as B
 import qualified Data.ByteString.Char8 as BC
 import qualified Data.ByteString.Lazy as BL
+import Data.Function
+import Data.Hashable
+import qualified Data.HashTable.IO as HT
 import Data.List
 import Data.Map (Map)
 import qualified Data.Map as M
 
+import Foreign.Storable (peek)
+
 import System.Directory
 import System.FilePath
 import System.INotify (INotify)
 import System.IO
 import System.IO.Error
+import System.IO.Unsafe (unsafePerformIO)
 import System.Posix.Files
 import System.Posix.IO
 import System.Posix.Types
@@ -32,8 +41,11 @@ import System.Posix.Types
 data Storage' c = Storage
     { stBacking :: StorageBacking c
     , stParent :: Maybe (Storage' Identity)
+    , stRefGeneration :: MVar (HT.BasicHashTable RefDigest Generation)
     }
-    deriving (Eq)
+
+instance Eq (Storage' c) where
+    (==) = (==) `on` (stBacking &&& stParent)
 
 instance Show (Storage' c) where
     show st@(Storage { stBacking = StorageDir { dirPath = path }}) = "dir" ++ showParentStorage st ++ ":" ++ path
@@ -55,7 +67,8 @@ data StorageBacking c
     deriving (Eq)
 
 
-type RefDigest = Digest Blake2b_256
+newtype RefDigest = RefDigest (Digest Blake2b_256)
+    deriving (Eq, Ord, NFData, ByteArrayAccess)
 
 data Ref' c = Ref (Storage' c) RefDigest
     deriving (Eq)
@@ -67,6 +80,12 @@ instance ByteArrayAccess (Ref' c) where
     length (Ref _ dgst) = BA.length dgst
     withByteArray (Ref _ dgst) = BA.withByteArray dgst
 
+instance Hashable RefDigest where
+    hashWithSalt salt ref = salt `xor` unsafePerformIO (BA.withByteArray ref peek)
+
+instance Hashable (Ref' c) where
+    hashWithSalt salt ref = salt `xor` unsafePerformIO (BA.withByteArray ref peek)
+
 refStorage :: Ref' c -> Storage' c
 refStorage (Ref st _) = st
 
@@ -82,6 +101,15 @@ showRefDigest = B.concat . map showHexByte . BA.unpack
                     | otherwise = x + 87
           showHexByte x = B.pack [ showHex (x `div` 16), showHex (x `mod` 16) ]
 
+refDigestFromByteString :: ByteArrayAccess ba => ba -> Maybe RefDigest
+refDigestFromByteString = fmap RefDigest . digestFromByteString
+
+hashToRefDigest :: BL.ByteString -> RefDigest
+hashToRefDigest = RefDigest . hashFinalize . hashUpdates hashInit . BL.toChunks
+
+
+newtype Generation = Generation Int
+    deriving (Eq, Show)
 
 data Head' c = Head String (Ref' c)
     deriving (Show)
diff --git a/src/Storage/Merge.hs b/src/Storage/Merge.hs
index f0eaf98..e9cb3d7 100644
--- a/src/Storage/Merge.hs
+++ b/src/Storage/Merge.hs
@@ -2,6 +2,10 @@ module Storage.Merge (
     Mergeable(..),
     merge, storeMerge,
 
+    Generation,
+    compareGeneration, generationMax,
+    storedGeneration,
+
     generations,
     ancestors,
     precedes,
@@ -10,12 +14,17 @@ module Storage.Merge (
     findProperty,
 ) where
 
+import Control.Concurrent.MVar
+
 import qualified Data.ByteString.Char8 as BC
+import qualified Data.HashTable.IO as HT
 import Data.List
 import Data.Maybe
 import Data.Set (Set)
 import qualified Data.Set as S
 
+import System.IO.Unsafe (unsafePerformIO)
+
 import Storage
 import Storage.Internal
 import Util
@@ -46,6 +55,34 @@ previous (Stored ref _) = case load ref of
     _ -> []
 
 
+nextGeneration :: [Generation] -> Generation
+nextGeneration = foldl' helper (Generation 0)
+    where helper (Generation c) (Generation n) | c <= n    = Generation (n + 1)
+                                               | otherwise = Generation c
+
+compareGeneration :: Generation -> Generation -> Maybe Ordering
+compareGeneration (Generation x) (Generation y) = Just $ compare x y
+
+generationMax :: Storable a => [Stored a] -> Maybe (Stored a)
+generationMax (x : xs) = Just $ snd $ foldl' helper (storedGeneration x, x) xs
+    where helper (mg, mx) y = let yg = storedGeneration y
+                               in case compareGeneration mg yg of
+                                       Just LT -> (yg, y)
+                                       _       -> (mg, mx)
+generationMax [] = Nothing
+
+storedGeneration :: Storable a => Stored a -> Generation
+storedGeneration x =
+    unsafePerformIO $ withMVar (stRefGeneration $ refStorage $ storedRef x) $ \ht -> do
+        let doLookup y = HT.lookup ht (refDigest $ storedRef y) >>= \case
+                Just gen -> return gen
+                Nothing -> do
+                    gen <- nextGeneration <$> mapM doLookup (previous y)
+                    HT.insert ht (refDigest $ storedRef y) gen
+                    return gen
+        doLookup x
+
+
 generations :: Storable a => [Stored a] -> [Set (Stored a)]
 generations = unfoldr gen . (,S.empty)
     where gen (hs, cur) = case filter (`S.notMember` cur) $ previous =<< hs of