Run typed splices in the zonker

Summary:
This fixes #15471

In the typechecker we check that the splice has the right type but we
crucially don't zonk the generated expression. This is because we might
end up unifying type variables from outer scopes later on.

Reviewers: simonpj, goldfire, bgamari

Subscribers: rwbarton, carter

GHC Trac Issues: #15471

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

10 files changed, 183 insertions, 44 deletions

diff --git a/compiler/deSugar/DsMeta.hs b/compiler/deSugar/DsMeta.hs
index 9906fc729b..02b6cbc21e 100644
--- a/compiler/deSugar/DsMeta.hs
+++ b/compiler/deSugar/DsMeta.hs
@@ -1255,6 +1255,7 @@ repSplice (HsTypedSplice   _ _ n _) = rep_splice n
 repSplice (HsUntypedSplice _ _ n _) = rep_splice n
 repSplice (HsQuasiQuote _ n _ _ _)  = rep_splice n
 repSplice e@(HsSpliced {})          = pprPanic "repSplice" (ppr e)
+repSplice e@(HsSplicedT {})         = pprPanic "repSpliceT" (ppr e)
 repSplice e@(XSplice {})            = pprPanic "repSplice" (ppr e)
 
 rep_splice :: Name -> DsM (Core a)
diff --git a/compiler/hieFile/HieAst.hs b/compiler/hieFile/HieAst.hs
index eafafbbde9..432dc36069 100644
--- a/compiler/hieFile/HieAst.hs
+++ b/compiler/hieFile/HieAst.hs
@@ -1504,6 +1504,8 @@ instance ( ToHie (LHsExpr a)
         ]
       HsSpliced _ _ _ ->
         []
+      HsSplicedT _ ->
+        []
       XSplice _ -> []
 
 instance ToHie (LRoleAnnotDecl GhcRn) where
diff --git a/compiler/hsSyn/HsExpr.hs b/compiler/hsSyn/HsExpr.hs
index a7d12c2704..41bbdb2a62 100644
--- a/compiler/hsSyn/HsExpr.hs
+++ b/compiler/hsSyn/HsExpr.hs
@@ -43,6 +43,8 @@ import Util
 import Outputable
 import FastString
 import Type
