[SyntaxParse] Prevent memory leak in Syntax parsing

Essentially re-applying b09f875

Author: rintaro

include/swift/Parse/ParsedRawSyntaxNode.h

@@ -51,7 +51,7 @@ class ParsedRawSyntaxNode {
     CharSourceRange Range;
   };
   struct DeferredLayoutNode {
-    ArrayRef<ParsedRawSyntaxNode> Children;
+    MutableArrayRef<ParsedRawSyntaxNode> Children;
   };
   struct DeferredTokenNode {
     const ParsedTriviaPiece *TriviaPieces;
@@ -73,8 +73,8 @@ class ParsedRawSyntaxNode {
   bool IsMissing = false;
 
   ParsedRawSyntaxNode(syntax::SyntaxKind k,
-                      ArrayRef<ParsedRawSyntaxNode> deferredNodes)
-    : DeferredLayout{deferredNodes},
+                      MutableArrayRef<ParsedRawSyntaxNode> deferredNodes)
+    : DeferredLayout({deferredNodes}),
       SynKind(uint16_t(k)), TokKind(uint16_t(tok::unknown)),
       DK(DataKind::DeferredLayout) {
     assert(getKind() == k && "Syntax kind with too large value!");
@@ -97,6 +97,8 @@ class ParsedRawSyntaxNode {
     assert(DeferredToken.NumTrailingTrivia == numTrailingTrivia &&
            "numLeadingTrivia is too large value!");
   }
+  ParsedRawSyntaxNode(ParsedRawSyntaxNode &other) = delete;
+  ParsedRawSyntaxNode &operator=(ParsedRawSyntaxNode &other) = delete;
 
 public:
   ParsedRawSyntaxNode()
@@ -115,6 +117,35 @@ class ParsedRawSyntaxNode {
     assert(getTokenKind() == tokKind && "Token kind with too large value!");
   }
 
+  ParsedRawSyntaxNode &operator=(ParsedRawSyntaxNode &&other) {
+    assert(DK != DataKind::Recorded);
+    switch (other.DK) {
+    case DataKind::Null:
+      break;
+    case DataKind::Recorded:
+      RecordedData = std::move(other.RecordedData);
+      break;
+    case DataKind::DeferredLayout:
+      DeferredLayout = std::move(other.DeferredLayout);
+      break;
+    case DataKind::DeferredToken:
+      DeferredToken = std::move(other.DeferredToken);
+      break;
+    }
+    SynKind = std::move(other.SynKind);
+    TokKind = std::move(other.TokKind);
+    DK = std::move(other.DK);
+    IsMissing = std::move(other.IsMissing);
+    other.reset();
+    return *this;
+  }
+  ParsedRawSyntaxNode(ParsedRawSyntaxNode &&other) : ParsedRawSyntaxNode() {
+    *this = std::move(other);
+  }
+  ~ParsedRawSyntaxNode() {
+    assert(DK != DataKind::Recorded);
+  }
+
   syntax::SyntaxKind getKind() const { return syntax::SyntaxKind(SynKind); }
   tok getTokenKind() const { return tok(TokKind); }
 
@@ -136,16 +167,52 @@ class ParsedRawSyntaxNode {
   /// Primary used for a deferred missing token.
   bool isMissing() const { return IsMissing; }
 
+  void reset() {
+    RecordedData = {};
+    SynKind = uint16_t(syntax::SyntaxKind::Unknown);
+    TokKind = uint16_t(tok::unknown);
+    DK = DataKind::Null;
+    IsMissing = false;
+  }
+
+  ParsedRawSyntaxNode unsafeCopy() const {
+    ParsedRawSyntaxNode copy;
+    switch (DK) {
+    case DataKind::DeferredLayout:
+      copy.DeferredLayout = DeferredLayout;
+      break;
+    case DataKind::DeferredToken:
+      copy.DeferredToken = DeferredToken;
+      break;
+    case DataKind::Recorded:
+      copy.RecordedData = RecordedData;
+      break;
+    case DataKind::Null:
+      break;
+    }
+    copy.SynKind = SynKind;
+    copy.TokKind = TokKind;
+    copy.DK = DK;
+    copy.IsMissing = IsMissing;
+    return copy;
+  }
+
   // Recorded Data ===========================================================//
 
   CharSourceRange getRange() const {
     assert(isRecorded());
     return RecordedData.Range;
   }
-  OpaqueSyntaxNode getOpaqueNode() const {
+  const OpaqueSyntaxNode &getOpaqueNode() const {
     assert(isRecorded());
     return RecordedData.OpaqueNode;
   }
+  OpaqueSyntaxNode takeOpaqueNode() {
+    assert(isRecorded());
+    auto opaque = RecordedData.OpaqueNode;
+    reset();
+    return opaque;
+  }
 
   // Deferred Layout Data ====================================================//
 
@@ -154,6 +221,30 @@ class ParsedRawSyntaxNode {
     return DeferredLayout.Children;
   }
 
+  MutableArrayRef<ParsedRawSyntaxNode> getDeferredChildren() {
+    assert(DK == DataKind::DeferredLayout);
+    return DeferredLayout.Children;
+  }
+
+  ParsedRawSyntaxNode copyDeferred() const {
+    ParsedRawSyntaxNode copy;
+    switch (DK) {
+    case DataKind::DeferredLayout:
+      copy.DeferredLayout = DeferredLayout;
+      break;
+    case DataKind::DeferredToken:
+      copy.DeferredToken = DeferredToken;
+      break;
+    default:
+      llvm_unreachable("node not deferred");
+    }
+    copy.SynKind = SynKind;
+    copy.TokKind = TokKind;
+    copy.DK = DK;
+    copy.IsMissing = IsMissing;
+    return copy;
+  }
+
   // Deferred Token Data =====================================================//
 
   CharSourceRange getDeferredTokenRange() const {
@@ -176,7 +267,7 @@ class ParsedRawSyntaxNode {
 
   /// Form a deferred syntax layout node.
   static ParsedRawSyntaxNode makeDeferred(syntax::SyntaxKind k,
-                        ArrayRef<ParsedRawSyntaxNode> deferredNodes,
+                        MutableArrayRef<ParsedRawSyntaxNode> deferredNodes,
                                           SyntaxParsingContext &ctx);
 
   /// Form a deferred token node.

include/swift/Parse/ParsedRawSyntaxRecorder.h

@@ -60,14 +60,16 @@ class ParsedRawSyntaxRecorder {
   /// \p kind. Missing optional elements are represented with a null
   /// ParsedRawSyntaxNode object.
   ParsedRawSyntaxNode recordRawSyntax(syntax::SyntaxKind kind,
-                                     ArrayRef<ParsedRawSyntaxNode> elements);
+                                      MutableArrayRef<ParsedRawSyntaxNode> elements);
 
   /// Record a raw syntax collecton without eny elements. \p loc can be invalid
   /// or an approximate location of where an element of the collection would be
   /// if not missing.
   ParsedRawSyntaxNode recordEmptyRawSyntaxCollection(syntax::SyntaxKind kind,
                                                      SourceLoc loc);
 
+  void discardRecordedNode(ParsedRawSyntaxNode &node);
+
   /// Used for incremental re-parsing.
   ParsedRawSyntaxNode lookupNode(size_t lexerOffset, SourceLoc loc,
                                  syntax::SyntaxKind kind);

include/swift/Parse/ParsedSyntax.h

@@ -26,6 +26,7 @@ class ParsedSyntax {
     : RawNode(std::move(rawNode)) {}
 
   const ParsedRawSyntaxNode &getRaw() const { return RawNode; }
+  ParsedRawSyntaxNode takeRaw() { return std::move(RawNode); }
   syntax::SyntaxKind getKind() const { return RawNode.getKind(); }
 
   /// Returns true if the syntax node is of the given type.
@@ -39,7 +40,7 @@ class ParsedSyntax {
   template <typename T>
   T castTo() const {
     assert(is<T>() && "castTo<T>() node of incompatible type!");
-    return T { RawNode };
+    return T { RawNode.copyDeferred() };
   }
 
   /// If this Syntax node is of the right kind, cast and return it,
@@ -52,6 +53,10 @@ class ParsedSyntax {
     return llvm::None;
   }
 
+  ParsedSyntax copyDeferred() const {
+    return ParsedSyntax { RawNode.copyDeferred() };
+  }
+
   static bool kindof(syntax::SyntaxKind Kind) {
     return true;
   }
@@ -65,7 +70,7 @@ class ParsedSyntax {
 class ParsedTokenSyntax final : public ParsedSyntax {
 public:
   explicit ParsedTokenSyntax(ParsedRawSyntaxNode rawNode)
-    : ParsedSyntax(rawNode) {}
+    : ParsedSyntax(std::move(rawNode)) {}
 
   tok getTokenKind() const {
     return getRaw().getTokenKind();

include/swift/Parse/ParsedSyntaxRecorder.h.gyb

@@ -53,15 +53,15 @@ public:
 %   elif node.is_syntax_collection():
 private:
   static Parsed${node.name} record${node.syntax_kind}(
-      ArrayRef<Parsed${node.collection_element_type}> elts,
+      MutableArrayRef<Parsed${node.collection_element_type}> elts,
       ParsedRawSyntaxRecorder &rec);
 
 public:
   static Parsed${node.name} defer${node.syntax_kind}(
-      ArrayRef<Parsed${node.collection_element_type}> elts,
+      MutableArrayRef<Parsed${node.collection_element_type}> elts,
       SyntaxParsingContext &SPCtx);
   static Parsed${node.name} make${node.syntax_kind}(
-      ArrayRef<Parsed${node.collection_element_type}> elts,
+      MutableArrayRef<Parsed${node.collection_element_type}> elts,
       SyntaxParsingContext &SPCtx);
 
   static Parsed${node.name} makeBlank${node.syntax_kind}(SourceLoc loc,
@@ -75,18 +75,21 @@ public:
   /// optional trailing comma.
   static ParsedTupleTypeElementSyntax
   makeTupleTypeElement(ParsedTypeSyntax Type,
-                         Optional<ParsedTokenSyntax> TrailingComma,
-                         SyntaxParsingContext &SPCtx);
+                       Optional<ParsedTokenSyntax> TrailingComma,
+                       SyntaxParsingContext &SPCtx);
 
   /// The provided \c elements are in the appropriate order for the syntax
   /// \c kind's layout but optional elements are not be included.
   /// This function will form the exact layout based on the provided elements,
   /// substituting missing parts with a null ParsedRawSyntaxNode object.
   ///
   /// \returns true if the layout could be formed, false otherwise.
-  static bool formExactLayoutFor(syntax::SyntaxKind kind,
-                                 ArrayRef<ParsedRawSyntaxNode> elements,
-         function_ref<void(syntax::SyntaxKind, ArrayRef<ParsedRawSyntaxNode>)> receiver);
+  static bool
+  formExactLayoutFor(syntax::SyntaxKind kind,
+                     MutableArrayRef<ParsedRawSyntaxNode> elements,
+                     function_ref<void(syntax::SyntaxKind,
+                                       MutableArrayRef<ParsedRawSyntaxNode>)>
+                         receiver);
 };
 }
 

include/swift/Parse/Parser.h

@@ -358,6 +358,22 @@ class Parser {
   };
   friend class StructureMarkerRAII;
 
+  /// A RAII object that tells the SyntaxParsingContext to defer Syntax nodes.
+  class DeferringContextRAII {
+    SyntaxParsingContext &Ctx;
+    bool WasDeferring;
+
+  public:
+    explicit DeferringContextRAII(SyntaxParsingContext &SPCtx)
+        : Ctx(SPCtx), WasDeferring(Ctx.shouldDefer()) {
+      Ctx.setShouldDefer();
+    }
+
+    ~DeferringContextRAII() {
+      Ctx.setShouldDefer(WasDeferring);
+    }
+  };
+
   /// The stack of structure markers indicating the locations of
   /// structural elements actively being parsed, including the start
   /// of declarations, statements, and opening operators of various
@@ -390,7 +406,7 @@ class Parser {
 
   /// Calling this function to finalize libSyntax tree creation without destroying
   /// the parser instance.
-  ParsedRawSyntaxNode finalizeSyntaxTree() {
+  OpaqueSyntaxNode finalizeSyntaxTree() {
     assert(Tok.is(tok::eof) && "not done parsing yet");
     return SyntaxContext->finalizeRoot();
   }

include/swift/Parse/SyntaxParseActions.h

@@ -55,6 +55,13 @@ class SyntaxParseActions {
                                            ArrayRef<OpaqueSyntaxNode> elements,
                                            CharSourceRange range) = 0;
 
+  /// Discard raw syntax node.
+  /// 
+  /// FIXME: This breaks invariant that any recorded node will be a part of the
+  /// result SourceFile syntax. This method is a temporary workaround, and
+  /// should be removed when we fully migrate to libSyntax parsing.
+  virtual void discardRecordedNode(OpaqueSyntaxNode node) = 0;
+
   /// Used for incremental re-parsing.
   virtual std::pair<size_t, OpaqueSyntaxNode>
   lookupNode(size_t lexerOffset, syntax::SyntaxKind kind) {

include/swift/Parse/SyntaxParserResult.h

@@ -29,17 +29,17 @@ template <typename Syntax, typename AST> class SyntaxParserResult {
       : SyntaxNode(None), ASTResult(nullptr) {}
   SyntaxParserResult(ParserStatus Status)
       : SyntaxNode(None), ASTResult(Status) {}
-  SyntaxParserResult(llvm::Optional<Syntax> SyntaxNode, AST *ASTNode)
-      : SyntaxNode(SyntaxNode), ASTResult(ASTNode) {}
-  SyntaxParserResult(ParserStatus Status, llvm::Optional<Syntax> SyntaxNode,
+  SyntaxParserResult(llvm::Optional<Syntax> &&SyntaxNode, AST *ASTNode)
+      : SyntaxNode(std::move(SyntaxNode)), ASTResult(ASTNode) {}
+  SyntaxParserResult(ParserStatus Status, llvm::Optional<Syntax> &&SyntaxNode,
                      AST *ASTNode)
-      : SyntaxNode(SyntaxNode), ASTResult(makeParserResult(Status, ASTNode)) {}
+      : SyntaxNode(std::move(SyntaxNode)), ASTResult(makeParserResult(Status, ASTNode)) {}
 
   /// Convert from a different but compatible parser result.
   template <typename U, typename Enabler = typename std::enable_if<
                             std::is_base_of<AST, U>::value>::type>
-  SyntaxParserResult(SyntaxParserResult<Syntax, U> Other)
-      : SyntaxNode(Other.SyntaxNode), ASTResult(Other.ASTResult) {}
+  SyntaxParserResult(SyntaxParserResult<Syntax, U> &&Other)
+      : SyntaxNode(std::move(Other.SyntaxNode)), ASTResult(Other.ASTResult) {}
 
 
   bool isNull() const { return ASTResult.isNull(); }
@@ -58,9 +58,9 @@ template <typename Syntax, typename AST> class SyntaxParserResult {
     return SyntaxNode.hasValue();
   }
 
-  Syntax getSyntax() const {
+  Syntax getSyntax() {
     assert(SyntaxNode.hasValue() && "getSyntax from None value");
-    return *SyntaxNode;
+    return std::move(*SyntaxNode);
   }
 
   SyntaxParserResult<Syntax, AST> &
@@ -73,16 +73,16 @@ template <typename Syntax, typename AST> class SyntaxParserResult {
 /// Create a successful parser result.
 template <typename Syntax, typename AST>
 static inline SyntaxParserResult<Syntax, AST>
-makeSyntaxResult(llvm::Optional<Syntax> SyntaxNode, AST *ASTNode) {
-  return SyntaxParserResult<Syntax, AST>(SyntaxNode, ASTNode);
+makeSyntaxResult(llvm::Optional<Syntax> &&SyntaxNode, AST *ASTNode) {
+  return SyntaxParserResult<Syntax, AST>(std::move(SyntaxNode), ASTNode);
 }
 
 /// Create a result with the specified status.
 template <typename Syntax, typename AST>
 static inline SyntaxParserResult<Syntax, AST>
-makeSyntaxResult(ParserStatus Status, llvm::Optional<Syntax> SyntaxNode,
+makeSyntaxResult(ParserStatus Status, llvm::Optional<Syntax> &&SyntaxNode,
                  AST *ASTNode) {
-  return SyntaxParserResult<Syntax, AST>(Status, SyntaxNode, ASTNode);
+  return SyntaxParserResult<Syntax, AST>(Status, std::move(SyntaxNode), ASTNode);
 }
 
 /// Create a result (null or non-null) with error and code completion bits set.