Implement addCStub in template-haskell.

Summary:
addCStub allows injecting C code in the current module to be included
in the final object file.

Test Plan: ./validate

Reviewers: simonpj, goldfire, austin, bgamari

Reviewed By: bgamari

Subscribers: bitonic, duncan, mboes, thomie

Differential Revision: https://phabricator.haskell.org/D3106

12 files changed, 105 insertions, 2 deletions

diff --git a/compiler/deSugar/Desugar.hs b/compiler/deSugar/Desugar.hs
index 1cd7979128..5111141770 100644
--- a/compiler/deSugar/Desugar.hs
+++ b/compiler/deSugar/Desugar.hs
@@ -289,6 +289,7 @@ deSugar hsc_env
                             tcg_imp_specs    = imp_specs,
                             tcg_dependent_files = dependent_files,
                             tcg_ev_binds     = ev_binds,
+                            tcg_th_cstubs    = th_cstubs_var,
                             tcg_fords        = fords,
                             tcg_rules        = rules,
                             tcg_vects        = vects,
@@ -373,6 +374,9 @@ deSugar hsc_env
         -- past desugaring. See Note [Identity versus semantic module].
         ; MASSERT( id_mod == mod )
 
+        ; cstubs <- readIORef th_cstubs_var
+        ; let ds_fords' = foldl' appendStubC ds_fords (map text cstubs)
+
         ; let mod_guts = ModGuts {
                 mg_module       = mod,
                 mg_hsc_src      = hsc_src,
@@ -393,7 +397,7 @@ deSugar hsc_env
                 mg_patsyns      = patsyns,
                 mg_rules        = ds_rules_for_imps,
                 mg_binds        = ds_binds,
-                mg_foreign      = ds_fords,
+                mg_foreign      = ds_fords',
                 mg_hpc_info     = ds_hpc_info,
                 mg_modBreaks    = modBreaks,
                 mg_vect_decls   = ds_vects,
diff --git a/compiler/typecheck/TcRnMonad.hs b/compiler/typecheck/TcRnMonad.hs
index 33cb4d14d0..1c84b4014e 100644
--- a/compiler/typecheck/TcRnMonad.hs
+++ b/compiler/typecheck/TcRnMonad.hs
@@ -217,6 +217,7 @@ initTc hsc_env hsc_src keep_rn_syntax mod loc do_this
         dependent_files_var <- newIORef [] ;
         static_wc_var       <- newIORef emptyWC ;
         th_topdecls_var      <- newIORef [] ;
+        th_cstubs_var        <- newIORef [] ;
         th_topnames_var      <- newIORef emptyNameSet ;
         th_modfinalizers_var <- newIORef [] ;
         th_state_var         <- newIORef Map.empty ;
@@ -231,6 +232,7 @@ initTc hsc_env hsc_src keep_rn_syntax mod loc do_this
 
              gbl_env = TcGblEnv {
                 tcg_th_topdecls      = th_topdecls_var,
+                tcg_th_cstubs        = th_cstubs_var,
                 tcg_th_topnames      = th_topnames_var,
                 tcg_th_modfinalizers = th_modfinalizers_var,
                 tcg_th_state         = th_state_var,
diff --git a/compiler/typecheck/TcRnTypes.hs b/compiler/typecheck/TcRnTypes.hs
index 782a99297b..766119635f 100644
--- a/compiler/typecheck/TcRnTypes.hs
+++ b/compiler/typecheck/TcRnTypes.hs
@@ -608,6 +608,9 @@ data TcGblEnv
         tcg_th_topdecls :: TcRef [LHsDecl RdrName],
         -- ^ Top-level declarations from addTopDecls
 
+        tcg_th_cstubs :: TcRef [String],
+        -- ^ C stubs from addCStub
+
         tcg_th_topnames :: TcRef NameSet,
         -- ^ Exact names bound in top-level declarations in tcg_th_topdecls
 
diff --git a/compiler/typecheck/TcSplice.hs b/compiler/typecheck/TcSplice.hs
index e142cae3b2..aba70aaaa1 100644
--- a/compiler/typecheck/TcSplice.hs
+++ b/compiler/typecheck/TcSplice.hs
@@ -909,6 +909,17 @@ instance TH.Quasi TcM where
           hang (text "The binder" <+> quotes (ppr name) <+> ptext (sLit "is not a NameU."))
              2 (text "Probable cause: you used mkName instead of newName to generate a binding.")
 
+  qAddCStub str = do
+      l <- getSrcSpanM
+      r <- case l of
+             UnhelpfulSpan _ -> pprPanic "qAddCStub: Unhelpful location" (ppr l)
+             RealSrcSpan s -> return s
+      let filename  = unpackFS (srcSpanFile r)
+          linePragma = "#line " ++ show (srcSpanStartLine r)
+                                ++ " " ++ show filename
+      th_cstubs_var <- fmap tcg_th_cstubs getGblEnv
+      updTcRef th_cstubs_var ([linePragma, str] ++)
+
   qAddModFinalizer fin = do
       r <- liftIO $ mkRemoteRef fin
       fref <- liftIO $ mkForeignRef r (freeRemoteRef r)
@@ -1100,6 +1111,7 @@ handleTHMessage msg = case msg of
     hsc_env <- env_top <$> getEnv
     wrapTHResult $ liftIO (mkFinalizedHValue hsc_env r) >>= addModFinalizerRef
   AddTopDecls decs -> wrapTHResult $ TH.qAddTopDecls decs
+  AddCStub str -> wrapTHResult $ TH.qAddCStub str
   IsExtEnabled ext -> wrapTHResult $ TH.qIsExtEnabled ext
   ExtsEnabled -> wrapTHResult $ TH.qExtsEnabled
   _ -> panic ("handleTHMessage: unexpected message " ++ show msg)
diff --git a/libraries/ghci/GHCi/Message.hs b/libraries/ghci/GHCi/Message.hs
index c336349daf..71da2287bb 100644
--- a/libraries/ghci/GHCi/Message.hs
+++ b/libraries/ghci/GHCi/Message.hs
@@ -237,6 +237,7 @@ data THMessage a where
   AddDependentFile :: FilePath -> THMessage (THResult ())
   AddModFinalizer :: RemoteRef (TH.Q ()) -> THMessage (THResult ())
   AddTopDecls :: [TH.Dec] -> THMessage (THResult ())
+  AddCStub :: String -> THMessage (THResult ())
   IsExtEnabled :: Extension -> THMessage (THResult Bool)
   ExtsEnabled :: THMessage (THResult [Extension])
 
@@ -272,7 +273,8 @@ getTHMessage = do
     14 -> THMsg <$> return StartRecover
     15 -> THMsg <$> EndRecover <$> get
     16 -> return (THMsg RunTHDone)
-    _  -> THMsg <$> AddModFinalizer <$> get
+    17 -> THMsg <$> AddModFinalizer <$> get
+    _  -> THMsg <$> AddCStub <$> get
 
 putTHMessage :: THMessage a -> Put
 putTHMessage m = case m of
@@ -294,6 +296,7 @@ putTHMessage m = case m of
   EndRecover a                -> putWord8 15 >> put a
   RunTHDone                   -> putWord8 16
   AddModFinalizer a           -> putWord8 17 >> put a
+  AddCStub a                  -> putWord8 18 >> put a
 
 
 data EvalOpts = EvalOpts
diff --git a/libraries/ghci/GHCi/TH.hs b/libraries/ghci/GHCi/TH.hs
index def6aee33e..8cb9accc5e 100644
--- a/libraries/ghci/GHCi/TH.hs
+++ b/libraries/ghci/GHCi/TH.hs
@@ -193,6 +193,7 @@ instance TH.Quasi GHCiQ where
   qRunIO m = GHCiQ $ \s -> fmap (,s) m
   qAddDependentFile file = ghcCmd (AddDependentFile file)
   qAddTopDecls decls = ghcCmd (AddTopDecls decls)
+  qAddCStub str = ghcCmd (AddCStub str)
   qAddModFinalizer fin = GHCiQ (\s -> mkRemoteRef fin >>= return . (, s)) >>=
                          ghcCmd . AddModFinalizer
   qGetQ = GHCiQ $ \s ->
diff --git a/libraries/template-haskell/Language/Haskell/TH/Syntax.hs b/libraries/template-haskell/Language/Haskell/TH/Syntax.hs
index 815e3fc5f3..c531eeffd7 100644
--- a/libraries/template-haskell/Language/Haskell/TH/Syntax.hs
+++ b/libraries/template-haskell/Language/Haskell/TH/Syntax.hs
@@ -92,6 +92,8 @@ class Monad m => Quasi m where
 
   qAddTopDecls :: [Dec] -> m ()
 
+  qAddCStub :: String -> m ()
+
   qAddModFinalizer :: Q () -> m ()
 
   qGetQ :: Typeable a => m (Maybe a)
@@ -131,6 +133,7 @@ instance Quasi IO where
   qRecover _ _          = badIO "recover" -- Maybe we could fix this?
   qAddDependentFile _   = badIO "addDependentFile"
   qAddTopDecls _        = badIO "addTopDecls"
+  qAddCStub    _        = badIO "addCStub"
   qAddModFinalizer _    = badIO "addModFinalizer"
   qGetQ                 = badIO "getQ"
   qPutQ _               = badIO "putQ"
@@ -456,6 +459,25 @@ addDependentFile fp = Q (qAddDependentFile fp)
 addTopDecls :: [Dec] -> Q ()
 addTopDecls ds = Q (qAddTopDecls ds)
 
+-- | Add an additional C stub. The added stub will be built and included in the
+-- object file of the current module.
+--
+-- Compilation errors in the given string are reported next to the line of the
+-- enclosing splice.
+--
+-- The accuracy of the error location can be improved by adding
+-- #line pragmas in the argument. e.g.
+--
+-- > {-# LANGUAGE CPP #-}
+-- > ...
+-- > addCStub $ unlines
+-- >   [ "#line " ++ show (__LINE__ + 1) ++ " " ++ show __FILE__
+-- >   , ...
+-- >   ]
+--
+addCStub :: String -> Q ()
+addCStub str = Q (qAddCStub str)
+
 -- | Add a finalizer that will run in the Q monad after the current module has
 -- been type checked. This only makes sense when run within a top-level splice.
 --
@@ -499,6 +521,7 @@ instance Quasi Q where
   qRunIO              = runIO
   qAddDependentFile   = addDependentFile
   qAddTopDecls        = addTopDecls
+  qAddCStub           = addCStub
   qAddModFinalizer    = addModFinalizer
   qGetQ               = getQ
   qPutQ               = putQ
diff --git a/testsuite/tests/th/TH_addCStub1.hs b/testsuite/tests/th/TH_addCStub1.hs
new file mode 100644
index 0000000000..3a2c5c3609
--- /dev/null
+++ b/testsuite/tests/th/TH_addCStub1.hs
@@ -0,0 +1,22 @@
+-- Tests that addCStub includes the C code in the final object file and that
+-- -optc options are passed when building it.
+
+{-# LANGUAGE ForeignFunctionInterface #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# OPTIONS_GHC -optc-DA_MACRO=1 #-}
+
+import Language.Haskell.TH.Syntax
+
+foreign import ccall f :: Int -> IO Int
+
+do addCStub $ unlines
+     [ "#include <stdio.h>"
+     , "int f(int x) {"
+     , "  printf(\"calling f(%d)\\n\",x);"
+     , "  return A_MACRO + x;"
+     , "}"
+     ]
+   return []
+
+main :: IO ()
+main = f 2 >>= print
diff --git a/testsuite/tests/th/TH_addCStub1.stdout b/testsuite/tests/th/TH_addCStub1.stdout
new file mode 100644
index 0000000000..e46825eb2b
--- /dev/null
+++ b/testsuite/tests/th/TH_addCStub1.stdout
@@ -0,0 +1,2 @@
+3
+calling f(2)
diff --git a/testsuite/tests/th/TH_addCStub2.hs b/testsuite/tests/th/TH_addCStub2.hs
new file mode 100644
index 0000000000..10119d9370
--- /dev/null
+++ b/testsuite/tests/th/TH_addCStub2.hs
@@ -0,0 +1,22 @@
+-- Tests that a reasonable error is reported when addCStub is used with
+-- incorrect C code.
+
+{-# LANGUAGE ForeignFunctionInterface #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# OPTIONS_GHC -optc-DA_MACRO=1 #-}
+
+import Language.Haskell.TH.Syntax
+
+foreign import ccall f :: Int -> IO Int
+
+do addCStub $ unlines
+     [ "#include <stdio.h>"
+     , "int f(int x {"
+     , "  printf(\"calling f(%d)\\n\",x);"
+     , "  return A_MACRO + x;"
+     , "}"
+     ]
+   return []
+
+main :: IO ()
+main = f 2 >>= print
diff --git a/testsuite/tests/th/TH_addCStub2.stderr b/testsuite/tests/th/TH_addCStub2.stderr
new file mode 100644
index 0000000000..ba3277b277
--- /dev/null
+++ b/testsuite/tests/th/TH_addCStub2.stderr
@@ -0,0 +1,6 @@
+
+TH_addCStub2.hs:13:13:
+     expected ‘;’, ‘,’ or ‘)’ before ‘{’ token
+          [ "#include <stdio.h>"
+                 ^
+`gcc' failed in phase `C Compiler'. (Exit code: 1)
diff --git a/testsuite/tests/th/all.T b/testsuite/tests/th/all.T
index f05a634301..9a08b6542c 100644
--- a/testsuite/tests/th/all.T
+++ b/testsuite/tests/th/all.T
@@ -63,6 +63,9 @@ test('TH_reifyDecl2', normal, compile, ['-v0'])
 test('TH_reifyLocalDefs', normal, compile, ['-v0'])
 test('TH_reifyLocalDefs2', normal, compile, ['-v0'])
 
+test('TH_addCStub1', normal, compile_and_run, ['-v0'])
+test('TH_addCStub2', normal, compile_fail, ['-v0'])
+
 test('TH_reifyMkName', normal, compile, ['-v0'])
 
 test('TH_reifyInstances', normal, compile, ['-v0'])