+import TcType (TcType)
+import {-# SOURCE #-} TcRnTypes (TcLclEnv)
 
 -- libraries:
 import Data.Data hiding (Fixity(..))
@@ -2403,6 +2405,8 @@ data HsSplice id
         (XSpliced id)
         ThModFinalizers     -- TH finalizers produced by the splice.
         (HsSplicedThing id) -- The result of splicing
+   | HsSplicedT
+      DelayedSplice
    | XSplice (XXSplice id)  -- Note [Trees that Grow] extension point
 
 type instance XTypedSplice   (GhcPass _) = NoExt
@@ -2442,6 +2446,21 @@ instance Data ThModFinalizers where
   toConstr  a   = mkConstr (dataTypeOf a) "ThModFinalizers" [] Data.Prefix
   dataTypeOf a  = mkDataType "HsExpr.ThModFinalizers" [toConstr a]
 
+-- See Note [Running typed splices in the zonker]
+-- These are the arguments that are passed to `TcSplice.runTopSplice`
+data DelayedSplice =
+  DelayedSplice
+    TcLclEnv          -- The local environment to run the splice in
+    (LHsExpr GhcRn)   -- The original renamed expression
+    TcType            -- The result type of running the splice, unzonked
+    (LHsExpr GhcTcId) -- The typechecked expression to run and splice in the result
+
+-- A Data instance which ignores the argument of 'DelayedSplice'.
+instance Data DelayedSplice where
+  gunfold _ _ _ = panic "DelayedSplice"
+  toConstr  a   = mkConstr (dataTypeOf a) "DelayedSplice" [] Data.Prefix
+  dataTypeOf a  = mkDataType "HsExpr.DelayedSplice" [toConstr a]
+
 -- | Haskell Spliced Thing
 --
 -- Values that can result from running a splice.
@@ -2573,6 +2592,7 @@ pprSplice (HsUntypedSplice _ NoParens n e)
   = ppr_splice empty  n e empty
 pprSplice (HsQuasiQuote _ n q _ s)      = ppr_quasi n q s
 pprSplice (HsSpliced _ _ thing)         = ppr thing
+pprSplice (HsSplicedT {})               = text "Unevaluated typed splice"
 pprSplice (XSplice x)                   = ppr x
 
 ppr_quasi :: OutputableBndr p => p -> p -> FastString -> SDoc
diff --git a/compiler/rename/RnSplice.hs b/compiler/rename/RnSplice.hs
index 6adee1c735..a0c926d4e7 100644
--- a/compiler/rename/RnSplice.hs
+++ b/compiler/rename/RnSplice.hs
@@ -55,6 +55,8 @@ import {-# SOURCE #-} TcSplice
     , tcTopSpliceExpr
     )
 
+import TcHsSyn
+
 import GHCi.RemoteTypes ( ForeignRef )
 import qualified Language.Haskell.TH as TH (Q)
 
@@ -300,12 +302,14 @@ runRnSplice flavour run_meta ppr_res splice
                 HsQuasiQuote _ _ q qs str -> mkQuasiQuoteExpr flavour q qs str
                 HsTypedSplice {}          -> pprPanic "runRnSplice" (ppr splice)
                 HsSpliced {}              -> pprPanic "runRnSplice" (ppr splice)
+                HsSplicedT {}             -> pprPanic "runRnSplice" (ppr splice)
                 XSplice {}                -> pprPanic "runRnSplice" (ppr splice)
 
              -- Typecheck the expression
        ; meta_exp_ty   <- tcMetaTy meta_ty_name
-       ; zonked_q_expr <- tcTopSpliceExpr Untyped $
-                          tcPolyExpr the_expr meta_exp_ty
+       ; zonked_q_expr <- zonkTopLExpr =<<
+                            tcTopSpliceExpr Untyped
+                              (tcPolyExpr the_expr meta_exp_ty)
 
              -- Run the expression
        ; mod_finalizers_ref <- newTcRef []
@@ -346,6 +350,8 @@ makePending _ splice@(HsTypedSplice {})
   = pprPanic "makePending" (ppr splice)
 makePending _ splice@(HsSpliced {})
   = pprPanic "makePending" (ppr splice)
+makePending _ splice@(HsSplicedT {})
+  = pprPanic "makePending" (ppr splice)
 makePending _ splice@(XSplice {})
   = pprPanic "makePending" (ppr splice)
 
@@ -400,6 +406,7 @@ rnSplice (HsQuasiQuote x splice_name quoter q_loc quote)
                                                              , unitFV quoter') }
 
 rnSplice splice@(HsSpliced {}) = pprPanic "rnSplice" (ppr splice)
+rnSplice splice@(HsSplicedT {}) = pprPanic "rnSplice" (ppr splice)
 rnSplice splice@(XSplice {})   = pprPanic "rnSplice" (ppr splice)
 
 ---------------------
@@ -709,6 +716,7 @@ spliceCtxt splice
              HsTypedSplice   {} -> text "typed splice:"
              HsQuasiQuote    {} -> text "quasi-quotation:"
              HsSpliced       {} -> text "spliced expression:"
+             HsSplicedT      {} -> text "spliced expression:"
              XSplice         {} -> text "spliced expression:"
 
 -- | The splice data to be logged
diff --git a/compiler/typecheck/TcHsSyn.hs b/compiler/typecheck/TcHsSyn.hs
index 6834af9f41..462b9245e5 100644
--- a/compiler/typecheck/TcHsSyn.hs
+++ b/compiler/typecheck/TcHsSyn.hs
@@ -81,6 +81,8 @@ import Util
 import UniqFM
 import CoreSyn
 
+import {-# SOURCE #-} TcSplice (runTopSplice)
+
 import Control.Monad
 import Data.List  ( partition )
 import Control.Arrow ( second )
@@ -773,6 +775,9 @@ zonkExpr env (HsTcBracketOut x body bs)
     zonk_b (PendingTcSplice n e) = do e' <- zonkLExpr env e
                                       return (PendingTcSplice n e')
 
+zonkExpr env (HsSpliceE _ (HsSplicedT s)) =
+  runTopSplice s >>= zonkExpr env
+
 zonkExpr _ (HsSpliceE x s) = WARN( True, ppr s ) -- Should not happen
                            return (HsSpliceE x s)
 
diff --git a/compiler/typecheck/TcPatSyn.hs b/compiler/typecheck/TcPatSyn.hs
index ea1f483268..822697faf6 100644
--- a/compiler/typecheck/TcPatSyn.hs
+++ b/compiler/typecheck/TcPatSyn.hs
@@ -1028,6 +1028,7 @@ tcPatToExpr name args pat = go pat
     go1 (SplicePat _ (HsSpliced _ _ (HsSplicedPat pat)))
                                     = go1 pat
     go1 (SplicePat _ (HsSpliced{})) = panic "Invalid splice variety"
+    go1 (SplicePat _ (HsSplicedT{})) = panic "Invalid splice variety"
 
     -- The following patterns are not invertible.
     go1 p@(BangPat {})                       = notInvertible p -- #14112
diff --git a/compiler/typecheck/TcRnDriver.hs b/compiler/typecheck/TcRnDriver.hs
index fae0c19394..36ec8dcd2e 100644
--- a/compiler/typecheck/TcRnDriver.hs
+++ b/compiler/typecheck/TcRnDriver.hs
@@ -421,12 +421,6 @@ tcRnSrcDecls explicit_mod_hdr decls
         -- Emit Typeable bindings
       ; tcg_env <- mkTypeableBinds
 
-        -- Finalizers must run after constraints are simplified, or some types
-        -- might not be complete when using reify (see #12777).
-      ; (tcg_env, tcl_env) <- setGblEnv tcg_env run_th_modfinalizers
-      ; setEnvs (tcg_env, tcl_env) $ do {
-
-      ; finishTH
 
       ; traceTc "Tc9" empty
 
@@ -438,32 +432,63 @@ tcRnSrcDecls explicit_mod_hdr decls
         -- Zonk the final code.  This must be done last.
         -- Even simplifyTop may do some unification.
         -- This pass also warns about missing type signatures
-      ; let { TcGblEnv { tcg_type_env  = type_env,
-                         tcg_binds     = binds,
-                         tcg_ev_binds  = cur_ev_binds,
-                         tcg_imp_specs = imp_specs,
-                         tcg_rules     = rules,
-                         tcg_fords     = fords } = tcg_env
-            ; all_ev_binds = cur_ev_binds `unionBags` new_ev_binds } ;
-
       ; (bind_env, ev_binds', binds', fords', imp_specs', rules')
-            <- {-# SCC "zonkTopDecls" #-}
-               zonkTopDecls all_ev_binds binds rules
-                            imp_specs fords ;
+            <- zonkTcGblEnv new_ev_binds tcg_env
+
+        -- Finalizers must run after constraints are simplified, or some types
+        -- might not be complete when using reify (see #12777).
+        -- and also after we zonk the first time because we run typed splices
+        -- in the zonker which gives rise to the finalisers.
+      ; (tcg_env_mf, _) <- setGblEnv (clearTcGblEnv tcg_env)
+                                     run_th_modfinalizers
+      ; finishTH
       ; traceTc "Tc11" empty
 
-      ; let { final_type_env = plusTypeEnv type_env bind_env
-            ; tcg_env' = tcg_env { tcg_binds    = binds',
-                                   tcg_ev_binds = ev_binds',
-                                   tcg_imp_specs = imp_specs',
-                                   tcg_rules    = rules',
-                                   tcg_fords    = fords' } } ;
+      ; -- zonk the new bindings arising from running the finalisers.
+        -- This won't give rise to any more finalisers as you can't nest
+        -- finalisers inside finalisers.
+      ; (bind_env_mf, ev_binds_mf, binds_mf, fords_mf, imp_specs_mf, rules_mf)
+            <- zonkTcGblEnv emptyBag tcg_env_mf
+
+
+      ; let { final_type_env = plusTypeEnv (tcg_type_env tcg_env)
+                                (plusTypeEnv bind_env_mf bind_env)
+            ; tcg_env' = tcg_env_mf
+                          { tcg_binds    = binds' `unionBags` binds_mf,
+                            tcg_ev_binds = ev_binds' `unionBags` ev_binds_mf ,
+                            tcg_imp_specs = imp_specs' ++ imp_specs_mf ,
+                            tcg_rules    = rules' ++ rules_mf ,
+                            tcg_fords    = fords' ++ fords_mf } } ;
 
       ; setGlobalTypeEnv tcg_env' final_type_env
 
-   }
    } }
 
+zonkTcGblEnv :: Bag EvBind -> TcGblEnv
+             -> TcM (TypeEnv, Bag EvBind, LHsBinds GhcTc,
+                       [LForeignDecl GhcTc], [LTcSpecPrag], [LRuleDecl GhcTc])
+zonkTcGblEnv new_ev_binds tcg_env =
+  let TcGblEnv {   tcg_binds     = binds,
+                   tcg_ev_binds  = cur_ev_binds,
+                   tcg_imp_specs = imp_specs,
+                   tcg_rules     = rules,
+                   tcg_fords     = fords } = tcg_env
+
+      all_ev_binds = cur_ev_binds `unionBags` new_ev_binds
+
+  in {-# SCC "zonkTopDecls" #-}
+      zonkTopDecls all_ev_binds binds rules imp_specs fords
+
+
+-- | Remove accumulated bindings, rules and so on from TcGblEnv
+clearTcGblEnv :: TcGblEnv -> TcGblEnv
+clearTcGblEnv tcg_env
+  = tcg_env { tcg_binds    = emptyBag,
+              tcg_ev_binds = emptyBag ,
+              tcg_imp_specs = [],
+              tcg_rules    = [],
+              tcg_fords    = [] }
+
 -- | Runs TH finalizers and renames and typechecks the top-level declarations
 -- that they could introduce.
 run_th_modfinalizers :: TcM (TcGblEnv, TcLclEnv)
diff --git a/compiler/typecheck/TcRnTypes.hs-boot b/compiler/typecheck/TcRnTypes.hs-boot
new file mode 100644
index 0000000000..b22f5c3ff8
--- /dev/null
+++ b/compiler/typecheck/TcRnTypes.hs-boot
@@ -0,0 +1,6 @@
+module TcRnTypes where
+
+-- Build ordering
+import GHC.Base()
+
+data TcLclEnv
diff --git a/compiler/typecheck/TcSplice.hs b/compiler/typecheck/TcSplice.hs
index 53df2bb03d..1bb844a77f 100644
--- a/compiler/typecheck/TcSplice.hs
+++ b/compiler/typecheck/TcSplice.hs
@@ -26,7 +26,7 @@ module TcSplice(
      runMetaE, runMetaP, runMetaT, runMetaD, runQuasi,
      tcTopSpliceExpr, lookupThName_maybe,
      defaultRunMeta, runMeta', runRemoteModFinalizers,
-     finishTH
+     finishTH, runTopSplice
       ) where
 
 #include "HsVersions.h"
@@ -58,7 +58,7 @@ import HscMain
         -- These imports are the reason that TcSplice
         -- is very high up the module hierarchy
 import FV
-import RnSplice( traceSplice, SpliceInfo(..) )
+import RnSplice( traceSplice, SpliceInfo(..))
 import RdrName
 import HscTypes
 import Convert
@@ -491,28 +491,44 @@ tcTopSplice expr res_ty
          -- making sure it has type Q (T res_ty)
          res_ty <- expTypeToType res_ty
        ; meta_exp_ty <- tcTExpTy res_ty
-       ; zonked_q_expr <- tcTopSpliceExpr Typed $
+       ; q_expr <- tcTopSpliceExpr Typed $
                           tcMonoExpr expr (mkCheckExpType meta_exp_ty)
+       ; lcl_env <- getLclEnv
+       ; let delayed_splice
+              = DelayedSplice lcl_env expr res_ty q_expr
+       ; return (HsSpliceE noExt (HsSplicedT delayed_splice))
 
-         -- See Note [Collecting modFinalizers in typed splices].
+       }
+
+
+-- This is called in the zonker
+-- See Note [Running typed splices in the zonker]
+runTopSplice :: DelayedSplice -> TcM (HsExpr GhcTc)
+runTopSplice (DelayedSplice lcl_env orig_expr res_ty q_expr)
+  = setLclEnv lcl_env $ do {
+         zonked_ty <- zonkTcType res_ty
+       ; zonked_q_expr <- zonkTopLExpr q_expr
+        -- See Note [Collecting modFinalizers in typed splices].
        ; modfinalizers_ref <- newTcRef []
          -- Run the expression
        ; expr2 <- setStage (RunSplice modfinalizers_ref) $
                     runMetaE zonked_q_expr
        ; mod_finalizers <- readTcRef modfinalizers_ref
        ; addModFinalizersWithLclEnv $ ThModFinalizers mod_finalizers
+       -- We use orig_expr here and not q_expr when tracing as a call to
+       -- unsafeTExpCoerce is added to the original expression by the
+       -- typechecker when typed quotes are type checked.
        ; traceSplice (SpliceInfo { spliceDescription = "expression"
                                  , spliceIsDecl      = False
-                                 , spliceSource      = Just expr
+                                 , spliceSource      = Just orig_expr
                                  , spliceGenerated   = ppr expr2 })
+        -- Rename and typecheck the spliced-in expression,
+        -- making sure it has type res_ty
+        -- These steps should never fail; this is a *typed* splice
+       ; addErrCtxt (spliceResultDoc zonked_q_expr) $ do
+         { (exp3, _fvs) <- rnLExpr expr2
+         ; unLoc <$> tcMonoExpr exp3 (mkCheckExpType zonked_ty)} }
 
-         -- Rename and typecheck the spliced-in expression,
-         -- making sure it has type res_ty
-         -- These steps should never fail; this is a *typed* splice
-       ; addErrCtxt (spliceResultDoc expr) $ do
-       { (exp3, _fvs) <- rnLExpr expr2
-       ; exp4 <- tcMonoExpr exp3 (mkCheckExpType res_ty)
-       ; return (unLoc exp4) } }
 
 {-
 ************************************************************************
@@ -527,7 +543,7 @@ spliceCtxtDoc splice
   = hang (text "In the Template Haskell splice")
          2 (pprSplice splice)
 
-spliceResultDoc :: LHsExpr GhcRn -> SDoc
+spliceResultDoc :: LHsExpr GhcTc -> SDoc
 spliceResultDoc expr
   = sep [ text "In the result of the splice:"
         , nest 2 (char '$' <> ppr expr)
@@ -559,7 +575,7 @@ tcTopSpliceExpr isTypedSplice tc_action
        ; const_binds     <- simplifyTop wanted
 
           -- Zonk it and tie the knot of dictionary bindings
-       ; zonkTopLExpr (mkHsDictLet (EvBinds const_binds) expr') }
+       ; return $ mkHsDictLet (EvBinds const_binds) expr' }
 
 {-
 ************************************************************************
@@ -578,7 +594,7 @@ runAnnotation target expr = do
     -- Check the instances we require live in another module (we want to execute it..)
     -- and check identifiers live in other modules using TH stage checks. tcSimplifyStagedExpr
     -- also resolves the LIE constraints to detect e.g. instance ambiguity
-    zonked_wrapped_expr' <- tcTopSpliceExpr Untyped $
+    zonked_wrapped_expr' <- zonkTopLExpr =<< tcTopSpliceExpr Untyped (
            do { (expr', expr_ty) <- tcInferRhoNC expr
                 -- We manually wrap the typechecked expression in a call to toAnnotationWrapper
                 -- By instantiating the call >here< it gets registered in the
@@ -589,7 +605,8 @@ runAnnotation target expr = do
                       = L loc (mkHsWrap wrapper
                                  (HsVar noExt (L loc to_annotation_wrapper_id)))
               ; return (L loc (HsApp noExt
-                                specialised_to_annotation_wrapper_expr expr')) }
+                                specialised_to_annotation_wrapper_expr expr'))
+                                })
 
     -- Run the appropriately wrapped expression to get the value of
     -- the annotation and its dictionaries. The return value is of
@@ -790,6 +807,58 @@ runMeta' show_code ppr_hs run_and_convert expr
         failWithTc msg
 
 {-
+Note [Running typed splices in the zonker]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+See #15471 for the full discussion.
+
+For many years typed splices were run immediately after they were type checked
+however, this is too early as it means to zonk some type variables before
+they can be unified with type variables in the surrounding context.
+
+For example,
+
+```
+module A where
+
+test_foo :: forall a . Q (TExp (a -> a))
+test_foo = [|| id ||]
+
+module B where
+
+import A
+
+qux = $$(test_foo)
+```
+
+We would expect `qux` to have inferred type `forall a . a -> a` but if
+we run the splices too early the unified variables are zonked to `Any`. The
+inferred type is the unusable `Any -> Any`.
+
+To run the splice, we must compile `test_foo` all the way to byte code.
+But at the moment when the type checker is looking at the splice, test_foo
+has type `Q (TExp (alpha -> alpha))` and we
+certainly can't compile code involving unification variables!
+
+We could default `alpha` to `Any` but then we infer `qux :: Any -> Any`
+which definitely is not what we want.  Moreover, if we had
+  qux = [$$(test_foo), (\x -> x +1::Int)]
+then `alpha` would have to be `Int`.
+
+Conclusion: we must defer taking decisions about `alpha` until the
+typechecker is done; and *then* we can run the splice.  It's fine to do it
+later, because we know it'll produce type-correct code.
+
+Deferring running the splice until later, in the zonker, means that the
+unification variables propagate upwards from the splice into the surrounding
+context and are unified correctly.
+
+This is implemented by storing the arguments we need for running the splice
+in a `DelayedSplice`. In the zonker, the arguments are passed to
+`TcSplice.runTopSplice` and the expression inserted into the AST as normal.
+
+
+
 Note [Exceptions in TH]
 ~~~~~~~~~~~~~~~~~~~~~~~
 Suppose we have something like this
diff --git a/compiler/typecheck/TcSplice.hs-boot b/compiler/typecheck/TcSplice.hs-boot
index be2c67d887..8fb294bfc6 100644
--- a/compiler/typecheck/TcSplice.hs-boot
+++ b/compiler/typecheck/TcSplice.hs-boot
@@ -5,11 +5,11 @@ module TcSplice where
 
 import GhcPrelude
 import Name
-import HsExpr   ( PendingRnSplice )
+import HsExpr   ( PendingRnSplice, DelayedSplice )
 import TcRnTypes( TcM , SpliceType )
 import TcType   ( ExpRhoType )
 import Annotations ( Annotation, CoreAnnTarget )
-import HsExtension ( GhcTcId, GhcRn, GhcPs )
+import HsExtension ( GhcTcId, GhcRn, GhcPs, GhcTc )
 
 import HsSyn      ( HsSplice, HsBracket, HsExpr, LHsExpr, LHsType, LPat,
                     LHsDecl, ThModFinalizers )
@@ -29,6 +29,8 @@ tcTypedBracket :: HsExpr GhcRn
                -> ExpRhoType
                -> TcM (HsExpr GhcTcId)
 
+runTopSplice :: DelayedSplice -> TcM (HsExpr GhcTc)
+
 runAnnotation     :: CoreAnnTarget -> LHsExpr GhcRn -> TcM Annotation
 
 tcTopSpliceExpr :: SpliceType -> TcM (LHsExpr GhcTcId) -> TcM (LHsExpr GhcTcId)