Improve hls graph

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

@@ -17,6 +17,7 @@ import           Control.Concurrent.Extra
 import           Control.Concurrent.STM.Stats         (STM, atomically,
                                                        atomicallyNamed,
                                                        modifyTVar', newTVarIO,
+                                                       putTMVar, readTMVar,
                                                        readTVarIO)
 import           Control.Exception
 import           Control.Monad
@@ -25,7 +26,6 @@ import           Control.Monad.Trans.Class            (lift)
 import           Control.Monad.Trans.Reader
 import qualified Control.Monad.Trans.State.Strict     as State
 import           Data.Dynamic
-import           Data.Either
 import           Data.Foldable                        (for_, traverse_)
 import           Data.IORef.Extra
 import           Data.Maybe
@@ -39,8 +39,10 @@ import           Development.IDE.Graph.Internal.Types
 import qualified Focus
 import qualified ListT
 import qualified StmContainers.Map                    as SMap
-import           System.IO.Unsafe
 import           System.Time.Extra                    (duration, sleep)
+import           UnliftIO                             (MonadUnliftIO (withRunInIO),
+                                                       newEmptyTMVarIO)
+import qualified UnliftIO.Exception                   as UE
 
 #if MIN_VERSION_base(4,19,0)
 import           Data.Functor                         (unzip)
@@ -78,7 +80,7 @@ incDatabase db Nothing = do
 updateDirty :: Monad m => Focus.Focus KeyDetails m ()
 updateDirty = Focus.adjust $ \(KeyDetails status rdeps) ->
             let status'
-                  | Running _ _ _ x <- status = Dirty x
+                  | Running _ _ x <- status = Dirty x
                   | Clean x <- status = Dirty (Just x)
                   | otherwise = status
             in KeyDetails status' rdeps
@@ -88,58 +90,60 @@ build
     => Database -> Stack -> f key -> IO (f Key, f value)
 -- build _ st k | traceShow ("build", st, k) False = undefined
 build db stack keys = do
-    built <- runAIO $ do
-        built <- builder db stack (fmap newKey keys)
-        case built of
-          Left clean  -> return clean
-          Right dirty -> liftIO dirty
+    !built <- runAIO $ builder db stack (fmap newKey keys)
     let (ids, vs) = unzip built
     pure (ids, fmap (asV . resultValue) vs)
     where
         asV :: Value -> value
         asV (Value x) = unwrapDynamic x
 
+data BuildArity = BuildUnary | BuildNary
 -- | Build a list of keys and return their results.
 --  If none of the keys are dirty, we can return the results immediately.
 --  Otherwise, a blocking computation is returned *which must be evaluated asynchronously* to avoid deadlock.
-builder
-    :: Traversable f => Database -> Stack -> f Key -> AIO (Either (f (Key, Result)) (IO (f (Key, Result))))
+builder :: (Traversable f) => Database -> Stack -> f Key -> AIO (f (Key, Result))
 -- builder _ st kk | traceShow ("builder", st,kk) False = undefined
