Custom GADT syntaxwip/gadt-custom-syntax

5 files changed, 121 insertions, 47 deletions

diff --git a/compiler/GHC/Hs/Type.hs b/compiler/GHC/Hs/Type.hs
index 0e67a4a94e..6cfa467c45 100644
--- a/compiler/GHC/Hs/Type.hs
+++ b/compiler/GHC/Hs/Type.hs
@@ -70,7 +70,7 @@ module GHC.Hs.Type (
         splitLHsInstDeclTy, getLHsInstDeclHead, getLHsInstDeclClass_maybe,
         splitLHsPatSynTy,
         splitLHsForAllTyInvis, splitLHsForAllTyInvis_KP, splitLHsQualTy,
-        splitLHsSigmaTyInvis, splitLHsGadtTy,
+        splitLHsSigmaTyInvis,
         splitHsFunType, hsTyGetAppHead_maybe,
         mkHsOpTy, mkHsAppTy, mkHsAppTys, mkHsAppKindTy,
         ignoreParens, hsSigWcType, hsPatSigType,
@@ -550,32 +550,6 @@ splitLHsSigmaTyInvis ty
   , (ctxt, ty2) <- splitLHsQualTy ty1
   = (tvs, ctxt, ty2)
 
--- | Decompose a GADT type into its constituent parts.
--- Returns @(outer_bndrs, mb_ctxt, body)@, where:
---
--- * @outer_bndrs@ are 'HsOuterExplicit' if the type has explicit, outermost
---   type variable binders. Otherwise, they are 'HsOuterImplicit'.
---
--- * @mb_ctxt@ is @Just@ the context, if it is provided.
---   Otherwise, it is @Nothing@.
---
--- * @body@ is the body of the type after the optional @forall@s and context.
---
--- This function is careful not to look through parentheses.
--- See @Note [GADT abstract syntax] (Wrinkle: No nested foralls or contexts)@
--- "GHC.Hs.Decls" for why this is important.
-splitLHsGadtTy ::
-     LHsSigType GhcPs
-  -> (HsOuterSigTyVarBndrs GhcPs, Maybe (LHsContext GhcPs), LHsType GhcPs)
-splitLHsGadtTy (L _ sig_ty)
-  | (outer_bndrs, rho_ty) <- split_bndrs sig_ty
-  , (mb_ctxt, tau_ty)     <- splitLHsQualTy_KP rho_ty
-  = (outer_bndrs, mb_ctxt, tau_ty)
-  where
-    split_bndrs :: HsSigType GhcPs -> (HsOuterSigTyVarBndrs GhcPs, LHsType GhcPs)
-    split_bndrs (HsSig{sig_bndrs = outer_bndrs, sig_body = body_ty}) =
-      (outer_bndrs, body_ty)
-
 -- | Decompose a type of the form @forall <tvs>. body@ into its constituent
 -- parts. Only splits type variable binders that
 -- were quantified invisibly (e.g., @forall a.@, with a dot).
diff --git a/compiler/GHC/Parser.y b/compiler/GHC/Parser.y
index c17444ddcb..0279677741 100644
--- a/compiler/GHC/Parser.y
+++ b/compiler/GHC/Parser.y
@@ -2126,6 +2126,10 @@ ftype :: { forall b. DisambTD b => PV (Located b) }
 tyarg :: { LHsType GhcPs }
         : atype                         { $1 }
         | unpackedness atype            {% addUnpackednessP $1 $2 }
+        | '{' fielddecls '}'            {% amms (checkRecordSyntax
+                                                   (sLL $1 $> $ HsRecTy noExtField $2))
+                                                       -- Constructor sigs only
+                                                [moc $1,mcc $3] }
 
 tyop :: { Located RdrName }
         : qtyconop                      { $1 }
@@ -2146,10 +2150,6 @@ atype :: { LHsType GhcPs }
         | PREFIX_TILDE atype             {% ams (sLL $1 $> (mkBangTy SrcLazy $2)) [mj AnnTilde $1] }
         | PREFIX_BANG  atype             {% ams (sLL $1 $> (mkBangTy SrcStrict $2)) [mj AnnBang $1] }
 
-        | '{' fielddecls '}'             {% amms (checkRecordSyntax
-                                                    (sLL $1 $> $ HsRecTy noExtField $2))
-                                                        -- Constructor sigs only
-                                                 [moc $1,mcc $3] }
         | '(' ')'                        {% ams (sLL $1 $> $ HsTupleTy noExtField
                                                     HsBoxedOrConstraintTuple [])
                                                 [mop $1,mcp $2] }
@@ -2322,11 +2322,46 @@ gadt_constrs :: { Located [LConDecl GhcPs] }
 gadt_constr :: { LConDecl GhcPs }
     -- see Note [Difference in parsing GADT and data constructors]
     -- Returns a list because of:   C,D :: ty
-        : optSemi con_list '::' sigtype
+        : optSemi con_list '::' gadt_con_sig
                 {% do { (decl, anns) <- mkGadtDecl (unLoc $2) $4
                       ; ams (sLL $2 $> decl)
                             (mu AnnDcolon $3:anns) } }
 
+gadt_con_sig :: { Located GadtConSig }
+  : btype %shift           { sL1 $1 (GadtConSigRes $1) }
+  | forall_telescope gadt_con_sig
+           {% do { let (forall_anns, forall_tele) = unLoc $1
+                 ; ams (sLL $1 $> $ GadtConSigForAll forall_tele $2) forall_anns
+                 } }
+  | context '=>' gadt_con_sig
+           {% do { addAnnotation (gl $1) (toUnicodeAnn AnnDarrow $2) (gl $2)
+                 ; return (sLL $1 $> $ GadtConSigQual $1 $3)
+                 } }
+  | '{' fielddecls '}' '->' gadt_con_sig
+           {% do { r <- amms (checkRecordSyntax (sLL $1 $3 $2)) [moc $1,mcc $3]
+                 ; ams r [mu AnnRarrow $4] -- See Note [GADT decl discards annotations]
+                 ; ams (sLL $1 $> $ GadtConSigRecSyn r $5)
+                       [mu AnnRarrow $4]
+                 } }
+  | btype '->' gadt_con_sig
+           {% do { ams $1 [mu AnnRarrow $2] -- See Note [GADT decl discards annotations]
+                 ; ams (sLL $1 $> $ GadtConSigFunArg (HsScaled (HsUnrestrictedArrow (toUnicode $2)) $1) $3)
+                       [mu AnnRarrow $2]
+                 } }
+  | btype mult '->' gadt_con_sig
+           {% do { hintLinear (getLoc $2)
+                 ; let (arr, ann) = (unLoc $2) (toUnicode $3)
+                 ; ams $1 [ann, mu AnnRarrow $3] -- See Note [GADT decl discards annotations]
+                 ; ams (sLL $1 $> $ GadtConSigFunArg (HsScaled arr $1) $4)
+                       [ann, mu AnnRarrow $3]
+                 } }
+  | btype '->.' gadt_con_sig
+           {% do { hintLinear (getLoc $2)
+                 ; ams $1 [mu AnnLollyU $2] -- See Note [GADT decl discards annotations]
+                 ; ams (sLL $1 $> $ GadtConSigFunArg (HsScaled (HsLinearArrow UnicodeSyntax) $1) $3)
+                       [mu AnnLollyU $2]
+                 } }
+
 {- Note [Difference in parsing GADT and data constructors]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 GADT constructors have simpler syntax than usual data constructors:
diff --git a/compiler/GHC/Parser/Errors.hs b/compiler/GHC/Parser/Errors.hs
index 83812f7673..3052020417 100644
--- a/compiler/GHC/Parser/Errors.hs
+++ b/compiler/GHC/Parser/Errors.hs
@@ -385,6 +385,8 @@ data PsErrorDesc
       --
       -- TODO: distinguish errors without using SDoc
 
+   | PsErrIllegalGadtConSig  -- TODO (int-index): add a description of the problem
+
 
 newtype StarIsType = StarIsType Bool
 
diff --git a/compiler/GHC/Parser/Errors/Ppr.hs b/compiler/GHC/Parser/Errors/Ppr.hs
index 47c8104fd1..48cac3af0a 100644
--- a/compiler/GHC/Parser/Errors/Ppr.hs
+++ b/compiler/GHC/Parser/Errors/Ppr.hs
@@ -587,6 +587,9 @@ pp_err = \case
          , text "Perhaps you intended to use QualifiedDo"
          ]
 
+   PsErrIllegalGadtConSig
+      -> text "Illegal GADT constructor signature"
+
 pp_unexpected_fun_app :: Outputable a => SDoc -> a -> SDoc
 pp_unexpected_fun_app e a =
    text "Unexpected " <> e <> text " in function application:"
diff --git a/compiler/GHC/Parser/PostProcess.hs b/compiler/GHC/Parser/PostProcess.hs
index 6a0f86aefe..9ba1aae853 100644
--- a/compiler/GHC/Parser/PostProcess.hs
+++ b/compiler/GHC/Parser/PostProcess.hs
@@ -89,6 +89,9 @@ module GHC.Parser.PostProcess (
 
         SumOrTuple (..),
 
+        LGadtConSig,
+        GadtConSig(..),
+
         -- Expression/command/pattern ambiguity resolution
         PV,
         runPV,
@@ -623,6 +626,69 @@ mkConDeclH98 name mb_forall mb_cxt args
                , con_args   = args
                , con_doc    = Nothing }
 
+type LGadtConSig = Located GadtConSig
+
+data GadtConSig
+  = GadtConSigRes !(LHsType GhcPs)
+  | GadtConSigRecSyn !(Located [LConDeclField GhcPs]) LGadtConSig
+  | GadtConSigFunArg !(HsScaled GhcPs (LHsType GhcPs)) LGadtConSig
+  | GadtConSigForAll !(HsForAllTelescope GhcPs) LGadtConSig
+  | GadtConSigQual !(LHsContext GhcPs) LGadtConSig
+
+data GadtConSigParts =
+  GadtConSigParts
+    !(HsOuterSigTyVarBndrs GhcPs)
+    !(Maybe (LHsContext GhcPs))
+    !(HsConDeclGADTDetails GhcPs)
+    !(LHsType GhcPs)
+
+-- | Decompose a GADT type into its constituent parts.
+--
+-- This function is careful not to look through parentheses.
+-- See @Note [GADT abstract syntax] (Wrinkle: No nested foralls or contexts)@
+-- "GHC.Hs.Decls" for why this is important.
+splitGadtSig :: LGadtConSig -> ([PsError], GadtConSigParts, [AddAnn])
+splitGadtSig gsig0 = do
+    let (bndrs, gsig_rho) = split_bndrs gsig0
+        (mctxt, gsig_tau) = split_ctxt gsig_rho
+    case unLoc gsig_tau of
+      GadtConSigRecSyn rec gsig' -> do
+        let details = RecConGADT rec
+        case gsig' of
+          L _ (GadtConSigRes res) -> ([], GadtConSigParts bndrs mctxt details res, [])
+          L l _ -> ( [PsError PsErrIllegalGadtConSig [] l]
+                   , GadtConSigParts bndrs mctxt details (L l (HsWildCardTy noExtField))
+                   , [] )
+      _ -> do
+        let (args1, gsig') = split_args [] gsig_tau
+        case gsig' of
+          L _ (GadtConSigRes res) -> do
+            let (args2, res', anns) = splitHsFunType res
+                details = PrefixConGADT (args1 ++ args2)
+            ([], GadtConSigParts bndrs mctxt details res', anns)
+          L l _ -> do
+            let details = PrefixConGADT args1
+            ( [PsError PsErrIllegalGadtConSig [] l],
+              GadtConSigParts bndrs mctxt details (L l (HsWildCardTy noExtField)),
+              [] )
+  where
+    -- See also: GHC.Hs.Utils.hsTypeToHsSigType
+    split_bndrs (L _ (GadtConSigForAll forall_tele gsig))
+      | HsForAllInvis { hsf_invis_bndrs = bndrs } <- forall_tele
+      = (mkHsOuterExplicit bndrs, gsig)
+    split_bndrs gsig
+      = (mkHsOuterImplicit, gsig)
+
+    -- See also: GHC.Hs.Type.splitLHsQualTy_KP
+    split_ctxt (L _ (GadtConSigQual ctxt gsig)) = (Just ctxt, gsig)
+    split_ctxt gsig = (Nothing, gsig)
+
+    -- See also: GHC.Hs.Type.splitHsFunType
+    split_args args gsig =
+      case gsig of
+        L _ (GadtConSigFunArg arg gsig') -> split_args (arg : args) gsig'
+        _ -> (reverse args, gsig)
+
 -- | Construct a GADT-style data constructor from the constructor names and
 -- their type. Some interesting aspects of this function:
 --
@@ -631,26 +697,20 @@ mkConDeclH98 name mb_forall mb_cxt args
 --   records whether this is a prefix or record GADT constructor. See
 --   Note [GADT abstract syntax] in "GHC.Hs.Decls" for more details.
 mkGadtDecl :: [Located RdrName]
-           -> LHsSigType GhcPs
+           -> Located GadtConSig
            -> P (ConDecl GhcPs, [AddAnn])
-mkGadtDecl names ty = do
-  let (args, res_ty, anns)
-        | L _ (HsFunTy _ _w (L loc (HsRecTy _ rf)) res_ty) <- body_ty
-        = (RecConGADT (L loc rf), res_ty, [])
-        | otherwise
-        = let (arg_types, res_type, anns) = splitHsFunType body_ty
-          in (PrefixConGADT arg_types, res_type, anns)
-
+mkGadtDecl names gsig = do
+  let (errs, parts, anns) = splitGadtSig gsig
+  mapM_ addError errs
+  let GadtConSigParts bndrs mctxt details res = parts
   pure ( ConDeclGADT { con_g_ext  = noExtField
                      , con_names  = names
-                     , con_bndrs  = L (getLoc ty) outer_bndrs
-                     , con_mb_cxt = mcxt
-                     , con_g_args = args
-                     , con_res_ty = res_ty
+                     , con_bndrs  = L (getLoc gsig) bndrs
+                     , con_mb_cxt = mctxt
+                     , con_g_args = details
+                     , con_res_ty = res
                      , con_doc    = Nothing }
        , anns )
-  where
-    (outer_bndrs, mcxt, body_ty) = splitLHsGadtTy ty
 
 setRdrNameSpace :: RdrName -> NameSpace -> RdrName
 -- ^ This rather gruesome function is used mainly by the parser.