Added byGrip map to KeyDB.

1 files changed, 39 insertions, 3 deletions

diff --git a/lib/KeyDB.hs b/lib/KeyDB.hs
index 1aef747..5e74ea4 100644
--- a/lib/KeyDB.hs
+++ b/lib/KeyDB.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE LambdaCase #-}
 module KeyDB
     ( TrustMap
     , SigAndTrust
@@ -8,6 +9,7 @@ module KeyDB
     , keyData
     , kkData
     , lookupKeyData
+    , lookupByGrip
     , transmute
     , transmuteAt
     , alterKeyDB
@@ -25,13 +27,21 @@ module KeyDB
     ) where
 
 import Control.Monad
+import Data.Binary
+import qualified Data.ByteString as S
+import qualified Data.ByteString.Lazy as L
 import Data.Functor
 import Data.List
 import qualified Data.Map.Strict as Map
 import Data.Maybe
 import Data.OpenPGP
 import Data.Ord
+import Foreign.Storable
 
+import Data.List.Merge
+import Data.OpenPGP.Util
+import qualified IntMapClass as I
+         ;import IntMapClass (IMap)
 import KeyRing.Types
 
 type TrustMap = Map.Map FilePath Packet
@@ -51,12 +61,20 @@ data KeyData = KeyData
     } deriving Show
 
 
+newtype KeyGrip = KeyInt Int
+
+fingerprintGrip :: Fingerprint -> KeyGrip
+fingerprintGrip (Fingerprint bs) =
+    case decode $ L.fromStrict $ S.drop (S.length bs - sizeOf (0::Int)) bs of
+        i -> KeyInt i
+
 data KeyDB = KeyDB
     { byKeyKey :: Map.Map KeyKey KeyData
+    , byGrip   :: IMap KeyGrip [KeyKey]
     } deriving Show
 
 emptyKeyDB :: KeyDB
-emptyKeyDB = KeyDB { byKeyKey = Map.empty }
+emptyKeyDB = KeyDB { byKeyKey = Map.empty, byGrip = I.empty }
 
 keyData :: KeyDB -> [KeyData]
 keyData db = Map.elems (byKeyKey db)
@@ -67,6 +85,11 @@ kkData db = Map.toList (byKeyKey db)
 lookupKeyData :: KeyKey -> KeyDB -> Maybe KeyData
 lookupKeyData kk db = Map.lookup kk (byKeyKey db)
 
+lookupByGrip :: KeyGrip -> KeyDB -> [KeyData]
+lookupByGrip k db = mapMaybe (`Map.lookup` byKeyKey db)
+                    $ concat . maybeToList
+                    $ I.lookup k (byGrip db)
+
 transmute :: (Monad m, Monad kiki, Traversable kiki) =>
        ((KeyData, [info]) -> opcode -> m (kiki (KeyData, [info]))) -- ^ interpreter
        -> (KeyData -> [opcode])                                    -- ^ instructions
@@ -78,10 +101,21 @@ transmute perform update db = do
                               (update kd)
     r <- sequenceA <$> mapM performAll (byKeyKey db)
     return $ r <&> \bkk -> ( db { byKeyKey = fst <$> bkk }
+                           -- Note: We currently leave deleted-keys in the byGrip map.
                            , concatMap snd $ Map.elems bkk )
 
 alterKeyDB :: (Maybe KeyData -> Maybe KeyData) -> KeyKey -> KeyDB -> KeyDB
-alterKeyDB update kk db = db { byKeyKey = Map.alter update kk (byKeyKey db) }
+alterKeyDB update kk db = db
+    { byKeyKey = Map.alter update kk (byKeyKey db)
+    , byGrip = case Map.lookup kk (byKeyKey db) of
+        Just _  -> byGrip db
+        Nothing -> case update Nothing of
+            Just kd -> I.alter (\case Nothing  -> Just [kk]
+                                      Just kks -> Just $ mergeL [kk] kks)
+                               (fingerprintGrip $ fingerprint $ packet $ keyMappedPacket kd)
+                               (byGrip db)
+            Nothing -> byGrip db
+    }
 
 transmuteAt :: ( Monad m
                , Functor kiki
@@ -93,7 +127,9 @@ transmuteAt go kk db = do
 
 mergeKeyDB :: (KeyData -> KeyData -> KeyData) -> KeyDB -> KeyDB -> KeyDB
 mergeKeyDB mergeKeyData db dbtrans =
-    KeyDB { byKeyKey = Map.unionWith mergeKeyData (byKeyKey db) (byKeyKey dbtrans) }
+    KeyDB { byKeyKey = Map.unionWith mergeKeyData (byKeyKey db) (byKeyKey dbtrans)
+          , byGrip   = I.unionWith mergeL (byGrip db) (byGrip dbtrans)
+          }
 
 mapKeyDB :: Monad m => (KeyData -> m KeyData) -> KeyDB -> m KeyDB
 mapKeyDB f db = fmap (\m -> db { byKeyKey = m }) $ mapM f (byKeyKey db)