copy the plugin's FastStringTable changes back into the host compiler

compiler/simplCore/CoreMonad.lhs

@@ -44,7 +44,7 @@ module CoreMonad (
     liftIO1, liftIO2, liftIO3, liftIO4,
     
     -- ** Global initialization
-    reinitializeGlobals,
+    reinitializeGlobals, bracketGlobals,
     
     -- ** Dealing with annotations
     getAnnotations, getFirstAnnotations,
@@ -947,6 +947,7 @@ domain of the GlobalRdrEnv is affected by the RdrName's OccName's FastString's
 unique.
 
 \begin{code}
+-- called by plugin
 reinitializeGlobals :: CoreM ()
 reinitializeGlobals = do
     (sf_globals, fs_table, linker_globals) <- read cr_globals
@@ -956,6 +957,16 @@ reinitializeGlobals = do
     liftIO $ restoreFSTable fs_table
     liftIO $ restoreLinkerGlobals linker_globals
     liftIO $ setUnsafeGlobalDynFlags dflags
+
+-- called by host compiler, assuming argument is code from a plugin
+bracketGlobals :: CoreM a -> CoreM a
+bracketGlobals (CoreM f) = do
+  tbl <- liftIO saveFSTable
+  let upd e = e {cr_globals=(x,tbl,z)}
+        where (x,_,z) = cr_globals e
+  x <- CoreM (\s -> updEnv upd (f s))
+  liftIO unsaveFSTable
+  return x
 \end{code}
 
 %************************************************************************

compiler/simplCore/SimplCore.lhs

@@ -311,7 +311,7 @@ addPluginPasses dflags builtin_passes
        ; foldM query_plug builtin_passes named_plugins }
   where
     query_plug todos (mod_nm, plug)
-       = installCoreToDos plug options todos
+       = bracketGlobals $ installCoreToDos plug options todos
        where
          options = [ option | (opt_mod_nm, option) <- pluginModNameOpts dflags
                             , opt_mod_nm == mod_nm ]
@@ -407,7 +407,7 @@ doCorePass _      CoreDoNothing                = return
 doCorePass _      (CoreDoPasses passes)        = runCorePasses passes
 
 #ifdef GHCI
-doCorePass _      (CoreDoPluginPass _ pass) = {-# SCC "Plugin" #-} pass
+doCorePass _      (CoreDoPluginPass _ pass) = {-# SCC "Plugin" #-} (bracketGlobals . pass)
 #endif
 
 doCorePass _      pass = pprPanic "doCorePass" (ppr pass)

compiler/utils/FastString.lhs

@@ -94,7 +94,7 @@ module FastString
         lengthLS,
 
         -- * Saving/restoring globals
-        saveFSTable, restoreFSTable, FastStringTable
+        saveFSTable, restoreFSTable, unsaveFSTable, FastStringTable
        ) where
 
 #include "HsVersions.h"
@@ -480,7 +480,7 @@ nilFS = mkFastString ""
 getFastStringTable :: IO [[FastString]]
 getFastStringTable = do
   tbl <- readIORef string_table
-  buckets <- mapM (lookupTbl tbl) [0 .. hASH_TBL_SIZE]
+  buckets <- mapM (lookupTbl tbl) [0 .. hASH_TBL_SIZE - 1]
   return buckets
 
 -- -----------------------------------------------------------------------------
@@ -581,9 +581,19 @@ fsLit x = mkFastString x
 --------------------
 -- for plugins; see Note [Initializing globals] in CoreMonad
 
+-- called by host compiler
 saveFSTable :: IO FastStringTable
 saveFSTable = readIORef string_table
 
+-- called by host compiler
+unsaveFSTable :: IO ()
+unsaveFSTable = do
+  tbl@(FastStringTable _ arr#) <- readIORef string_table
+  buckets <- mapM (lookupTbl tbl) [0 .. hASH_TBL_SIZE - 1]
+  let size = sum $ map length $ buckets
+  writeIORef string_table (FastStringTable size arr#)
+
+-- called by plugin
 restoreFSTable :: FastStringTable -> IO ()
 restoreFSTable = writeIORef string_table
 \end{code}