Track the typed/untyped distinction in the current TH stage.

Also check for illegal typed/untyped bracket/splice combinations.

5 files changed, 77 insertions, 36 deletions

diff --git a/compiler/hsSyn/HsExpr.lhs b/compiler/hsSyn/HsExpr.lhs
index fc862713a4..36d6ceec80 100644
--- a/compiler/hsSyn/HsExpr.lhs
+++ b/compiler/hsSyn/HsExpr.lhs
@@ -1213,6 +1213,10 @@ data HsBracket id = ExpBr (LHsExpr id)   -- [|  expr  |]
                   | TExpBr (LHsExpr id)  -- [||  expr  ||]
   deriving (Data, Typeable)
 
+isTypedBracket :: HsBracket id -> Bool
+isTypedBracket (TExpBr {}) = True
+isTypedBracket _           = False
+
 instance OutputableBndr id => Outputable (HsBracket id) where
   ppr = pprHsBracket
 
diff --git a/compiler/typecheck/TcEnv.lhs b/compiler/typecheck/TcEnv.lhs
index fce17affaa..4f3a15d0b8 100644
--- a/compiler/typecheck/TcEnv.lhs
+++ b/compiler/typecheck/TcEnv.lhs
@@ -558,7 +558,7 @@ thRnBrack :: ThStage
 -- Used *only* to indicate that we are inside a TH bracket during renaming
 -- Tested by TcEnv.isBrackStage
 -- See Note [Top-level Names in Template Haskell decl quotes]
-thRnBrack = Brack (panic "thRnBrack1") (panic "thRnBrack2") (panic "thRnBrack3") 
+thRnBrack = Brack False (panic "thRnBrack1") (panic "thRnBrack2") (panic "thRnBrack3") 
 
 isBrackStage :: ThStage -> Bool
 isBrackStage (Brack {}) = True
