1 files changed, 25 insertions, 4 deletions

diff --git a/lib/KeyDB.hs b/lib/KeyDB.hs
index 0bc0fb3..c92f614 100644
--- a/lib/KeyDB.hs
+++ b/lib/KeyDB.hs
@@ -5,11 +5,13 @@ module KeyDB
     , SubKey(..)
     , KeyData(..)
     , KeyDB
+    , KeyGrip(..)
     , emptyKeyDB
     , keyData
     , kkData
     , lookupKeyData
     , lookupByGrip
+    , associatedKeys
     , fingerprintGrip
     , smallprGrip
     , transmute
@@ -26,6 +28,7 @@ module KeyDB
     , flattenKeys
     , flattenFiltered
     , UidString(..)
+    , buildGripMap
     ) where
 
 import Control.Monad
@@ -79,6 +82,17 @@ data KeyDB = KeyDB
     , byGrip   :: IMap KeyGrip [KeyKey]
     } deriving Show
 
+
+-- | TODO: This is an optimization to legacy (pre-KeyDB) code.  Ultimately it
+-- should be unneccessary.
+buildGripMap :: [Packet] -> IMap KeyGrip [Packet]
+buildGripMap ps = foldr go I.empty ps
+ where
+    go pkt m =  I.alter (\case Just ks -> Just (pkt:ks)
+                               Nothing -> Just [pkt])
+                        (fingerprintGrip . fingerprint $ pkt)
+                        m
+
 emptyKeyDB :: KeyDB
 emptyKeyDB = KeyDB { byKeyKey = Map.empty, byGrip = I.empty }
 
@@ -110,16 +124,23 @@ transmute perform update db = do
                            -- Note: We currently leave deleted-keys in the byGrip map.
                            , concatMap snd $ Map.elems bkk )
 
+mpGrip :: MappedPacket -> KeyGrip
+mpGrip mp = fingerprintGrip $ fingerprint $ packet mp
+
+associatedKeys :: KeyData -> [MappedPacket]
+associatedKeys kd = keyMappedPacket kd : [ k | SubKey k _ <- Map.elems (keySubKeys kd) ]
+
 alterKeyDB :: (Maybe KeyData -> Maybe KeyData) -> KeyKey -> KeyDB -> KeyDB
 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)
+            Just kd -> let go g m = I.alter (\case Nothing  -> Just [kk]
+                                                   Just kks -> Just $ mergeL [kk] kks)
+                                            g
+                                            m
+                        in foldr go (byGrip db) $ map mpGrip $ associatedKeys kd
             Nothing -> byGrip db
     }