Async.hs

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
-- | Async API for heap traversals or computations that can be executed
-- asynchronously.
--
-- Async variations of heap traverals produce a result of 'AsyncTraceT' or 'AsyncTrace'
-- which can be periodically queried for progress.
-- This useful for implementing responsive UIs, as the user doesn't have to wait for the
-- heap traversal to finish before seeing the result.
module GHC.Debug.Async (
  -- * Async API
  AsyncTraceT(..),
  AsyncTrace,
  mkAsyncTrace,
  withAsyncTrace,
  waitForAsyncTraceResultM,
  AsyncTraceResult(..),
  expectAsyncTraceResultIsDone,
  expectAsyncTraceResultIsDoneM,
) where

import GHC.Debug.Client.Monad
import GHC.Stack.Types (HasCallStack)
import Data.IORef

data AsyncTraceResult s = Done s  -- ^ The traversal has finished, the result is the accumulated state.
                        | Running s -- ^ The traversal is running, the current result
                        | NotStarted -- ^ The traversal has not started yet.
                        deriving (Foldable, Traversable, Functor, Show)

data AsyncTraceT m s = AsyncTrace
  { asyncResult :: IO (AsyncTraceResult s)
  -- ^ The current status of the computation.
  -- This should not be blocking indefinitely.
  , asyncWait :: m ()
  -- ^ Wait for the result to complete.
  -- Once completed, 'asyncResult' must contain a 'Done' value.
  } deriving (Functor)

type AsyncTrace = AsyncTraceT DebugM

mkAsyncTrace :: IO (AsyncTraceResult s) -> m () -> AsyncTraceT m s
mkAsyncTrace res waitRes = AsyncTrace
  { asyncResult = res
  , asyncWait = waitRes
  }

waitForAsyncTraceResultM :: DebugMonad m => AsyncTraceT m s -> m s
waitForAsyncTraceResultM asyncTrace = do
  asyncWait asyncTrace
  expectAsyncTraceResultIsDoneM $ asyncResult asyncTrace

expectAsyncTraceResultIsDone :: HasCallStack => IO (AsyncTraceResult a) -> IO a
expectAsyncTraceResultIsDone read_results = do
  res <- read_results
  case res of
    Done s -> return s
    NotStarted -> error "expectAsyncTraceResultIsDone: NotStarted"
    Running _ -> error "expectAsyncTraceResultIsDone: Running"

expectAsyncTraceResultIsDoneM :: (HasCallStack, DebugMonad m) => IO (AsyncTraceResult a) -> m a
expectAsyncTraceResultIsDoneM = unsafeLiftIO . expectAsyncTraceResultIsDone

-- | Simplified API, provide an action that be used to peek at the current progress of the traversal
-- and an action that is executed asynchronously. It is the responsibility of the caller to make sure that
-- the asynchronous action updates the progress.
withAsyncTrace :: forall a m . DebugMonad m => IO a -> m () -> m (AsyncTraceT m a)
withAsyncTrace read_results runHeapTrace = do
  progressVar <- unsafeLiftIO $ newIORef NotStarted

  let
    reportProgress :: IO (AsyncTraceResult a)
    reportProgress = do
      progress <- readIORef progressVar
      traverse (\_ -> read_results) progress

    performAction :: m ()
    performAction = do
        unsafeLiftIO $ writeIORef progressVar (Running ())
        runHeapTrace
        unsafeLiftIO $ writeIORef progressVar (Done ())

  pure $ mkAsyncTrace reportProgress performAction