@@ -762,7 +762,7 @@ notFound name
   = do { lcl_env <- getLclEnv
        ; let stage = tcl_th_ctxt lcl_env
        ; case stage of   -- See Note [Out of scope might be a staging error]
-           Splice -> stageRestrictionError (quotes (ppr name))
+           Splice {} -> stageRestrictionError (quotes (ppr name))
            _ -> failWithTc $
                 vcat[ptext (sLit "GHC internal error:") <+> quotes (ppr name) <+> 
                      ptext (sLit "is not in scope during type checking, but it passed the renamer"),
diff --git a/compiler/typecheck/TcExpr.lhs b/compiler/typecheck/TcExpr.lhs
index 30185fcab5..41e9dc2a28 100644
--- a/compiler/typecheck/TcExpr.lhs
+++ b/compiler/typecheck/TcExpr.lhs
@@ -1235,10 +1235,10 @@ checkCrossStageLifting :: Id -> ThLevel -> ThStage -> TcM ()
 -- Examples   \x -> [| x |]
 --            [| map |]
 
-checkCrossStageLifting _ _ Comp   = return ()
-checkCrossStageLifting _ _ Splice = return ()
+checkCrossStageLifting _ _ Comp      = return ()
+checkCrossStageLifting _ _ (Splice _) = return ()
 
-checkCrossStageLifting id _ (Brack _ ps_var lie_var)
+checkCrossStageLifting id _ (Brack _ _ ps_var lie_var)
   | thTopLevelId id
   =     -- Top-level identifiers in this module,
         -- (which have External Names)
diff --git a/compiler/typecheck/TcRnTypes.lhs b/compiler/typecheck/TcRnTypes.lhs
index 0bf05cffc4..9df3e48286 100644
--- a/compiler/typecheck/TcRnTypes.lhs
+++ b/compiler/typecheck/TcRnTypes.lhs
@@ -470,24 +470,26 @@ data ThStage    -- See Note [Template Haskell state diagram] in TcSplice
                 -- This code will be run *at compile time*;
                 --   the result replaces the splice
                 -- Binding level = 0
+      Bool      -- True if in a typed splice, False otherwise
 
   | Comp        -- Ordinary Haskell code
                 -- Binding level = 1
 
   | Brack                       -- Inside brackets
+      Bool                      --   True if inside a typed bracket, False otherwise
       ThStage                   --   Binding level = level(stage) + 1
       (TcRef [PendingSplice])   --   Accumulate pending splices here
       (TcRef WantedConstraints) --     and type constraints here
 
 topStage, topAnnStage, topSpliceStage :: ThStage
 topStage       = Comp
-topAnnStage    = Splice
-topSpliceStage = Splice
+topAnnStage    = Splice False
+topSpliceStage = Splice False
 
 instance Outputable ThStage where
-   ppr Splice        = text "Splice"
-   ppr Comp          = text "Comp"
-   ppr (Brack s _ _) = text "Brack" <> parens (ppr s)
+   ppr (Splice _)      = text "Splice"
+   ppr Comp            = text "Comp"
+   ppr (Brack _ s _ _) = text "Brack" <> parens (ppr s)
 
 type ThLevel = Int
         -- See Note [Template Haskell levels] in TcSplice
@@ -508,9 +510,9 @@ outerLevel = 1  -- Things defined outside brackets
 --      g2 = $(f ...)           is not OK; because we havn't compiled f yet
 
 thLevel :: ThStage -> ThLevel
-thLevel Splice        = 0
-thLevel Comp          = 1
-thLevel (Brack s _ _) = thLevel s + 1
+thLevel (Splice _)      = 0
+thLevel Comp            = 1
+thLevel (Brack _ s _ _) = thLevel s + 1
 
 ---------------------------
 -- Arrow-notation context
diff --git a/compiler/typecheck/TcSplice.lhs b/compiler/typecheck/TcSplice.lhs
index 535ef3080f..ac7aa7c843 100644
--- a/compiler/typecheck/TcSplice.lhs
+++ b/compiler/typecheck/TcSplice.lhs
@@ -328,7 +328,12 @@ tcBracket brack res_ty
                    2 (ppr brack)) $
     do {        -- Check for nested brackets
          cur_stage <- getStage
-       ; checkTc (not (isBrackStage cur_stage)) illegalBracket
+       ; case cur_stage of
+           { Splice True  -> checkTc (isTypedBracket brack) illegalUntypedBracket
+           ; Splice False -> checkTc (not (isTypedBracket brack)) illegalTypedBracket
+           ; Comp         -> return ()
+           ; Brack {}     -> failWithTc illegalBracket
+           }
 
         -- Brackets are desugared to code that mentions the TH package
        ; recordThUse
@@ -338,7 +343,7 @@ tcBracket brack res_ty
         -- it again when we actually use it.
        ; pending_splices <- newMutVar []
        ; lie_var <- getConstraintVar
-       ; let brack_stage = Brack cur_stage pending_splices lie_var
+       ; let brack_stage = Brack (isTypedBracket brack) cur_stage pending_splices lie_var
 
           -- We want to check that there aren't any constraints that
           -- can't be satisfied (e.g. Show Foo, where Foo has no Show
@@ -419,8 +424,13 @@ tc_bracket _ (DecBrL _)
 tc_bracket _ (TExpBr expr)
   = do  { any_ty <- newFlexiTyVarTy openTypeKind
         ; _ <- tcMonoExprNC expr any_ty  -- NC for no context; tcBracket does that
-        ; tcMetaTy expQTyConName }
-        -- Result type is ExpQ (= Q Exp)
+        ; tcTExpTy any_ty }
+        -- Result type is TExp tau
+
+tcTExpTy :: TcType -> TcM TcType
+tcTExpTy tau = do
+    t <- tcLookupTyCon tExpTyConName
+    return (mkTyConApp t [tau])
 
 quotedNameStageErr :: HsBracket Name -> SDoc
 quotedNameStageErr br
@@ -436,14 +446,14 @@ quotedNameStageErr br
 %************************************************************************
 
 \begin{code}
-tcSpliceExpr (HsSplice _ name expr) res_ty
+tcSpliceExpr (HsSplice isTypedSplice name expr) res_ty
   = setSrcSpan (getLoc expr)    $ do
     { stage <- getStage
     ; case stage of {
-        Splice -> tcTopSplice expr res_ty ;
-        Comp   -> tcTopSplice expr res_ty ;
+        Splice {} -> tcTopSplice isTypedSplice expr res_ty ;
+        Comp      -> tcTopSplice isTypedSplice expr res_ty ;
 
-        Brack pop_stage ps_var lie_var -> do
+        Brack isTypedBrack pop_stage ps_var lie_var -> do
 
         -- See Note [How brackets and nested splices are handled]
         -- A splice inside brackets
@@ -452,7 +462,16 @@ tcSpliceExpr (HsSplice _ name expr) res_ty
         -- Here (h 4) :: Q Exp
         -- but $(h 4) :: forall a.a     i.e. anything!
 
-     { meta_exp_ty <- tcMetaTy expQTyConName
+     { when (isTypedBrack && not isTypedSplice) $
+           failWithTc illegalUntypedSplice
+     ; when (not isTypedBrack && isTypedSplice) $
+           failWithTc illegalTypedSplice
+     ; meta_exp_ty <- if isTypedSplice
+                      then do { any_ty <- newFlexiTyVarTy openTypeKind
+                              ; tcTExpTy any_ty
+                              }
+                      else tcMetaTy expQTyConName
+
      ; expr' <- setStage pop_stage $
                 setConstraintVar lie_var    $
                 tcMonoExpr expr meta_exp_ty
@@ -464,13 +483,17 @@ tcSpliceExpr (HsSplice _ name expr) res_ty
      ; return (panic "tcSpliceExpr")    -- The returned expression is ignored
      }}}
 
-tcTopSplice :: LHsExpr Name -> TcRhoType -> TcM (HsExpr Id)
+tcTopSplice :: Bool -> LHsExpr Name -> TcRhoType -> TcM (HsExpr Id)
 -- Note [How top-level splices are handled]
-tcTopSplice expr res_ty
-  = do { meta_exp_ty <- tcMetaTy expQTyConName
+tcTopSplice isTypedSplice expr res_ty
+  = do { meta_exp_ty <- if isTypedSplice
+                        then do { any_ty <- newFlexiTyVarTy openTypeKind
+                                ; tcTExpTy any_ty
+                                }
+                        else tcMetaTy expQTyConName
 
         -- Typecheck the expression
-       ; zonked_q_expr <- tcTopSpliceExpr (tcMonoExpr expr meta_exp_ty)
+       ; zonked_q_expr <- tcTopSpliceExpr isTypedSplice (tcMonoExpr expr meta_exp_ty)
 
         -- Run the expression
        ; expr2 <- runMetaE zonked_q_expr
@@ -489,7 +512,7 @@ spliceResultDoc expr
         , ptext (sLit "To see what the splice expanded to, use -ddump-splices")]
 
 -------------------
-tcTopSpliceExpr :: TcM (LHsExpr Id) -> TcM (LHsExpr Id)
+tcTopSpliceExpr :: Bool -> TcM (LHsExpr Id) -> TcM (LHsExpr Id)
 -- Note [How top-level splices are handled]
 -- Type check an expression that is the body of a top-level splice
 --   (the caller will compile and run it)
@@ -499,7 +522,7 @@ tcTopSpliceExpr :: TcM (LHsExpr Id) -> TcM (LHsExpr Id)
 -- The recursive call to tcMonoExpr will simply expand the
 -- inner escape before dealing with the outer one
 
-tcTopSpliceExpr tc_action
+tcTopSpliceExpr isTypedSplice tc_action
   = checkNoErrs $  -- checkNoErrs: must not try to run the thing
                    -- if the type checker fails!
     unsetDOptM Opt_DeferTypeErrors $
@@ -508,7 +531,7 @@ tcTopSpliceExpr tc_action
                    -- coerce, so we get a seg-fault if, say we
                    -- splice a type into a place where an expression
                    -- is expected (Trac #7276)
-    setStage Splice $
+    setStage (Splice isTypedSplice) $
     do {    -- Typecheck the expression
          (expr', lie) <- captureConstraints tc_action
 
@@ -540,10 +563,10 @@ tcSpliceType (HsSplice _ name hs_expr) _
   = setSrcSpan (getLoc hs_expr) $ do
     { stage <- getStage
     ; case stage of {
-        Splice -> tcTopSpliceType hs_expr ;
-        Comp   -> tcTopSpliceType hs_expr ;
+        Splice {} -> tcTopSpliceType hs_expr ;
+        Comp      -> tcTopSpliceType hs_expr ;
 
-        Brack pop_level ps_var lie_var -> do
+        Brack _ pop_level ps_var lie_var -> do
            -- See Note [How brackets and nested splices are handled]
            -- A splice inside brackets
     { meta_ty <- tcMetaTy typeQTyConName
@@ -570,7 +593,7 @@ tcTopSpliceType expr
   = do  { meta_ty <- tcMetaTy typeQTyConName
 
         -- Typecheck the expression
-        ; zonked_q_expr <- tcTopSpliceExpr (tcMonoExpr expr meta_ty)
+        ; zonked_q_expr <- tcTopSpliceExpr False (tcMonoExpr expr meta_ty)
 
         -- Run the expression
         ; hs_ty2 <- runMetaT zonked_q_expr
@@ -596,7 +619,7 @@ tcTopSpliceType expr
 --      tcSpliceDecls :: LHsExpr Name -> TcM [LHsDecl RdrName]
 tcSpliceDecls expr
   = do  { list_q <- tcMetaTy decsQTyConName     -- Q [Dec]
-        ; zonked_q_expr <- tcTopSpliceExpr (tcMonoExpr expr list_q)
+        ; zonked_q_expr <- tcTopSpliceExpr False (tcMonoExpr expr list_q)
 
                 -- Run the expression
         ; decls <- runMetaD zonked_q_expr
@@ -623,7 +646,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 $
+    zonked_wrapped_expr' <- tcTopSpliceExpr False $
            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
@@ -731,7 +754,7 @@ runQuasiQuote (HsQuasiQuote quoter q_span quote) quote_selector meta_ty meta_ops
         ; meta_exp_ty <- tcMetaTy meta_ty
 
         -- Typecheck the expression
-        ; zonked_q_expr <- tcTopSpliceExpr (tcMonoExpr expr meta_exp_ty)
+        ; zonked_q_expr <- tcTopSpliceExpr False (tcMonoExpr expr meta_exp_ty)
 
         -- Run the expression
         ; result <- runMetaQ meta_ops zonked_q_expr
@@ -1000,6 +1023,18 @@ showSplice what before after
 
 illegalBracket :: SDoc
 illegalBracket = ptext (sLit "Template Haskell brackets cannot be nested (without intervening splices)")
+
+illegalTypedBracket :: SDoc
+illegalTypedBracket = ptext (sLit "Typed brackets may only appear in typed slices.")
+
+illegalUntypedBracket :: SDoc
+illegalUntypedBracket = ptext (sLit "Untyped brackets may only appear in untyped slices.")
+
+illegalTypedSplice :: SDoc
+illegalTypedSplice = ptext (sLit "Typed splices may not appear in untyped brackets")
+
+illegalUntypedSplice :: SDoc
+illegalUntypedSplice = ptext (sLit "Untyped splices may not appear in typed brackets")
 #endif  /* GHCI */
 \end{code}