Implement sharing for hls-graph Keys

ghcide/src/Development/IDE/Core/Shake.hs

@@ -324,7 +324,7 @@ getPluginConfig plugin = do
 addPersistentRule :: IdeRule k v => k -> (NormalizedFilePath -> IdeAction (Maybe (v,PositionDelta,TextDocumentVersion))) -> Rules ()
 addPersistentRule k getVal = do
   ShakeExtras{persistentKeys} <- getShakeExtrasRules
-  void $ liftIO $ atomically $ modifyTVar' persistentKeys $ HMap.insert (Key k) (fmap (fmap (first3 toDyn)) . getVal)
+  void $ liftIO $ atomically $ modifyTVar' persistentKeys $ HMap.insert (newKey k) (fmap (fmap (first3 toDyn)) . getVal)
 
 class Typeable a => IsIdeGlobal a where
 
@@ -399,7 +399,7 @@ lastValueIO s@ShakeExtras{positionMapping,persistentKeys,state} k file = do
           pmap <- readTVarIO persistentKeys
           mv <- runMaybeT $ do
             liftIO $ Logger.logDebug (logger s) $ T.pack $ "LOOKUP PERSISTENT FOR: " ++ show k
-            f <- MaybeT $ pure $ HMap.lookup (Key k) pmap
+            f <- MaybeT $ pure $ HMap.lookup (newKey k) pmap
             (dv,del,ver) <- MaybeT $ runIdeAction "lastValueIO" s $ f file
             MaybeT $ pure $ (,del,ver) <$> fromDynamic dv
           case mv of
