1 files changed, 38 insertions, 40 deletions

diff --git a/src/Ghci.hs b/src/Ghci.hs
index 4c6dca1..7f06f99 100644
--- a/src/Ghci.hs
+++ b/src/Ghci.hs
@@ -13,7 +13,7 @@ module Ghci (
   parseSettingsPaste,
 ) where
 
-import Control.Exception (catch, SomeException)
+import Control.Exception (catch, catchNoPropagate, rethrowIO, ExceptionWithContext, SomeException)
 import Control.Monad (replicateM, when)
 import Data.Bifunctor (first)
 import qualified Data.ByteString as BS
@@ -26,7 +26,7 @@ import qualified Data.ByteString.Short as BSS
 import Data.Char (isSpace, toLower)
 import Data.List (nub)
 import Foreign (allocaBytes)
-import System.IO (hFlush, hIsClosed, hGetBufSome, hGetLine, hPutStrLn, stdout, stderr, Handle)
+import System.IO (hFlush, hIsClosed, hGetBufSome, hPutStrLn, stdout, stderr, Handle)
 import System.Process
 import System.Random (getStdRandom, uniformR)
 import System.Timeout (timeout, Timeout)
@@ -64,6 +64,9 @@ parseSettingsPaste = ParseSettings 50000 False
 data Result a = Error String | Ignored | Return a | Exited
   deriving (Show)
 
+data ShouldTerminate = TermNow | StillAlive
+  deriving (Show)
+
 makeGhci :: IO Ghci
 makeGhci = do
   (pipeOut, pipeIn) <- createPipe
@@ -111,9 +114,28 @@ runStmtClever ghci pset line =
     long `startsWith` short = take (length short) long == short
 
 runStmt :: Ghci -> ParseSettings -> String -> IO (Ghci, Result String)
-runStmt ghci pset line = timeouting 2_000_000 (restarting (\g -> runStmt' g pset line)) ghci
+runStmt ghci pset line = do
+  closed <- hIsClosed (ghciStdin ghci)
+  ghci' <- if closed
+             then do
+               putStrLn "ghci: restarting due to closed stdin"
+               hFlush stdout
+               terminateGhci ghci
+               makeGhci
+             else return ghci
+
+  (sterm, res) <- timeouting 2_000_000 (restarting (\g -> runStmt' g pset line)) ghci'
+
+  ghci'' <- case sterm of
+    TermNow -> do
+      terminateGhci ghci'
+      ghci'' <- makeGhci
+      return ghci''
+    StillAlive -> return ghci'
 
-timeouting :: Int -> (Ghci -> IO (Ghci, Result a)) -> Ghci -> IO (Ghci, Result a)
+  return (ghci'', res)
+
+timeouting :: Int -> (Ghci -> IO (ShouldTerminate, Result a)) -> Ghci -> IO (ShouldTerminate, Result a)
 timeouting microseconds f ghci =
   -- TODO: The timeout handling code never actually runs, because the timeout
   -- exception is already handled by the catch-all exception handler in
@@ -121,62 +143,38 @@ timeouting microseconds f ghci =
   timeout microseconds (f ghci) >>= \case
     Nothing -> do putStrLn "ghci: restarting due to timeout"
                   hFlush stdout
-                  terminateGhci ghci
-                  ghci' <- makeGhci
-                  return (ghci', Error "<timeout>")
+                  return (TermNow, Error "<timeout>")
     Just pair -> return pair
 
-restarting :: (Ghci -> IO a) -> Ghci -> IO (Ghci, Result a)
-restarting f ghci = do
-  closed <- hIsClosed (ghciStdin ghci)
-  ghci' <- if closed
-             then do
-               putStrLn "ghci: restarting due to closed stdin"
-               hFlush stdout
-               terminateGhci ghci
-               makeGhci
-             else return ghci
-
-  fmap (\x -> (ghci', Return x)) (f ghci')
-    `catch` (\e -> do let _ = e :: Timeout
-                      putStrLn "ghci: restarting due to timeout (caught in restarting)"
-                      hFlush stdout
-                      terminateGhci ghci
-                      -- putStrLn $ "ghci: terminated"
-                      -- hFlush stdout
-                      ghci'' <- makeGhci
-                      -- putStrLn $ "ghci: new made"
-                      -- hFlush stdout
-                      return (ghci'', Error "<timeout>"))
+restarting :: (Ghci -> IO (ShouldTerminate, a)) -> Ghci -> IO (ShouldTerminate, Result a)
+restarting f ghci =
+  fmap (\(st, x) -> (st, Return x)) (f ghci)
+    -- propagate Timeout exceptions to 'timeouting'
+    `catchNoPropagate` (\e -> rethrowIO (e :: ExceptionWithContext Timeout))
     `catch` (\e -> do let _ = e :: SomeException
                       putStrLn $ "ghci: restarting due to exception: " ++ show e
                       hFlush stdout
-                      terminateGhci ghci'
-                      -- putStrLn $ "ghci: terminated"
-                      -- hFlush stdout
-                      ghci'' <- makeGhci
-                      -- putStrLn $ "ghci: new made"
-                      -- hFlush stdout
-                      return (ghci'', Error "Oops, something went wrong"))
+                      return (TermNow, Error "Oops, something went wrong"))
 
 terminateGhci :: Ghci -> IO ()
 terminateGhci ghci = terminateProcess (ghciProc ghci)
 
-runStmt' :: Ghci -> ParseSettings -> String -> IO String
+runStmt' :: Ghci -> ParseSettings -> String -> IO (ShouldTerminate, String)
 runStmt' ghci pset stmt = do
   tag <- updatePrompt ghci
   ghciPutStrLn (ghciStdin ghci) stmt
   hFlush (ghciStdin ghci)
   let readmax = psMaxOutputLen pset + 8192
   (output, reason) <- hGetUntilUTF8 (ghciStdout ghci) (Just readmax) tag
+  hPutStrLn stderr $ "runStmt': reason = " ++ show reason
   case reason of
     ReachedMaxLen -> do
       terminateGhci ghci  -- because we lost the new prompt
-      return (formatOutput output)  -- don't need to strip tag because we read more than the max output len
-    ReachedTag -> return (formatOutput $ take (length output - length tag) output)
+      return (TermNow, formatOutput output)  -- don't need to strip tag because we read more than the max output len
+    ReachedTag -> return (StillAlive, formatOutput $ take (length output - length tag) output)
     ReachedEOF -> do
       terminateGhci ghci
-      return (formatOutput output)
+      return (TermNow, formatOutput output)
   where
     formatOutput output =
       let output' | psJoinLines pset = replaceNewlines (dropBothSlow isSpace output)