-builder db@Database{..} stack keys = withRunInIO $ \(RunInIO run) -> do
-    -- Things that I need to force before my results are ready
-    toForce <- liftIO $ newTVarIO []
-    current <- liftIO $ readTVarIO databaseStep
-    results <- liftIO $ for keys $ \id ->
-        -- Updating the status of all the dependencies atomically is not necessary.
-        -- Therefore, run one transaction per dep. to avoid contention
-        atomicallyNamed "builder" $ do
-            -- Spawn the id if needed
-            status <- SMap.lookup id databaseValues
-            val <- case viewDirty current $ maybe (Dirty Nothing) keyStatus status of
-                Clean r -> pure r
-                Running _ force val _
-                  | memberStack id stack -> throw $ StackException stack
-                  | otherwise -> do
-                    modifyTVar' toForce (Wait force :)
-                    pure val
-                Dirty s -> do
-                    let act = run (refresh db stack id s)
-                        (force, val) = splitIO (join act)
-                    SMap.focus (updateStatus $ Running current force val s) id databaseValues
-                    modifyTVar' toForce (Spawn force:)
-                    pure val
-
-            pure (id, val)
-
-    toForceList <- liftIO $ readTVarIO toForce
-    let waitAll = run $ waitConcurrently_ toForceList
-    case toForceList of
-        [] -> return $ Left results
-        _ -> return $ Right $ do
-                waitAll
-                pure results
-
-
+builder db stack keys = do
+    let ba = if length keys == 1 then BuildUnary else BuildNary
+    keyWaits <- for keys $ \k -> builderOne ba db stack k
+    !res <- for keyWaits $ \(k, waitR) -> do
+        !v<- liftIO waitR
+        return (k, v)
+    return res
+
+builderOne :: BuildArity -> Database -> Stack -> Key -> AIO (Key, IO Result)
+builderOne ba db@Database {..} stack id = UE.mask $ \restore -> do
+  current <- liftIO $ readTVarIO databaseStep
+  barrier <- newEmptyTMVarIO
+  (k, registerWaitResult) <- liftIO $ atomicallyNamed "builder" $ do
+    -- Spawn the id if needed
+    status <- SMap.lookup id databaseValues
+    val <-
+      let refreshRsult s = do
+            let putResult act = do
+                  res <- act
+                  liftIO $ atomically $ putTMVar barrier res
+                  return res
+            let act = restore $ (case ba of
+                  BuildNary ->
+                    asyncWithCleanUp $
+                      putResult $ refresh db stack id s
+                  BuildUnary -> fmap return $ putResult $ refresh db stack id s)
+                        `UE.onException` (UE.uninterruptibleMask_ $ liftIO (atomicallyNamed "builder - onException" (SMap.focus updateDirty id databaseValues)))
+            -- Mark the key as running
+            SMap.focus (updateStatus $ Running current (atomically $ readTMVar barrier) s) id databaseValues
+            return act
+       in case viewDirty current $ maybe (Dirty Nothing) keyStatus status of
+            Dirty mbr -> refreshRsult mbr
+            Running step ba _mbr
+              | step /= current -> error $ "Inconsistent database state: key " ++ show id ++ " is marked Running at step " ++ show step ++ " but current step is " ++ show current
+              | memberStack id stack -> throw $ StackException stack
+              | otherwise -> pure . pure $ ba
+            Clean r -> pure . pure . pure $ r
+    -- force here might contains async exceptions from previous runs
+    pure (id, val)
+  waitR <- registerWaitResult
+  return (k, waitR)
 -- | isDirty
 -- only dirty when it's build time is older than the changed time of one of its dependencies
 isDirty :: Foldable t => Result -> t (a, Result) -> Bool
@@ -155,31 +159,27 @@ isDirty me = any (\(_,dep) -> resultBuilt me < resultChanged dep)
 refreshDeps :: KeySet -> Database -> Stack -> Key -> Result -> [KeySet] -> AIO Result
 refreshDeps visited db stack key result = \case
     -- no more deps to refresh
-    [] -> liftIO $ compute db stack key RunDependenciesSame (Just result)
+    [] -> compute' db stack key RunDependenciesSame (Just result)
     (dep:deps) -> do
         let newVisited = dep <> visited
         res <- builder db stack (toListKeySet (dep `differenceKeySet` visited))
-        case res of
-            Left res ->  if isDirty result res
+        if isDirty result res
                 -- restart the computation if any of the deps are dirty
-                then liftIO $ compute db stack key RunDependenciesChanged (Just result)
+                then compute' db stack key RunDependenciesChanged (Just result)
                 -- else kick the rest of the deps
                 else refreshDeps newVisited db stack key result deps
-            Right iores -> do
-                res <- liftIO iores
-                if isDirty result res
-                    then liftIO $ compute db stack key RunDependenciesChanged (Just result)
-                    else refreshDeps newVisited db stack key result deps
 
--- | Refresh a key:
-refresh :: Database -> Stack -> Key -> Maybe Result -> AIO (IO Result)
+
+-- refresh :: Database -> Stack -> Key -> Maybe Result -> IO Result
 -- refresh _ st k _ | traceShow ("refresh", st, k) False = undefined
+refresh :: Database -> Stack -> Key -> Maybe Result -> AIO Result
 refresh db stack key result = case (addStack key stack, result) of
     (Left e, _) -> throw e
-    (Right stack, Just me@Result{resultDeps = ResultDeps deps}) -> asyncWithCleanUp $ refreshDeps mempty db stack key me (reverse deps)
-    (Right stack, _) ->
-        asyncWithCleanUp $ liftIO $ compute db stack key RunDependenciesChanged result
+    (Right stack, Just me@Result{resultDeps = ResultDeps deps}) -> refreshDeps mempty db stack key me (reverse deps)
+    (Right stack, _) -> compute' db stack key RunDependenciesChanged result
 