@@ -1068,7 +1068,7 @@ defineEarlyCutoff recorder (Rule op) = addRule $ \(Q (key, file)) (old :: Maybe
     extras <- getShakeExtras
     let diagnostics ver diags = do
             traceDiagnostics diags
-            updateFileDiagnostics recorder file ver (Key key) extras . map (\(_,y,z) -> (y,z)) $ diags
+            updateFileDiagnostics recorder file ver (newKey key) extras . map (\(_,y,z) -> (y,z)) $ diags
     defineEarlyCutoff' diagnostics (==) key file old mode $ const $ op key file
 defineEarlyCutoff recorder (RuleNoDiagnostics op) = addRule $ \(Q (key, file)) (old :: Maybe BS.ByteString) mode -> otTracedAction key file mode traceA $ \traceDiagnostics -> do
     let diagnostics _ver diags = do
@@ -1087,7 +1087,7 @@ defineEarlyCutoff recorder (RuleWithOldValue op) = addRule $ \(Q (key, file)) (o
     extras <- getShakeExtras
     let diagnostics ver diags = do
             traceDiagnostics diags
-            updateFileDiagnostics recorder file ver (Key key) extras . map (\(_,y,z) -> (y,z)) $ diags
+            updateFileDiagnostics recorder file ver (newKey key) extras . map (\(_,y,z) -> (y,z)) $ diags
     defineEarlyCutoff' diagnostics (==) key file old mode $ op key file
 
 defineNoFile :: IdeRule k v => Recorder (WithPriority Log) -> (k -> Action v) -> Rules ()

ghcide/src/Development/IDE/Types/Shake.hs

@@ -26,7 +26,7 @@ import           Data.Typeable                        (cast)
 import           Data.Vector                          (Vector)
 import           Development.IDE.Core.PositionMapping
 import           Development.IDE.Core.RuleTypes       (FileVersion)
-import           Development.IDE.Graph                (Key (..), RuleResult)
+import           Development.IDE.Graph                (Key (..), RuleResult, newKey)
 import qualified Development.IDE.Graph                as Shake
 import           Development.IDE.Types.Diagnostics
 import           Development.IDE.Types.Location
@@ -75,7 +75,7 @@ isBadDependency x
     | otherwise = False
 
 toKey :: Shake.ShakeValue k => k -> NormalizedFilePath -> Key
-toKey = (Key.) . curry Q
+toKey = (newKey.) . curry Q
 
 fromKey :: Typeable k => Key -> Maybe (k, NormalizedFilePath)
 fromKey (Key k)
@@ -91,7 +91,7 @@ fromKeyType (Key k) = case typeOf k of
     _ -> Nothing
 
 toNoFileKey :: (Show k, Typeable k, Eq k, Hashable k) => k -> Key
-toNoFileKey k = Key $ Q (k, emptyFilePath)
+toNoFileKey k = newKey $ Q (k, emptyFilePath)
 
 newtype Q k = Q (k, NormalizedFilePath)
     deriving newtype (Eq, Hashable, NFData)

hls-graph/src/Development/IDE/Graph.hs

@@ -1,8 +1,10 @@
+{-# LANGUAGE PatternSynonyms #-}
 module Development.IDE.Graph(
-    shakeOptions,
+      shakeOptions,
     Rules,
     Action, action,
-    Key(..),
+    Key(.., Key),
+    newKey,
     actionFinally, actionBracket, actionCatch, actionFork,
     -- * Configuration
     ShakeOptions(shakeAllowRedefineRules, shakeExtra),

hls-graph/src/Development/IDE/Graph/Database.hs

@@ -79,7 +79,7 @@ shakeGetBuildEdges :: ShakeDatabase -> IO Int
 shakeGetBuildEdges (ShakeDatabase _ _ db) = do
     keys <- getDatabaseValues db
     let ress = mapMaybe (getResult . snd) keys
-    return $ sum $ map (length . getResultDepsDefault [] . resultDeps) ress
+    return $ sum $ map (length . getResultDepsDefault mempty . resultDeps) ress
 
 -- | Returns an approximation of the database keys,
 --   annotated with how long ago (in # builds) they were visited

hls-graph/src/Development/IDE/Graph/Internal/Action.hs

@@ -26,6 +26,7 @@ import           Control.Monad.Trans.Class
 import           Control.Monad.Trans.Reader
 import           Data.Foldable                           (toList)
 import           Data.Functor.Identity
+import qualified Data.HashSet                         as HSet
 import           Data.IORef
 import           Development.IDE.Graph.Classes
 import           Development.IDE.Graph.Internal.Database
@@ -39,7 +40,7 @@ type ShakeValue a = (Show a, Typeable a, Eq a, Hashable a, NFData a)
 alwaysRerun :: Action ()
 alwaysRerun = do
     ref <- Action $ asks actionDeps
-    liftIO $ modifyIORef ref (AlwaysRerunDeps [] <>)
+    liftIO $ modifyIORef ref (AlwaysRerunDeps mempty <>)
 
 -- No-op for now
 reschedule :: Double -> Action ()
@@ -121,7 +122,7 @@ apply ks = do
     stack <- Action $ asks actionStack
     (is, vs) <- liftIO $ build db stack ks
     ref <- Action $ asks actionDeps
-    liftIO $ modifyIORef ref (ResultDeps (toList is) <>)
+    liftIO $ modifyIORef ref (ResultDeps (HSet.fromList $ toList is) <>)
     pure vs
 
 -- | Evaluate a list of keys without recording any dependencies.

hls-graph/src/Development/IDE/Graph/Internal/Database.hs

@@ -9,6 +9,7 @@
 {-# LANGUAGE ScopedTypeVariables        #-}
 {-# LANGUAGE TupleSections              #-}
 {-# LANGUAGE TypeFamilies               #-}
+{-# LANGUAGE ViewPatterns               #-}
 
 module Development.IDE.Graph.Internal.Database (newDatabase, incDatabase, build, getDirtySet, getKeysAndVisitAge) where
 
@@ -87,7 +88,7 @@ build
 -- build _ st k | traceShow ("build", st, k) False = undefined
 build db stack keys = do
     built <- runAIO $ do
-        built <- builder db stack (fmap Key keys)
+        built <- builder db stack (fmap newKey keys)
         case built of
           Left clean  -> return clean
           Right dirty -> liftIO dirty
@@ -145,7 +146,7 @@ refresh :: Database -> Stack -> Key -> Maybe Result -> AIO (IO Result)
 -- refresh _ st k _ | traceShow ("refresh", st, k) False = undefined
 refresh db stack key result = case (addStack key stack, result) of
     (Left e, _) -> throw e
-    (Right stack, Just me@Result{resultDeps = ResultDeps deps}) -> do
+    (Right stack, Just me@Result{resultDeps = ResultDeps (HSet.toList -> deps)}) -> do
         res <- builder db stack deps
         let isDirty = any (\(_,dep) -> resultBuilt me < resultChanged dep)
         case res of
@@ -176,7 +177,7 @@ compute db@Database{..} stack key mode result = do
         actualDeps = if runChanged /= ChangedNothing then deps else previousDeps
         previousDeps= maybe UnknownDeps resultDeps result
     let res = Result runValue built' changed built actualDeps execution runStore
-    case getResultDepsDefault [] actualDeps of
+    case getResultDepsDefault mempty actualDeps of
         deps | not(null deps)
             && runChanged /= ChangedNothing
                     -> do
@@ -186,8 +187,8 @@ compute db@Database{..} stack key mode result = do
             -- on the next build.
             void $
                 updateReverseDeps key db
-                    (getResultDepsDefault [] previousDeps)
-                    (HSet.fromList deps)
+                    (getResultDepsDefault mempty previousDeps)
+                    deps
         _ -> pure ()
     atomicallyNamed "compute" $ SMap.focus (updateStatus $ Clean res) key databaseValues
     pure res
@@ -235,14 +236,13 @@ splitIO act = do
 updateReverseDeps
     :: Key        -- ^ Id
     -> Database
-    -> [Key] -- ^ Previous direct dependencies of Id
+    -> HashSet Key -- ^ Previous direct dependencies of Id
     -> HashSet Key -- ^ Current direct dependencies of Id
     -> IO ()
 -- mask to ensure that all the reverse dependencies are updated
 updateReverseDeps myId db prev new = do
-    forM_ prev $ \d ->
-        unless (d `HSet.member` new) $
-            doOne (HSet.delete myId) d
+    forM_ (HSet.toList $ prev `HSet.difference` new) $ \d ->
+         doOne (HSet.delete myId) d
     forM_ (HSet.toList new) $
         doOne (HSet.insert myId)
     where

hls-graph/src/Development/IDE/Graph/Internal/Profile.hs

@@ -60,7 +60,7 @@ data ProfileEntry = ProfileEntry
 -- resultsOnly :: Map.HashMap Id (Key, Status) -> Map.HashMap Id (Key, Result (Either BS.ByteString Value))
 resultsOnly :: [(Key, Status)] -> Map.HashMap Key Result
 resultsOnly mp = Map.map (\r ->
-      r{resultDeps = mapResultDeps (filter (isJust . flip Map.lookup keep)) $ resultDeps r}
+      r{resultDeps = mapResultDeps (Set.filter (isJust . flip Map.lookup keep)) $ resultDeps r}
     ) keep
     where
         keep = Map.fromList $ mapMaybe (traverse getResult) mp
@@ -103,15 +103,15 @@ dependencyOrder shw status =
 prepareForDependencyOrder :: Database -> IO (HashMap Key Result)
 prepareForDependencyOrder db = do
     current <- readTVarIO $ databaseStep db
-    Map.insert (Key "alwaysRerun") (alwaysRerunResult current) .  resultsOnly
+    Map.insert (newKey "alwaysRerun") (alwaysRerunResult current) .  resultsOnly
         <$> getDatabaseValues db
 
 -- | Returns a list of profile entries, and a mapping linking a non-error Id to its profile entry
 toReport :: Database -> IO ([ProfileEntry], HashMap Key Int)
 toReport db = do
     status <- prepareForDependencyOrder db
     let order = dependencyOrder show
-                $ map (second (getResultDepsDefault [Key "alwaysRerun"] . resultDeps))
+                $ map (second (Set.toList . getResultDepsDefault (Set.singleton $ newKey "alwaysRerun") . resultDeps))
                 $ Map.toList status
         ids = Map.fromList $ zip order [0..]
 
@@ -124,14 +124,14 @@ toReport db = do
             ,prfBuilt = fromStep resultBuilt
             ,prfVisited = fromStep resultVisited
             ,prfChanged = fromStep resultChanged
-            ,prfDepends = map pure $ mapMaybe (`Map.lookup` ids) $ getResultDepsDefault [Key "alwaysRerun"] resultDeps
+            ,prfDepends = map pure $ Map.elems $ Map.intersectionWith const ids $ Set.toMap $ getResultDepsDefault (Set.singleton $ newKey "alwaysRerun") resultDeps
             ,prfExecution = resultExecution
             }
             where fromStep i = fromJust $ Map.lookup i steps
     pure ([maybe (error "toReport") (f i) $ Map.lookup i status | i <- order], ids)
 
 alwaysRerunResult :: Step -> Result
-alwaysRerunResult current = Result (Value $ toDyn "<alwaysRerun>") (Step 0) (Step 0) current (ResultDeps []) 0 mempty
+alwaysRerunResult current = Result (Value $ toDyn "<alwaysRerun>") (Step 0) (Step 0) current (ResultDeps mempty) 0 mempty
 
 generateHTML :: Maybe [Int] -> [ProfileEntry] -> IO LBS.ByteString
 generateHTML dirtyKeys xs = do

hls-graph/src/Development/IDE/Graph/Internal/Types.hs

@@ -7,6 +7,9 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE RecordWildCards            #-}
 {-# LANGUAGE ScopedTypeVariables        #-}
+{-# LANGUAGE PatternSynonyms            #-}
+{-# LANGUAGE BangPatterns               #-}
+{-# LANGUAGE ViewPatterns               #-}
 
 module Development.IDE.Graph.Internal.Types where
 
@@ -27,6 +30,7 @@ import qualified Data.ByteString               as BS
 import           Data.Dynamic
 import qualified Data.HashMap.Strict           as Map
 import           Data.HashSet                  (HashSet, member)
+import qualified Data.IntMap                   as IM
 import qualified Data.HashSet                  as Set
 import           Data.IORef
 import           Data.List                     (intercalate)
@@ -38,6 +42,7 @@ import qualified ListT
 import qualified StmContainers.Map             as SMap
 import           StmContainers.Map             (Map)
 import           System.Time.Extra             (Seconds)
+import           System.IO.Unsafe
 import           UnliftIO                      (MonadUnliftIO)
 
 
@@ -84,16 +89,52 @@ data ShakeDatabase = ShakeDatabase !Int [Action ()] Database
 newtype Step = Step Int
     deriving newtype (Eq,Ord,Hashable)
 
-data Key = forall a . (Typeable a, Eq a, Hashable a, Show a) => Key a
+---------------------------------------------------------------------
+-- Keys
 
-instance Eq Key where
-    Key a == Key b = Just a == cast b
+data KeyValue = forall a . (Typeable a, Hashable a, Show a) => KeyValue a
 
-instance Hashable Key where
-    hashWithSalt i (Key x) = hashWithSalt i (typeOf x, x)
+newtype Key = UnsafeMkKey Int
+
+pattern Key a <- (lookupKeyValue -> KeyValue a)
+
+data KeyMap = KeyMap !(Map.HashMap KeyValue Key) !(IM.IntMap KeyValue) {-# UNPACK #-} !Int
+
+keyMap :: IORef KeyMap
+keyMap = unsafePerformIO $ newIORef (KeyMap Map.empty IM.empty 0)
 
+{-# NOINLINE keyMap #-}
+
+newKey :: (Typeable a, Hashable a, Show a) => a -> Key
+newKey k = unsafePerformIO $ do
+  let !newKey = KeyValue k
+  atomicModifyIORef' keyMap $ \km@(KeyMap hm im n) ->
+    let new_key = Map.lookup newKey hm
+    in case new_key of
+          Just v  -> (km, v)
+          Nothing ->
+            let !new_index = UnsafeMkKey n
+            in (KeyMap (Map.insert newKey new_index hm) (IM.insert n newKey im) (n+1), new_index)
+{-# NOINLINE newKey #-}
+
+lookupKeyValue :: Key -> KeyValue
+lookupKeyValue (UnsafeMkKey x) = unsafePerformIO $ do
+  KeyMap _ im _ <- readIORef keyMap
+  pure $! fromJust (IM.lookup x im)
+
+instance Eq Key where
+  UnsafeMkKey a == UnsafeMkKey b = a == b
+instance Hashable Key where
+  hashWithSalt i (UnsafeMkKey x) = hashWithSalt i x
 instance Show Key where
-    show (Key x) = show x
+  show (Key x) = show x
+
+instance Eq KeyValue where
+    KeyValue a == KeyValue b = Just a == cast b
+instance Hashable KeyValue where
+    hashWithSalt i (KeyValue x) = hashWithSalt i (typeOf x, x)
+instance Show KeyValue where
+    show (KeyValue x) = show x
 
 newtype Value = Value Dynamic
 
@@ -149,24 +190,24 @@ data Result = Result {
     resultData      :: !BS.ByteString
     }
 
-data ResultDeps = UnknownDeps | AlwaysRerunDeps ![Key] | ResultDeps ![Key]
+data ResultDeps = UnknownDeps | AlwaysRerunDeps !(HashSet Key) | ResultDeps !(HashSet Key)
   deriving (Eq, Show)
 
-getResultDepsDefault :: [Key] -> ResultDeps -> [Key]
+getResultDepsDefault :: (HashSet Key) -> ResultDeps -> (HashSet Key)
 getResultDepsDefault _ (ResultDeps ids)      = ids
 getResultDepsDefault _ (AlwaysRerunDeps ids) = ids
 getResultDepsDefault def UnknownDeps         = def
 
-mapResultDeps :: ([Key] -> [Key]) -> ResultDeps -> ResultDeps
+mapResultDeps :: (HashSet Key -> HashSet Key) -> ResultDeps -> ResultDeps
 mapResultDeps f (ResultDeps ids)      = ResultDeps $ f ids
 mapResultDeps f (AlwaysRerunDeps ids) = AlwaysRerunDeps $ f ids
 mapResultDeps _ UnknownDeps           = UnknownDeps
 
 instance Semigroup ResultDeps where
     UnknownDeps <> x = x
     x <> UnknownDeps = x
-    AlwaysRerunDeps ids <> x = AlwaysRerunDeps (ids <> getResultDepsDefault [] x)
-    x <> AlwaysRerunDeps ids = AlwaysRerunDeps (getResultDepsDefault [] x <> ids)
+    AlwaysRerunDeps ids <> x = AlwaysRerunDeps (ids <> getResultDepsDefault mempty x)
+    x <> AlwaysRerunDeps ids = AlwaysRerunDeps (getResultDepsDefault mempty x <> ids)
     ResultDeps ids <> ResultDeps ids' = ResultDeps (ids <> ids')
 
 instance Monoid ResultDeps where