[SyntaxParse] Parse PoundAssertStmt

Also added parseExpressionSyntax as an intermediate syntax parser until
expression parsing is completely implemented.

Author: Nathan Hawes

include/swift/Parse/ASTGen.h

@@ -38,6 +38,8 @@ class ASTGen {
   // them
   /// Types that cannot be represented by Syntax or generated by ASTGen.
   llvm::DenseMap<SourceLoc, TypeRepr *> Types;
+  /// Exprs that cannot be represented by Syntax or generated by ASTGen.
+  llvm::DenseMap<SourceLoc, Expr *> Exprs;
 
   llvm::DenseMap<SourceLoc, DeclAttributes> ParsedDeclAttrs;
 
@@ -74,6 +76,13 @@ class ASTGen {
   static MagicIdentifierLiteralExpr::Kind
   getMagicIdentifierLiteralKind(tok Kind);
 
+public:
+  //===--------------------------------------------------------------------===//
+  // Statements
+
+  Stmt *generate(const syntax::PoundAssertStmtSyntax &Stmt,
+                 const SourceLoc Loc);
+
 public:
   //===--------------------------------------------------------------------===//
   // Types.
@@ -160,6 +169,13 @@ class ASTGen {
   static SourceLoc advanceLocBegin(const SourceLoc &Loc,
                                    const syntax::Syntax &Node);
 
+  /// Advance \p Loc to the last non-missing token of the \p Node or, if it
+  /// doesn't contain any, the last non-missing token preceding it in the tree.
+  /// \p Loc must be the leading trivia of the first token in the tree in which
+  /// \p Node resides
+  static SourceLoc advanceLocEnd(const SourceLoc &Loc,
+                                 const syntax::Syntax &Node);
+
   ValueDecl *lookupInScope(DeclName Name);
 
   void addToScope(ValueDecl *D, bool diagnoseRedefinitions = true);
@@ -173,6 +189,10 @@ class ASTGen {
   bool hasType(const SourceLoc Loc) const;
   TypeRepr *getType(const SourceLoc Loc) const;
 
+  void addExpr(Expr *Expr, const SourceLoc Loc);
+  bool hasExpr(const SourceLoc Loc) const;
+  Expr *getExpr(const SourceLoc Loc) const;
+
   void addDeclAttributes(DeclAttributes attrs, const SourceLoc Loc);
   bool hasDeclAttributes(SourceLoc Loc) const;
   DeclAttributes getDeclAttributes(const SourceLoc Loc) const;

include/swift/Parse/Parser.h

@@ -1420,6 +1420,7 @@ class Parser {
 
   //===--------------------------------------------------------------------===//
   // Expression Parsing
+  ParsedSyntaxResult<ParsedExprSyntax> parseExpressionSyntax(Diag<> ID);
   ParserResult<Expr> parseExpr(Diag<> ID) {
     return parseExprImpl(ID, /*isExprBasic=*/false);
   }
@@ -1603,6 +1604,7 @@ class Parser {
   ParserStatus parseStmtCases(SmallVectorImpl<ASTNode> &cases, bool IsActive);
   ParserResult<CaseStmt> parseStmtCase(bool IsActive);
   ParserResult<Stmt> parseStmtPoundAssert();
+  ParsedSyntaxResult<ParsedPoundAssertStmtSyntax> parseStmtPoundAssertSyntax();
 
   //===--------------------------------------------------------------------===//
   // Generics Parsing

include/swift/Syntax/Syntax.h

@@ -169,6 +169,9 @@ class Syntax {
   /// Returns true if the node is "present" in the source.
   bool isPresent() const;
 
+  /// Get the node immediately before this current node that does contain a
+  /// non-missing token. Return nullptr if we cannot find such node.
+  Optional<Syntax> getPreviousNode() const;
 
   /// Returns the first non-missing token in this syntax. Returns None if there
   /// is no non-missing token.

lib/Parse/ASTGen.cpp

@@ -711,10 +711,50 @@ LayoutConstraint ASTGen::generate(const LayoutConstraintSyntax &constraint,
                                                Context);
 }
 
+Stmt *ASTGen::generate(const syntax::PoundAssertStmtSyntax &Stmt,
+                       const SourceLoc Loc) {
+  // Don't form a PoundAssertStmt without a condition
+  if (Stmt.getCondition().isUnknown())
+    return nullptr;
+  SourceLoc CondLoc = advanceLocBegin(Loc, Stmt.getCondition());
+  if (!hasExpr(CondLoc))
+    return nullptr;
+
+  Expr *CondExpr = getExpr(CondLoc);
+  SourceLoc Start = advanceLocBegin(Loc, Stmt.getPoundAssert());
+  SourceLoc End = advanceLocEnd(Loc, Stmt.getRightParen());
+
+  StringRef MessageText = "";
+  if (auto Message = Stmt.getMessage()) {
+    auto Tok = P.L->getTokenAt(generate(*Message, Loc));
+    SmallVector<Lexer::StringSegment, 1> Segments;
+    P.L->getStringLiteralSegments(Tok, Segments);
+    if (Segments.size() == 1 &&
+      Segments.front().Kind == Lexer::StringSegment::Literal &&
+      // FIXME: Support extended escaping string literal.
+      Tok.getCustomDelimiterLen() == 0) {
+      MessageText = P.SourceMgr.extractText(
+          CharSourceRange(Segments.front().Loc, Segments.front().Length));
+    }
+  }
+  return new (Context) PoundAssertStmt(SourceRange(Start, End), CondExpr,
+                                       MessageText);
+}
+
 SourceLoc ASTGen::advanceLocBegin(const SourceLoc &Loc, const Syntax &Node) {
   return Loc.getAdvancedLoc(Node.getAbsolutePosition().getOffset());
 }
 
+SourceLoc ASTGen::advanceLocEnd(const SourceLoc &Loc,
+                                const syntax::Syntax &Node) {
+  if (auto Tok = Node.getLastToken())
+    return advanceLocBegin(Loc, *Tok);
+  if (auto Prev = Node.getPreviousNode())
+    return advanceLocBegin(Loc, *Prev->getLastToken());
+  assert(false && "No tokens in tree?");
+  return Loc;
+}
+
 StringRef ASTGen::copyAndStripUnderscores(StringRef Orig) {
   return copyAndStripUnderscores(Orig, Context);
 }
@@ -780,6 +820,30 @@ TypeRepr *ASTGen::getType(const SourceLoc Loc) const {
   return Types.find(Loc)->second;
 }
 
+void ASTGen::addExpr(Expr *E, const SourceLoc Loc) {
+#ifndef NDEBUG
+  if (hasExpr(Loc)) {
+    bool PrevIsSubExpr = false;
+    Expr *Prev = Exprs.find(Loc)->second;
+    E->forEachChildExpr([&](Expr *Child) {
+      if (Child == Prev)
+        PrevIsSubExpr = true;
+      return Child;
+    });
+    assert(PrevIsSubExpr);
+  }
+#endif
+  Exprs.insert({Loc, E});
+}
+
+bool ASTGen::hasExpr(const SourceLoc Loc) const {
+  return Exprs.find(Loc) != Exprs.end();
+}
+
+Expr *ASTGen::getExpr(const SourceLoc Loc) const {
+  return Exprs.find(Loc)->second;
+}
+
 void ASTGen::addDeclAttributes(DeclAttributes attrs, SourceLoc Loc) {
   ParsedDeclAttrs.insert({Loc, attrs});
 }

lib/Parse/ParseExpr.cpp

@@ -34,6 +34,18 @@
 using namespace swift;
 using namespace swift::syntax;
 
+ParsedSyntaxResult<ParsedExprSyntax> Parser::parseExpressionSyntax(Diag<> ID) {
+  SourceLoc ExprLoc = Tok.getLoc();
+  SyntaxParsingContext ExprParsingContext(SyntaxContext, SyntaxContextKind::Expr);
+  ExprParsingContext.setTransparent();
+  ParserResult<Expr> Result = parseExpr(ID);
+  if (auto ParsedExpr = ExprParsingContext.popIf<ParsedExprSyntax>()) {
+    Generator.addExpr(Result.getPtrOrNull(), ExprLoc);
+    return makeParsedResult(std::move(*ParsedExpr), Result.getStatus());
+  }
+  return Result.getStatus();
+}
+
 /// parseExpr
 ///
 ///   expr:

lib/Parse/ParseStmt.cpp

@@ -23,6 +23,8 @@
 #include "swift/Parse/Lexer.h"
 #include "swift/Parse/Parser.h"
 #include "swift/Parse/SyntaxParsingContext.h"
+#include "swift/Parse/ParsedSyntaxBuilders.h"
+#include "swift/Parse/ParsedSyntaxRecorder.h"
 #include "swift/Subsystems.h"
 #include "swift/Syntax/TokenSyntax.h"
 #include "llvm/ADT/PointerUnion.h"
@@ -2529,49 +2531,76 @@ ParserResult<CaseStmt> Parser::parseStmtCase(bool IsActive) {
 /// stmt-pound-assert:
 ///   '#assert' '(' expr (',' string_literal)? ')'
 ParserResult<Stmt> Parser::parseStmtPoundAssert() {
-  SyntaxContext->setCreateSyntax(SyntaxKind::PoundAssertStmt);
+  auto leadingLoc = leadingTriviaLoc();
+  auto parsed = parseStmtPoundAssertSyntax();
+  SyntaxContext->addSyntax(parsed.get());
+  auto poundAssertSyntax = SyntaxContext->topNode<PoundAssertStmtSyntax>();
+  auto poundAssert = Generator.generate(poundAssertSyntax, leadingLoc);
+  return makeParserResult(parsed.getStatus(), poundAssert);
+}
 
-  SourceLoc startLoc = consumeToken(tok::pound_assert);
-  SourceLoc endLoc;
+ParsedSyntaxResult<ParsedPoundAssertStmtSyntax>
+Parser::parseStmtPoundAssertSyntax() {
+  ParsedPoundAssertStmtSyntaxBuilder builder(*SyntaxContext);
+  ParserStatus status;
 
-  if (Tok.isNot(tok::l_paren)) {
-    diagnose(Tok, diag::pound_assert_expected_lparen);
-    return makeParserError();
+  SourceLoc startLoc = Tok.getLoc();
+  builder.usePoundAssert(consumeTokenSyntax(tok::pound_assert));
+  if (!Tok.isFollowingLParen()) {
+    diagnose(Tok.getLoc(), diag::pound_assert_expected_lparen);
+    status.setIsParseError();
+    builder.useCondition(ParsedSyntaxRecorder::makeUnknownExpr({}, *SyntaxContext));
+    return makeParsedResult(builder.build(), status);
   }
-  SourceLoc LBLoc = consumeToken(tok::l_paren);
+  SourceLoc lParenLoc = Tok.getLoc();
+  builder.useLeftParen(consumeTokenSyntax(tok::l_paren));
 
-  auto conditionExprResult = parseExpr(diag::pound_assert_expected_expression);
-  if (conditionExprResult.isParseError())
-    return ParserStatus(conditionExprResult);
+  auto conditionExprResult = parseExpressionSyntax(diag::pound_assert_expected_expression);
+  status |= conditionExprResult.getStatus();
+  builder.useCondition(conditionExprResult.isNull()
+      ? ParsedSyntaxRecorder::makeUnknownExpr({}, *SyntaxContext)
+      : conditionExprResult.get());
+
+  if (auto Comma = consumeTokenSyntaxIf(tok::comma)) {
+    builder.useComma(std::move(*Comma));
 
-  StringRef message;
-  if (consumeIf(tok::comma)) {
     if (Tok.isNot(tok::string_literal)) {
       diagnose(Tok.getLoc(), diag::pound_assert_expected_string_literal);
-      return makeParserError();
+      status.setIsParseError();
+      return makeParsedResult(builder.build(), status);
     }
-
-    auto messageOpt = getStringLiteralIfNotInterpolated(Tok.getLoc(),
-                                                        "'#assert' message");
-    consumeToken();
-    if (!messageOpt)
-      return makeParserError();
-
-    message = *messageOpt;
+    auto MessageLoc = Tok.getLoc();
+    // FIXME: Support extended escaping string literal.
+    if (Tok.getCustomDelimiterLen()) {
+      diagnose(MessageLoc, diag::forbidden_extended_escaping_string, "'#assert' message");
+      status.setIsParseError();
+    } else {
+      SmallVector<Lexer::StringSegment, 1> Segments;
+      L->getStringLiteralSegments(Tok, Segments);
+      if (Segments.size() != 1 ||
+          Segments.front().Kind == Lexer::StringSegment::Expr) {
+        diagnose(MessageLoc, diag::forbidden_interpolated_string, "'#assert' message");
+        status.setIsParseError();
+      }
+    }
+    builder.useMessage(consumeTokenSyntax(tok::string_literal));
   }
 
-  if (parseMatchingToken(tok::r_paren, endLoc,
-                         diag::pound_assert_expected_rparen, LBLoc)) {
-    return makeParserError();
+  SourceLoc EndLoc;
+  if (auto rParen = parseMatchingTokenSyntax(tok::r_paren, EndLoc,
+                                             diag::pound_assert_expected_rparen,
+                                             lParenLoc)) {
+    builder.useRightParen(std::move(*rParen));
+  } else {
+    status.setIsParseError();
   }
 
   // We check this after consuming everything, so that the SyntaxContext
   // understands this statement even when the feature is disabled.
   if (!Context.LangOpts.EnableExperimentalStaticAssert) {
     diagnose(startLoc, diag::pound_assert_disabled);
-    return makeParserError();
+    status.setIsParseError();
   }
 
-  return makeParserResult<Stmt>(new (Context) PoundAssertStmt(
-      SourceRange(startLoc, endLoc), conditionExprResult.get(), message));
+  return makeParsedResult(builder.build(), status);
 }

lib/Syntax/Syntax.cpp

@@ -98,6 +98,12 @@ llvm::Optional<Syntax> Syntax::getChild(const size_t N) const {
   return Syntax {Root, ChildData.get()};
 }
 
+Optional<Syntax> Syntax::getPreviousNode() const {
+  if (auto prev = getData().getPreviousNode())
+    return TokenSyntax(Root, prev.get());
+  return None;
+}
+
 Optional<TokenSyntax> Syntax::getFirstToken() const {
   if (auto tok = getData().getFirstToken())
     return TokenSyntax(Root, tok.get());