+compute' :: Database -> Stack -> Key -> RunMode -> Maybe Result -> AIO Result
+compute' db stack key mode result = liftIO $ compute db stack key mode result
 -- | Compute a key.
 compute :: Database -> Stack -> Key -> RunMode -> Maybe Result -> IO Result
 -- compute _ st k _ _ | traceShow ("compute", st, k) False = undefined
@@ -247,18 +247,6 @@ getKeysAndVisitAge db = do
         getAge Result{resultVisited = Step s} = curr - s
     return keysWithVisitAge
 --------------------------------------------------------------------------------
--- Lazy IO trick
-
-data Box a = Box {fromBox :: a}
-
--- | Split an IO computation into an unsafe lazy value and a forcing computation
-splitIO :: IO a -> (IO (), a)
-splitIO act = do
-    let act2 = Box <$> act
-    let res = unsafePerformIO act2
-    (void $ evaluate res, fromBox res)
-
---------------------------------------------------------------------------------
 -- Reverse dependencies
 
 -- | Update the reverse dependencies of an Id
@@ -307,8 +295,12 @@ newtype AIO a = AIO { unAIO :: ReaderT (IORef [Async ()]) IO a }
 -- | Run the monadic computation, cancelling all the spawned asyncs if an exception arises
 runAIO :: AIO a -> IO a
 runAIO (AIO act) = do
-    asyncs <- newIORef []
-    runReaderT act asyncs `onException` cleanupAsync asyncs
+    asyncsRef <- newIORef []
+    -- Log the exact exception (including async exceptions) before cleanup,
+    -- then rethrow to preserve previous semantics.
+    runReaderT act asyncsRef `onException` do
+        asyncs <- atomicModifyIORef' asyncsRef ([],)
+        cleanupAsync asyncs
 
 -- | Like 'async' but with built-in cancellation.
 --   Returns an IO action to wait on the result.
@@ -319,25 +311,22 @@ asyncWithCleanUp act = do
     -- mask to make sure we keep track of the spawned async
     liftIO $ uninterruptibleMask $ \restore -> do
         a <- async $ restore io
-        atomicModifyIORef'_ st (void a :)
+        atomicModifyIORef'_ st (void a:)
         return $ wait a
 
 unliftAIO :: AIO a -> AIO (IO a)
 unliftAIO act = do
     st <- AIO ask
     return $ runReaderT (unAIO act) st
 
-newtype RunInIO = RunInIO (forall a. AIO a -> IO a)
-
-withRunInIO :: (RunInIO -> AIO b) -> AIO b
-withRunInIO k = do
-    st <- AIO ask
-    k $ RunInIO (\aio -> runReaderT (unAIO aio) st)
+instance MonadUnliftIO AIO where
+    withRunInIO k = do
+        st <- AIO ask
+        liftIO $ k (\aio -> runReaderT (unAIO aio) st)
 
-cleanupAsync :: IORef [Async a] -> IO ()
+cleanupAsync :: [Async a] -> IO ()
 -- mask to make sure we interrupt all the asyncs
-cleanupAsync ref = uninterruptibleMask $ \unmask -> do
-    asyncs <- atomicModifyIORef' ref ([],)
+cleanupAsync asyncs = uninterruptibleMask $ \unmask -> do
     -- interrupt all the asyncs without waiting
     mapM_ (\a -> throwTo (asyncThreadId a) AsyncCancelled) asyncs
     -- Wait until all the asyncs are done
@@ -348,32 +337,3 @@ cleanupAsync ref = uninterruptibleMask $ \unmask -> do
                 traceM "cleanupAsync: waiting for asyncs to finish"
         withAsync warnIfTakingTooLong $ \_ ->
             mapM_ waitCatch asyncs
