1 files changed, 78 insertions, 25 deletions

diff --git a/Omni/Log/Concurrent.hs b/Omni/Log/Concurrent.hs
index 5187367..f5a420f 100644
--- a/Omni/Log/Concurrent.hs
+++ b/Omni/Log/Concurrent.hs
@@ -10,11 +10,13 @@ module Omni.Log.Concurrent
     reserveLine,
     updateLine,
     releaseLine,
+    updateCurrentLine,
+    releaseCurrentLine,
   )
 where
 
 import Alpha
-import Data.IORef (IORef, atomicModifyIORef', modifyIORef', newIORef, readIORef)
+import Data.IORef (IORef, atomicModifyIORef', modifyIORef', newIORef, readIORef, writeIORef)
 import qualified Data.Map as Map
 import qualified Data.Text as Text
 import Omni.Namespace (Namespace)
@@ -22,6 +24,7 @@ import qualified Omni.Namespace as Namespace
 import qualified System.Console.ANSI as ANSI
 import qualified System.Environment as Env
 import qualified System.IO as IO
+import System.IO.Unsafe (unsafePerformIO)
 
 data BuildState = Building | Success | Failed
   deriving (Eq, Show)
@@ -39,6 +42,14 @@ data LineManager = LineManager
     lmSupportsANSI :: Bool
   }
 
+{-# NOINLINE currentLineManager #-}
+currentLineManager :: IORef (Maybe LineManager)
+currentLineManager = unsafePerformIO (newIORef Nothing)
+
+{-# NOINLINE namespaceLines #-}
+namespaceLines :: IORef (Map Namespace Int)
+namespaceLines = unsafePerformIO (newIORef Map.empty)
+
 withLineManager :: Int -> (LineManager -> IO a) -> IO a
 withLineManager maxLines action = do
   supportsANSI <- checkANSISupport
@@ -47,13 +58,18 @@ withLineManager maxLines action = do
     then do
       linesRef <- newIORef Map.empty
       currentRef <- newIORef 0
-      action
-        LineManager
-          { lmLines = linesRef,
-            lmMaxLines = 1,
-            lmCurrentLine = currentRef,
-            lmSupportsANSI = False
-          }
+      let mgr =
+            LineManager
+              { lmLines = linesRef,
+                lmMaxLines = 1,
+                lmCurrentLine = currentRef,
+                lmSupportsANSI = False
+              }
+      writeIORef currentLineManager (Just mgr)
+      result <- action mgr
+      writeIORef currentLineManager Nothing
+      writeIORef namespaceLines Map.empty
+      pure result
     else do
       replicateM_ maxLines (IO.hPutStrLn IO.stderr "")
       ANSI.hCursorUp IO.stderr maxLines
@@ -61,16 +77,20 @@ withLineManager maxLines action = do
       linesRef <- newIORef (Map.fromList [(i, Nothing) | i <- [0 .. maxLines - 1]])
       currentRef <- newIORef maxLines
 
-      result <-
-        action
-          LineManager
-            { lmLines = linesRef,
-              lmMaxLines = maxLines,
-              lmCurrentLine = currentRef,
-              lmSupportsANSI = True
-            }
+      let mgr =
+            LineManager
+              { lmLines = linesRef,
+                lmMaxLines = maxLines,
+                lmCurrentLine = currentRef,
+                lmSupportsANSI = True
+              }
+      writeIORef currentLineManager (Just mgr)
+
+      result <- action mgr
 
       ANSI.hCursorDown IO.stderr maxLines
+      writeIORef currentLineManager Nothing
+      writeIORef namespaceLines Map.empty
       pure result
 
 checkANSISupport :: IO Bool
@@ -87,14 +107,19 @@ reserveLine :: LineManager -> Namespace -> IO (Maybe Int)
 reserveLine LineManager {..} ns =
   if not lmSupportsANSI
     then pure Nothing
-    else
-      atomicModifyIORef' lmLines <| \linesMap ->
-        case findFirstFree linesMap of
-          Nothing -> (linesMap, Nothing)
-          Just lineNum ->
-            let status = BuildStatus ns "" Building
-                linesMap' = Map.insert lineNum (Just status) linesMap
-             in (linesMap', Just lineNum)
+    else do
+      mLine <-
+        atomicModifyIORef' lmLines <| \linesMap ->
+          case findFirstFree linesMap of
+            Nothing -> (linesMap, Nothing)
+            Just lineNum ->
+              let status = BuildStatus ns "" Building
+                  linesMap' = Map.insert lineNum (Just status) linesMap
+               in (linesMap', Just lineNum)
+      case mLine of
+        Just lineNum -> modifyIORef' namespaceLines (Map.insert ns lineNum)
+        Nothing -> pure ()
+      pure mLine
   where
     findFirstFree :: Map Int (Maybe BuildStatus) -> Maybe Int
     findFirstFree m =
@@ -104,7 +129,7 @@ reserveLine LineManager {..} ns =
         |> listToMaybe
 
 updateLine :: LineManager -> Maybe Int -> Namespace -> Text -> IO ()
-updateLine LineManager {..} mLineNum ns output =
+updateLine LineManager {..} mLineNum _ output =
   if not lmSupportsANSI
     then do
       IO.hPutStr IO.stderr (Text.unpack <| output <> "\n")
@@ -153,3 +178,31 @@ releaseLine LineManager {..} mLineNum buildState =
         IO.hFlush IO.stderr
 
         ANSI.hRestoreCursor IO.stderr
+
+updateCurrentLine :: Namespace -> Text -> IO ()
+updateCurrentLine ns output = do
+  mMgr <- readIORef currentLineManager
+  case mMgr of
+    Nothing -> do
+      IO.hPutStr IO.stderr (Text.unpack <| output <> "\r")
+      IO.hFlush IO.stderr
+    Just mgr -> do
+      nsMap <- readIORef namespaceLines
+      case Map.lookup ns nsMap of
+        Nothing -> do
+          IO.hPutStr IO.stderr (Text.unpack <| output <> "\r")
+          IO.hFlush IO.stderr
+        Just lineNum -> updateLine mgr (Just lineNum) ns output
+
+releaseCurrentLine :: Namespace -> BuildState -> IO ()
+releaseCurrentLine ns buildState = do
+  mMgr <- readIORef currentLineManager
+  case mMgr of
+    Nothing -> pure ()
+    Just mgr -> do
+      nsMap <- readIORef namespaceLines
+      case Map.lookup ns nsMap of
+        Nothing -> pure ()
+        Just lineNum -> do
+          releaseLine mgr (Just lineNum) buildState
+          modifyIORef' namespaceLines (Map.delete ns)