From 4245ef519126d991f365d600363b7d3ba37a4ba6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= <sagemueller@geo.uni-koeln.de>
Date: Tue, 14 Mar 2017 18:51:01 +0100
Subject: [PATCH] Include RTS flags in kernelspec when installing.

---
 main/Main.hs            |  2 ++
 src/IHaskell/Flags.hs   | 18 +++++++++++++++++-
 src/IHaskell/IPython.hs |  5 +++++
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/main/Main.hs b/main/Main.hs
index af4f7b57..abe33953 100644
--- a/main/Main.hs
+++ b/main/Main.hs
@@ -113,6 +113,8 @@ parseKernelArgs = foldl' addFlag defaultKernelSpecOptions
       kernelSpecOpts { kernelSpecDebug = True }
     addFlag kernelSpecOpts (GhcLibDir libdir) =
       kernelSpecOpts { kernelSpecGhcLibdir = libdir }
+    addFlag kernelSpecOpts (RTSFlags rts) =
+      kernelSpecOpts { kernelSpecRTSOptions = rts }
     addFlag kernelSpecOpts (KernelspecInstallPrefix prefix) =
       kernelSpecOpts { kernelSpecInstallPrefix = Just prefix }
     addFlag kernelSpecOpts KernelspecUseStack =
diff --git a/src/IHaskell/Flags.hs b/src/IHaskell/Flags.hs
index 2ba9d9aa..9a56bf19 100644
--- a/src/IHaskell/Flags.hs
+++ b/src/IHaskell/Flags.hs
@@ -31,6 +31,8 @@ data Args = Args IHaskellMode [Argument]
 data Argument = ConfFile String     -- ^ A file with commands to load at startup.
               | OverwriteFiles      -- ^ Present when output should overwrite existing files. 
               | GhcLibDir String    -- ^ Where to find the GHC libraries.
+              | RTSFlags [String]   -- ^ Options for the GHC runtime (e.g. heap-size limit
+                                    --     or number of threads).
               | KernelDebug         -- ^ Spew debugging output from the kernel.
               | Help                -- ^ Display help text.
               | Version             -- ^ Display version text.
@@ -96,6 +98,20 @@ help mode = showText (Wrap 100) $ helpText [] HelpFormatAll $ chooseMode mode
 ghcLibFlag :: Flag Args
 ghcLibFlag = flagReq ["ghclib", "l"] (store GhcLibDir) "<path>" "Library directory for GHC."
 
+ghcRTSFlag :: Flag Args
+ghcRTSFlag = flagReq ["use-rtsopts"] storeRTS "\"<flags>\""
+                  "Runtime options (multithreading etc.). See `ghc +RTS -?`."
+ where storeRTS allRTSFlags (Args mode prev)
+          = fmap (Args mode . (:prev) . RTSFlags)
+              . parseRTS . words $ filter (/='"') allRTSFlags
+       parseRTS ("+RTS":fs)  -- Ignore if this is included (we already wrap
+           = parseRTS fs     -- the ihaskell-kernel call in +RTS <flags> -RTS anyway)
+       parseRTS ["-RTS"] = Right []
+       parseRTS ("-RTS":_)  -- Evil injection of extra arguments? Unlikely, but...
+           = Left "Adding non-RTS options to --use-rtsopts not permitted."
+       parseRTS (f:fs) = (f:) <$> parseRTS fs
+       parseRTS [] = Right []
+
 kernelDebugFlag :: Flag Args
 kernelDebugFlag = flagNone ["debug"] addDebug "Print debugging output from the kernel."
   where
@@ -125,7 +141,7 @@ store constructor str (Args mode prev) = Right $ Args mode $ constructor str : p
 installKernelSpec :: Mode Args
 installKernelSpec =
   mode "install" (Args InstallKernelSpec []) "Install the Jupyter kernelspec." noArgs
-    [ghcLibFlag, kernelDebugFlag, confFlag, installPrefixFlag, helpFlag, kernelStackFlag]
+    [ghcLibFlag, ghcRTSFlag, kernelDebugFlag, confFlag, installPrefixFlag, helpFlag, kernelStackFlag]
 
 kernel :: Mode Args
 kernel = mode "kernel" (Args (Kernel Nothing) []) "Invoke the IHaskell kernel." kernelArg
diff --git a/src/IHaskell/IPython.hs b/src/IHaskell/IPython.hs
index 345b007d..2669d469 100644
--- a/src/IHaskell/IPython.hs
+++ b/src/IHaskell/IPython.hs
@@ -45,6 +45,7 @@ import           StringUtils (replace, split)
 data KernelSpecOptions =
        KernelSpecOptions
          { kernelSpecGhcLibdir :: String           -- ^ GHC libdir.
+         , kernelSpecRTSOptions :: [String]        -- ^ Runtime options to use.
          , kernelSpecDebug :: Bool                 -- ^ Spew debugging output?
          , kernelSpecConfFile :: IO (Maybe String) -- ^ Filename of profile JSON file.
          , kernelSpecInstallPrefix :: Maybe String
@@ -54,6 +55,7 @@ data KernelSpecOptions =
 defaultKernelSpecOptions :: KernelSpecOptions
 defaultKernelSpecOptions = KernelSpecOptions
   { kernelSpecGhcLibdir = GHC.Paths.libdir
+  , kernelSpecRTSOptions = []
   , kernelSpecDebug = False
   , kernelSpecConfFile = defaultConfFile
   , kernelSpecInstallPrefix = Nothing
@@ -191,6 +193,9 @@ installKernelspec replace opts = void $ do
            Nothing   -> []
            Just file -> ["--conf", file])
         ++ ["--ghclib", kernelSpecGhcLibdir opts]
+        ++ (case kernelSpecRTSOptions opts of
+             [] -> []
+             rtsOpts -> "+RTS" : kernelSpecRTSOptions opts ++ ["-RTS"])
            ++ ["--stack" | kernelSpecUseStack opts]
 
   let kernelSpec = KernelSpec