-
-data Wait
-    = Wait {justWait :: !(IO ())}
-    | Spawn {justWait :: !(IO ())}
-
-fmapWait :: (IO () -> IO ()) -> Wait -> Wait
-fmapWait f (Wait io)  = Wait (f io)
-fmapWait f (Spawn io) = Spawn (f io)
-
-waitOrSpawn :: Wait -> IO (Either (IO ()) (Async ()))
-waitOrSpawn (Wait io)  = pure $ Left io
-waitOrSpawn (Spawn io) = Right <$> async io
-
-waitConcurrently_ :: [Wait] -> AIO ()
-waitConcurrently_ [] = pure ()
-waitConcurrently_ [one] = liftIO $ justWait one
-waitConcurrently_ many = do
-    ref <- AIO ask
-    -- spawn the async computations.
-    -- mask to make sure we keep track of all the asyncs.
-    (asyncs, syncs) <- liftIO $ uninterruptibleMask $ \unmask -> do
-        waits <- liftIO $ traverse (waitOrSpawn . fmapWait unmask) many
-        let (syncs, asyncs) = partitionEithers waits
-        liftIO $ atomicModifyIORef'_ ref (asyncs ++)
-        return (asyncs, syncs)
-    -- work on the sync computations
-    liftIO $ sequence_ syncs
-    -- wait for the async computations before returning
-    liftIO $ traverse_ wait asyncs

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

@@ -6,7 +6,7 @@
 module Development.IDE.Graph.Internal.Types where
 
 import           Control.Concurrent.STM             (STM)
-import           Control.Monad                      ((>=>))
+import           Control.Monad                      (void, (>=>))
 import           Control.Monad.Catch
 import           Control.Monad.IO.Class
 import           Control.Monad.Trans.Reader
@@ -89,7 +89,7 @@ waitForDatabaseRunningKeysAction = getDatabase >>= liftIO . waitForDatabaseRunni
 data ShakeDatabase = ShakeDatabase !Int [Action ()] Database
 
 newtype Step = Step Int
-    deriving newtype (Eq,Ord,Hashable,Show)
+    deriving newtype (Eq,Ord,Hashable,Show,Num,Enum,Real,Integral)
 
 ---------------------------------------------------------------------
 -- Keys
@@ -129,23 +129,23 @@ data Status
     = Clean !Result
     | Dirty (Maybe Result)
     | Running {
-        runningStep   :: !Step,
-        runningWait   :: !(IO ()),
-        runningResult :: Result,     -- LAZY
-        runningPrev   :: !(Maybe Result)
+        runningStep :: !Step,
+        runningWait :: !(IO Result),
+        -- runningResult :: Result,     -- LAZY
+        runningPrev :: !(Maybe Result)
         }
 
 viewDirty :: Step -> Status -> Status
-viewDirty currentStep (Running s _ _ re) | currentStep /= s = Dirty re
+viewDirty currentStep (Running s _ re) | currentStep /= s = Dirty re
 viewDirty _ other = other
 
 getResult :: Status -> Maybe Result
-getResult (Clean re)           = Just re
-getResult (Dirty m_re)         = m_re
-getResult (Running _ _ _ m_re) = m_re -- watch out: this returns the previous result
+getResult (Clean re)         = Just re
+getResult (Dirty m_re)       = m_re
+getResult (Running _ _ m_re) = m_re -- watch out: this returns the previous result
 
 waitRunning :: Status -> IO ()
-waitRunning Running{..} = runningWait
+waitRunning Running{..} = void runningWait
 waitRunning _           = return ()
 
 data Result = Result {

plugins/hls-signature-help-plugin/src/Ide/Plugin/SignatureHelp.hs

@@ -94,7 +94,7 @@ help when a space char is inserted, we probably have to use up-to-date results.
 
 {-
 Here is a brief description of the algorithm of finding relevant bits from HIE AST
-1. let 'hsAppNode' = the smallest 'HsApp' AST node which contains the cursor postion
+1. let 'hsAppNode' = the smallest 'HsApp' AST node which contains the cursor position
    See 'extractInfoFromSmallestContainingFunctionApplicationAst'
 2. let 'functionNode' = the left-most node of 'hsAppNode'
    See 'getLeftMostNode'

plugins/hls-signature-help-plugin/test/Main.hs

@@ -5,7 +5,6 @@ import           Control.Arrow                            ((>>>))
 import           Control.Exception                        (throw)
 import           Control.Lens                             ((^.))
 import           Data.Maybe                               (fromJust)
-import           Data.String.Interpolate                  (__i)
 import           Data.Text                                (Text)
 import qualified Data.Text                                as T
 import           Development.IDE.Plugin.Completions.Types (PosPrefixInfo (PosPrefixInfo))