Generation number for stored objects

5 files changed, 84 insertions, 13 deletions

diff --git a/src/Attach.hs b/src/Attach.hs
index 2ce6110..761da0f 100644
--- a/src/Attach.hs
+++ b/src/Attach.hs
@@ -6,14 +6,12 @@ module Attach (
 import Control.Monad.Except
 import Control.Monad.Reader
 
-import Crypto.Hash
 import Crypto.Random
 
 import Data.Bits
 import Data.ByteArray (Bytes, ScrubbedBytes, convert)
 import qualified Data.ByteArray as BA
 import qualified Data.ByteString.Char8 as BC
-import qualified Data.ByteString.Lazy as BL
 import Data.Maybe
 import qualified Data.Text as T
 import Data.Word
@@ -60,7 +58,7 @@ instance Storable AttachService where
         skeys <- loadBinaries "skey"
         (decline :: Maybe T.Text) <- loadMbText "decline"
         let res = catMaybes
-                [ AttachRequest <$> (digestFromByteString =<< req)
+                [ AttachRequest <$> (refDigestFromByteString =<< req)
                 , AttachResponse <$> rsp
                 , AttachRequestNonce <$> rnonce
                 , AttachIdentity <$> aid <*> pure skeys
@@ -183,8 +181,7 @@ attachAccept printMsg self peer = do
 
 
 nonceDigest :: UnifiedIdentity -> UnifiedIdentity -> Bytes -> Bytes -> RefDigest
-nonceDigest id1 id2 nonce1 nonce2 = hashFinalize $ hashUpdates hashInit $
-    BL.toChunks $ serializeObject $ Rec
+nonceDigest id1 id2 nonce1 nonce2 = hashToRefDigest $ serializeObject $ Rec
         [ (BC.pack "id", RecRef $ storedRef $ idData id1)
         , (BC.pack "id", RecRef $ storedRef $ idData id2)
         , (BC.pack "nonce", RecBinary $ convert nonce1)
diff --git a/src/Main.hs b/src/Main.hs
index b692357..0398233 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -32,6 +32,7 @@ import PubKey
 import Service
 import State
 import Storage
+import Storage.Merge
 import Sync
 
 main :: IO ()
@@ -69,6 +70,10 @@ main = do
                         Nothing -> putStrLn $ "Identity verification failed"
                     _ -> error $ "unknown object type '" ++ objtype ++ "'"
 
+        ["show-generation", sref] -> readRef st (BC.pack sref) >>= \case
+            Nothing -> error "ref does not exist"
+            Just ref -> print $ storedGeneration (wrappedLoad ref :: Stored Object)
+
         ["update-identity"] -> updateSharedIdentity st
 
         ("update-identity" : srefs) -> do
diff --git a/src/Storage.hs b/src/Storage.hs
index 5af34b7..5a5d992 100644
--- a/src/Storage.hs
+++ b/src/Storage.hs
@@ -6,6 +6,7 @@ module Storage (
     Ref, PartialRef, RefDigest,
     refStorage, refDigest,
     readRef, showRef, showRefDigest,
+    refDigestFromByteString, hashToRefDigest,
     copyRef, partialRef, partialRefFromDigest,
 
     Object, PartialObject, Object'(..), RecItem, RecItem'(..),
@@ -72,6 +73,7 @@ import qualified Data.ByteString.Lazy as BL
 import qualified Data.ByteString.Lazy.Char8 as BLC
 import Data.Char
 import Data.Function
+import qualified Data.HashTable.IO as HT
 import Data.List
 import qualified Data.Map as M
 import Data.Maybe
@@ -105,12 +107,14 @@ openStorage path = do
     createDirectoryIfMissing True $ path ++ "/objects"
     createDirectoryIfMissing True $ path ++ "/heads"
     watchers <- newMVar (Nothing, [])
-    return $ Storage { stBacking = StorageDir path watchers, stParent = Nothing }
+    refgen <- newMVar =<< HT.new
+    return $ Storage { stBacking = StorageDir path watchers, stParent = Nothing, stRefGeneration = refgen }
 
 memoryStorage' :: IO (Storage' c')
 memoryStorage' = do
     backing <- StorageMemory <$> newMVar [] <*> newMVar M.empty <*> newMVar M.empty <*> newMVar []
-    return $ Storage { stBacking = backing, stParent = Nothing }
+    refgen <- newMVar =<< HT.new
+    return $ Storage { stBacking = backing, stParent = Nothing, stRefGeneration = refgen }
 
 memoryStorage :: IO Storage
 memoryStorage = memoryStorage'
@@ -129,7 +133,7 @@ type Ref = Ref' Complete
 type PartialRef = Ref' Partial
 
 zeroRef :: Storage' c -> Ref' c
-zeroRef s = Ref s h
+zeroRef s = Ref s (RefDigest h)
     where h = case digestFromByteString $ B.replicate (hashDigestSize $ digestAlgo h) 0 of
                    Nothing -> error $ "Failed to create zero hash"
                    Just h' -> h'
@@ -141,7 +145,7 @@ isZeroRef (Ref _ h) = all (==0) $ BA.unpack h
 
 
 readRefDigest :: ByteString -> Maybe RefDigest
-readRefDigest = digestFromByteString . B.concat <=< readHex
+readRefDigest = refDigestFromByteString . B.concat <=< readHex
     where readHex bs | B.null bs = Just []
           readHex bs = do (bx, bs') <- B.uncons bs
                           (by, bs'') <- B.uncons bs'
@@ -239,7 +243,7 @@ storeRawBytes = unsafeStoreRawBytes
 
 unsafeStoreRawBytes :: Storage' c -> BL.ByteString -> IO (Ref' c)
 unsafeStoreRawBytes st raw = do
-    let dgst = hashFinalize $ hashUpdates hashInit $ BL.toChunks raw
+    let dgst = hashToRefDigest raw
     case stBacking st of
          StorageDir { dirPath = sdir } -> writeFileOnce (refPath sdir dgst) $ compress raw
          StorageMemory { memObjs = tobjs } ->
@@ -269,7 +273,7 @@ ioLoadObject ref@(Ref st rhash) = do
     file' <- ioLoadBytes ref
     return $ do
         file <- file'
-        let chash = hashFinalize $ hashUpdates hashInit $ BL.toChunks file
+        let chash = hashToRefDigest file
         when (chash /= rhash) $ error $ "Hash mismatch on object " ++ BC.unpack (showRef ref) {- TODO throw -}
         return $ case runExcept $ unsafeDeserializeObject st file of
                       Left err -> error $ err ++ ", ref " ++ BC.unpack (showRef ref) {- TODO throw -}
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