From b5e675eb262e4ed700c30e5a97f73f4d2eda9a7e Mon Sep 17 00:00:00 2001
From: Dougal <d.maclaurin@gmail.com>
Date: Mon, 8 Jan 2024 11:45:17 -0500
Subject: [PATCH] Drop source info from `import` statements.

The source info from evaluating the module itself doesn't apply to the
`import foo` line.
---
 src/lib/MonadUtil.hs | 23 +++++++++++++++++++++--
 src/lib/TopLevel.hs  | 18 +++++++++++++++++-
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/src/lib/MonadUtil.hs b/src/lib/MonadUtil.hs
index 17a21bd95..2d6acfbf8 100644
--- a/src/lib/MonadUtil.hs
+++ b/src/lib/MonadUtil.hs
@@ -8,12 +8,14 @@
 
 module MonadUtil (
   DefuncState (..), LabelReader (..), SingletonLabel (..), FreshNames (..),
-  runFreshNameT, FreshNameT (..), Logger (..), LogLevel (..), getIOLogger,
-  IOLoggerT (..), runIOLoggerT, LoggerT (..), runLoggerT, IOLogger (..), HasIOLogger (..)) where
+  runFreshNameT, FreshNameT (..), Logger (..), LogLevel (..), getIOLogger, CanSetIOLogger (..),
+  IOLoggerT (..), runIOLoggerT, LoggerT (..), runLoggerT,
+  IOLogger (..), HasIOLogger (..), captureIOLogs) where
 
 import Control.Monad.Reader
 import Control.Monad.State.Strict
 import Control.Monad.Writer.Strict
+import Data.IORef
 
 import Err
 
@@ -70,9 +72,16 @@ newtype IOLoggerT w m a = IOLoggerT { runIOLoggerT' :: ReaderT (IOLogger w) m a
 class Monad m => HasIOLogger w m | m -> w where
   getIOLogAction :: Monad m => m (w -> IO ())
 
+class Monad m => CanSetIOLogger w m | m -> w where
+  withIOLogAction :: Monad m => (w -> IO ()) -> m a -> m a
+
 instance (Monoid w, MonadIO m) => HasIOLogger w (IOLoggerT w m) where
   getIOLogAction = IOLoggerT $ asks ioLogAction
 
+instance (Monoid w, MonadIO m) => CanSetIOLogger w (IOLoggerT w m) where
+  withIOLogAction logger (IOLoggerT m) = IOLoggerT do
+    local (\r -> r { ioLogAction = logger }) m
+
 instance (Monoid w, MonadIO m) => Logger w (IOLoggerT w m) where
   emitLog w = do
     logger <- getIOLogAction
@@ -94,3 +103,13 @@ instance (Monoid w, Monad m) => Logger w (LoggerT w m) where
 
 runLoggerT :: (Monoid w, Monad m) => LoggerT w m a -> m (a, w)
 runLoggerT cont = runWriterT (runLoggerT' cont)
+
+captureIOLogs
+  :: forall w m a. (Monoid w, MonadIO m, HasIOLogger w m, CanSetIOLogger w m)
+  => m a -> m (a, w)
+captureIOLogs cont = do
+  ref <- liftIO $ newIORef (mempty :: w)
+  ans <- withIOLogAction (\w -> modifyIORef ref (<> w)) cont
+  w <- liftIO $ readIORef ref
+  return (ans, w)
+
diff --git a/src/lib/TopLevel.hs b/src/lib/TopLevel.hs
index aa0f94bd4..9ab73bf7d 100644
--- a/src/lib/TopLevel.hs
+++ b/src/lib/TopLevel.hs
@@ -103,6 +103,7 @@ type TopLogger m = (MonadIO m, Logger Outputs m)
 class ( forall n. Fallible (m n)
       , forall n. Logger Outputs (m n)
       , forall n. HasIOLogger Outputs (m n)
+      , forall n. CanSetIOLogger Outputs (m n)
       , forall n. Catchable (m n)
       , forall n. ConfigReader (m n)
       , forall n. RuntimeEnvReader (m n)
@@ -421,7 +422,7 @@ evalPartiallyParsedUModule partiallyParsed = do
 
 -- Assumes all module dependencies have been loaded already
 evalUModule :: (Topper m, Mut n) => UModule -> m n (Module n)
-evalUModule (UModule name _ blocks) = do
+evalUModule (UModule name _ blocks) = dropSourceInfoLogging do
   Abs topFrag UnitE <- localTopBuilder $ mapM_ (evalSourceBlock' name) blocks >> return UnitE
   TopEnvFrag envFrag moduleEnvFrag otherUpdates <- return topFrag
   ModuleEnv (ImportStatus directDeps transDeps) sm scs <- return moduleEnvFrag
@@ -429,6 +430,17 @@ evalUModule (UModule name _ blocks) = do
   let evaluatedModule = Module name directDeps transDeps sm scs
   emitEnv $ Abs fragToReEmit evaluatedModule
 
+dropSourceInfoLogging :: Topper m => m n a -> m n a
+dropSourceInfoLogging cont = do
+  (ans, Outputs logs) <- captureIOLogs cont
+  let logs' = filter isNotSourceInfo logs
+  emitLog $ Outputs logs'
+  return ans
+  where
+    isNotSourceInfo = \case
+      SourceInfo _ -> False
+      _ -> True
+
 importModule :: (Mut n, TopBuilder m, Fallible1 m) => ModuleSourceName -> m n ()
 importModule name = do
   lookupLoadedModule name >>= \case
@@ -812,6 +824,10 @@ instance Logger Outputs (TopperM n) where
 instance HasIOLogger Outputs (TopperM n) where
   getIOLogAction = TopperM $ asks topperLogAction
 
+instance CanSetIOLogger Outputs (TopperM n) where
+  withIOLogAction logger (TopperM m) = TopperM do
+    local (\r -> r { topperLogAction = logger }) m
+
 instance Generic TopStateEx where
   type Rep TopStateEx = Rep (Env UnsafeS, RuntimeEnv)
   from (TopStateEx env rtEnv) = from ((unsafeCoerceE env :: Env UnsafeS), rtEnv)