[libSyntax] Store raw trivia inside RawSyntax and only lex into pieces when requested

Author: ahoppen

include/swift/Syntax/RawSyntax.h

@@ -150,8 +150,7 @@ typedef unsigned SyntaxNodeId;
 ///
 /// This is implementation detail - do not expose it in public API.
 class RawSyntax final
-    : private llvm::TrailingObjects<RawSyntax, RC<RawSyntax>, OwnedString,
-                                    TriviaPiece> {
+    : private llvm::TrailingObjects<RawSyntax, RC<RawSyntax>, OwnedString> {
   friend TrailingObjects;
 
   /// The ID that shall be used for the next node that is created and does not
@@ -198,10 +197,8 @@ class RawSyntax final
       uint64_t : bitmax(NumRawSyntaxBits, 64); // align to 16 bits
       /// The kind of token this "token" node represents.
       unsigned TokenKind : 16;
-      /// Number of leading  trivia pieces.
-      unsigned NumLeadingTrivia : 16;
-      /// Number of trailing trivia pieces.
-      unsigned NumTrailingTrivia : 16;
+      StringRef LeadingTrivia;
+      StringRef TrailingTrivia;
     } Token;
   } Bits;
 
@@ -211,11 +208,6 @@ class RawSyntax final
   size_t numTrailingObjects(OverloadToken<OwnedString>) const {
     return isToken() ? 1 : 0;
   }
-  size_t numTrailingObjects(OverloadToken<TriviaPiece>) const {
-    return isToken()
-             ? Bits.Token.NumLeadingTrivia + Bits.Token.NumTrailingTrivia
-             : 0;
-  }
 
   /// Constructor for creating layout nodes.
   /// If the node has been allocated inside the bump allocator of a
@@ -232,9 +224,9 @@ class RawSyntax final
   /// If \p NodeId is \c None, the next free NodeId is used, if it is passed,
   /// the caller needs to assure that the NodeId has not been used yet.
   RawSyntax(tok TokKind, OwnedString Text, size_t TextLength,
-            ArrayRef<TriviaPiece> LeadingTrivia,
-            ArrayRef<TriviaPiece> TrailingTrivia, SourcePresence Presence,
-            const RC<SyntaxArena> &Arena, llvm::Optional<SyntaxNodeId> NodeId);
+            StringRef LeadingTrivia, StringRef TrailingTrivia,
+            SourcePresence Presence, const RC<SyntaxArena> &Arena,
+            llvm::Optional<SyntaxNodeId> NodeId);
 
   /// Compute the node's text length by summing up the length of its childern
   size_t computeTextLength() {
@@ -295,27 +287,22 @@ class RawSyntax final
 
   /// Make a raw "token" syntax node.
   static RC<RawSyntax> make(tok TokKind, OwnedString Text, size_t TextLength,
-                            ArrayRef<TriviaPiece> LeadingTrivia,
-                            ArrayRef<TriviaPiece> TrailingTrivia,
+                            StringRef LeadingTrivia, StringRef TrailingTrivia,
                             SourcePresence Presence,
                             const RC<SyntaxArena> &Arena = SyntaxArena::make(),
                             llvm::Optional<SyntaxNodeId> NodeId = llvm::None);
 
   /// Make a raw "token" syntax node that was allocated in \p Arena.
-  static RC<RawSyntax> makeAndCalcLength(
-      tok TokKind, OwnedString Text, ArrayRef<TriviaPiece> LeadingTrivia,
-      ArrayRef<TriviaPiece> TrailingTrivia, SourcePresence Presence,
-      const RC<SyntaxArena> &Arena = SyntaxArena::make(),
-      llvm::Optional<SyntaxNodeId> NodeId = llvm::None) {
+  static RC<RawSyntax>
+  makeAndCalcLength(tok TokKind, OwnedString Text, StringRef LeadingTrivia,
+                    StringRef TrailingTrivia, SourcePresence Presence,
+                    const RC<SyntaxArena> &Arena = SyntaxArena::make(),
+                    llvm::Optional<SyntaxNodeId> NodeId = llvm::None) {
     size_t TextLength = 0;
     if (Presence != SourcePresence::Missing) {
-      for (auto Trivia : LeadingTrivia) {
-        TextLength += Trivia.getTextLength();
-      }
+      TextLength += LeadingTrivia.size();
       TextLength += Text.size();
-      for (auto Trivia : TrailingTrivia) {
-        TextLength += Trivia.getTextLength();
-      }
+      TextLength += TrailingTrivia.size();
     }
     return make(TokKind, Text, TextLength, LeadingTrivia, TrailingTrivia,
                 Presence, Arena, NodeId);
@@ -414,18 +401,28 @@ class RawSyntax final
   /// disappear when the syntax node gets freed.
   StringRef getTokenText() const { return getOwnedTokenText().str(); }
 
-  /// Return the leading trivia list of the token.
-  ArrayRef<TriviaPiece> getLeadingTrivia() const {
+  /// Return the unparsed leading trivia of the token.
+  StringRef getLeadingTrivia() const {
     assert(isToken());
-    return {getTrailingObjects<TriviaPiece>(), Bits.Token.NumLeadingTrivia};
+    return Bits.Token.LeadingTrivia;
   }
-  /// Return the trailing trivia list of the token.
-  ArrayRef<TriviaPiece> getTrailingTrivia() const {
+
+  /// Return the unparsed trailing trivia of the token.
+  StringRef getTrailingTrivia() const {
     assert(isToken());
-    return {getTrailingObjects<TriviaPiece>() + Bits.Token.NumLeadingTrivia,
-            Bits.Token.NumTrailingTrivia};
+    return Bits.Token.TrailingTrivia;
   }
 
+  /// Return pieces that make up the leading trivia of the token.
+  /// Note that this triggers trivia parsing which may be expensive. If the
+  /// trivia pieces are required multiple times, consider caching them.
+  Trivia getLeadingTriviaPieces() const;
+
+  /// Return the trailing trivia list of the token.
+  /// Note that this triggers trivia parsing which may be expensive. If the
+  /// trivia pieces are required multiple times, consider caching them.
+  Trivia getTrailingTriviaPieces() const;
+
   /// Return \c true if this is the given kind of token.
   bool isToken(tok K) const { return isToken() && getTokenKind() == K; }
 
@@ -436,30 +433,20 @@ class RawSyntax final
 
   /// Return a new token like this one, but with the given leading
   /// trivia instead.
-  RC<RawSyntax>
-  withLeadingTrivia(ArrayRef<TriviaPiece> NewLeadingTrivia) const {
+  RC<RawSyntax> withLeadingTrivia(StringRef NewLeadingTrivia) const {
     return makeAndCalcLength(getTokenKind(), getOwnedTokenText(),
                              NewLeadingTrivia, getTrailingTrivia(),
                              getPresence());
   }
 
-  RC<RawSyntax> withLeadingTrivia(Trivia NewLeadingTrivia) const {
-    return withLeadingTrivia(NewLeadingTrivia.Pieces);
-  }
-
   /// Return a new token like this one, but with the given trailing
   /// trivia instead.
-  RC<RawSyntax>
-  withTrailingTrivia(ArrayRef<TriviaPiece> NewTrailingTrivia) const {
+  RC<RawSyntax> withTrailingTrivia(StringRef NewTrailingTrivia) const {
     return makeAndCalcLength(getTokenKind(), getOwnedTokenText(),
                              getLeadingTrivia(), NewTrailingTrivia,
                              getPresence());
   }
 
-  RC<RawSyntax> withTrailingTrivia(Trivia NewTrailingTrivia) const {
-    return withTrailingTrivia(NewTrailingTrivia.Pieces);
-  }
-
   /// @}
 
   /// \name Getter routines for "layout" nodes.
@@ -503,21 +490,9 @@ class RawSyntax final
 
   /// @}
 
-  size_t getLeadingTriviaLength() {
-    size_t Length = 0;
-    for (auto Trivia : getLeadingTrivia()) {
-      Length += Trivia.getTextLength();
-    }
-    return Length;
-  }
+  size_t getLeadingTriviaLength() const { return getLeadingTrivia().size(); }
 
-  size_t getTrailingTriviaLength() {
-    size_t Length = 0;
-    for (auto Trivia : getTrailingTrivia()) {
-      Length += Trivia.getTextLength();
-    }
-    return Length;
-  }
+  size_t getTrailingTriviaLength() const { return getTrailingTrivia().size(); }
 
   /// Print this piece of syntax recursively.
   void print(llvm::raw_ostream &OS, SyntaxPrintOptions Opts) const;
@@ -529,8 +504,7 @@ class RawSyntax final
   void dump(llvm::raw_ostream &OS, unsigned Indent = 0) const;
 
   static void Profile(llvm::FoldingSetNodeID &ID, tok TokKind, OwnedString Text,
-                      ArrayRef<TriviaPiece> LeadingTrivia,
-                      ArrayRef<TriviaPiece> TrailingTrivia);
+                      StringRef LeadingTrivia, StringRef TrailingTrivia);
 };
 
 } // end namespace syntax

include/swift/Syntax/Serialization/SyntaxDeserialization.h

@@ -150,9 +150,9 @@ template <> struct MappingTraits<swift::RC<swift::RawSyntax>> {
     if (description.hasValue) {
       swift::tok tokenKind = description.Kind;
       StringRef text = description.Text;
-      std::vector<swift::TriviaPiece> leadingTrivia;
+      StringRef leadingTrivia;
       in.mapRequired("leadingTrivia", leadingTrivia);
-      std::vector<swift::TriviaPiece> trailingTrivia;
+      StringRef trailingTrivia;
       in.mapRequired("trailingTrivia", trailingTrivia);
       swift::SourcePresence presence;
       in.mapRequired("presence", presence);

include/swift/Syntax/SyntaxFactory.h.gyb

@@ -46,8 +46,8 @@ class SyntaxArena;
 struct SyntaxFactory {
   /// Make any kind of token.
   static TokenSyntax makeToken(tok Kind,
-      OwnedString Text, const Trivia &LeadingTrivia,
-      const Trivia &TrailingTrivia, SourcePresence Presence,
+      OwnedString Text, StringRef LeadingTrivia, StringRef TrailingTrivia,
+      SourcePresence Presence,
       const RC<SyntaxArena> &Arena = SyntaxArena::make()
   );
 
@@ -100,16 +100,16 @@ struct SyntaxFactory {
 
 % for token in SYNTAX_TOKENS:
 %   if token.is_keyword:
-  static TokenSyntax make${token.name}Keyword(const Trivia &LeadingTrivia,
-      const Trivia &TrailingTrivia,
+  static TokenSyntax make${token.name}Keyword(StringRef LeadingTrivia,
+      StringRef TrailingTrivia,
       const RC<SyntaxArena> &Arena = SyntaxArena::make());
 %   elif token.text:
-  static TokenSyntax make${token.name}Token(const Trivia &LeadingTrivia,
-      const Trivia &TrailingTrivia,
+  static TokenSyntax make${token.name}Token(StringRef LeadingTrivia,
+            StringRef TrailingTrivia,
       const RC<SyntaxArena> &Arena = SyntaxArena::make());
 %   else:
   static TokenSyntax make${token.name}(OwnedString Text,
-      const Trivia &LeadingTrivia, const Trivia &TrailingTrivia,
+            StringRef LeadingTrivia, StringRef TrailingTrivia,
       const RC<SyntaxArena> &Arena = SyntaxArena::make());
 %   end
 % end
@@ -140,7 +140,7 @@ struct SyntaxFactory {
   /// Creates a TypeIdentifierSyntax with the provided name and leading/trailing
   /// trivia.
   static TypeSyntax makeTypeIdentifier(OwnedString TypeName,
-      const Trivia &LeadingTrivia = {}, const Trivia &TrailingTrivia = {},
+      StringRef LeadingTrivia = {}, StringRef TrailingTrivia = {},
       const RC<SyntaxArena> &Arena = SyntaxArena::make()
   );
 
@@ -153,32 +153,32 @@ struct SyntaxFactory {
   );
 
   /// Creates a TypeIdentifierSyntax for the `Any` type.
-  static TypeSyntax makeAnyTypeIdentifier(const Trivia &LeadingTrivia = {},
-      const Trivia &TrailingTrivia = {},
+  static TypeSyntax makeAnyTypeIdentifier(StringRef LeadingTrivia = {},
+      StringRef TrailingTrivia = {},
       const RC<SyntaxArena> &Arena = SyntaxArena::make()
   );
 
   /// Creates a TypeIdentifierSyntax for the `Self` type.
-  static TypeSyntax makeSelfTypeIdentifier(const Trivia &LeadingTrivia = {},
-      const Trivia &TrailingTrivia = {},
+  static TypeSyntax makeSelfTypeIdentifier(StringRef LeadingTrivia = {},
+      StringRef TrailingTrivia = {},
       const RC<SyntaxArena> &Arena = SyntaxArena::make()
   );
 
   /// Creates a TokenSyntax for the `Type` identifier.
-  static TokenSyntax makeTypeToken(const Trivia &LeadingTrivia = {},
-      const Trivia &TrailingTrivia = {},
+  static TokenSyntax makeTypeToken(StringRef LeadingTrivia = {},
+      StringRef TrailingTrivia = {},
       const RC<SyntaxArena> &Arena = SyntaxArena::make()
   );
 
   /// Creates a TokenSyntax for the `Protocol` identifier.
-  static TokenSyntax makeProtocolToken(const Trivia &LeadingTrivia = {},
-      const Trivia &TrailingTrivia = {},
+  static TokenSyntax makeProtocolToken(StringRef LeadingTrivia = {},
+      StringRef TrailingTrivia = {},
       const RC<SyntaxArena> &Arena = SyntaxArena::make()
   );
 
   /// Creates an `==` operator token.
-  static TokenSyntax makeEqualityOperator(const Trivia &LeadingTrivia = {},
-      const Trivia &TrailingTrivia = {},
+  static TokenSyntax makeEqualityOperator(StringRef LeadingTrivia = {},
+      StringRef TrailingTrivia = {},
       const RC<SyntaxArena> &Arena = SyntaxArena::make()
   );
 

include/swift/Syntax/TokenSyntax.h

@@ -39,22 +39,22 @@ class TokenSyntax final : public Syntax {
     return makeRoot<TokenSyntax>(RawSyntax::missing(Kind, Text));
   }
 
-  Trivia getLeadingTrivia() const {
-    return Trivia { getRaw()->getLeadingTrivia().vec() };
-  }
+  StringRef getLeadingTrivia() const { return getRaw()->getLeadingTrivia(); }
+  Trivia getLeadingTriviaPieces() const { return getRaw()->getLeadingTriviaPieces(); }
 
-  Trivia getTrailingTrivia() const {
-    return Trivia { getRaw()->getTrailingTrivia().vec() };
+  StringRef getTrailingTrivia() const { return getRaw()->getTrailingTrivia(); }
+  Trivia getTrailingTriviaPieces() const {
+    return getRaw()->getTrailingTriviaPieces();
   }
 
-  TokenSyntax withLeadingTrivia(const Trivia &Trivia) const {
-    auto NewRaw = getRaw()->withLeadingTrivia(Trivia.Pieces);
-    return TokenSyntax(Data.replacingSelf(NewRaw));
+  TokenSyntax withLeadingTrivia(StringRef Trivia) const {
+    auto NewRaw = getRaw()->withLeadingTrivia(Trivia);
+    return TokenSyntax(getData().replacingSelf(NewRaw));
   }
 
-  TokenSyntax withTrailingTrivia(const Trivia &Trivia) const {
-    auto NewRaw = getRaw()->withTrailingTrivia(Trivia.Pieces);
-    return TokenSyntax(Data.replacingSelf(NewRaw));
+  TokenSyntax withTrailingTrivia(StringRef Trivia) const {
+    auto NewRaw = getRaw()->withTrailingTrivia(Trivia);
+    return TokenSyntax(getData().replacingSelf(NewRaw));
   }
 
   /* TODO: If we really need them.

lib/Parse/Parser.cpp

@@ -84,7 +84,8 @@ void tokenize(const LangOptions &LangOpts, const SourceManager &SM,
     if (F != ResetTokens.end()) {
       assert(F->isNot(tok::string_literal));
 
-      DestFunc(*F, ParsedTrivia(), ParsedTrivia());
+      DestFunc(*F, /*LeadingTrivia=*/StringRef(),
+               /*TrailingTrivia=*/StringRef());
 
       auto NewState = L.getStateForBeginningOfTokenLoc(
           F->getLoc().getAdvancedLoc(F->getLength()));
@@ -96,12 +97,11 @@ void tokenize(const LangOptions &LangOpts, const SourceManager &SM,
       std::vector<Token> StrTokens;
       getStringPartTokens(Tok, LangOpts, SM, BufferID, StrTokens);
       for (auto &StrTok : StrTokens) {
-        DestFunc(StrTok, ParsedTrivia(), ParsedTrivia());
+        DestFunc(StrTok, /*LeadingTrivia=*/StringRef(),
+                 /*TrailingTrivia=*/StringRef());
       }
     } else {
-      auto LeadingTriviaPieces = TriviaLexer::lexTrivia(LeadingTrivia);
-      auto TrailingTriviaPieces = TriviaLexer::lexTrivia(TrailingTrivia);
-      DestFunc(Tok, LeadingTriviaPieces, TrailingTriviaPieces);
+      DestFunc(Tok, LeadingTrivia, TrailingTrivia);
     }
 
   } while (Tok.getKind() != tok::eof);
@@ -308,14 +308,13 @@ std::vector<Token> swift::tokenize(const LangOptions &LangOpts,
                                    ArrayRef<Token> SplitTokens) {
   std::vector<Token> Tokens;
 
-  tokenize(LangOpts, SM, BufferID, Offset, EndOffset,
-           Diags,
+  tokenize(LangOpts, SM, BufferID, Offset, EndOffset, Diags,
            KeepComments ? CommentRetentionMode::ReturnAsTokens
                         : CommentRetentionMode::AttachToNextToken,
            TriviaRetentionMode::WithoutTrivia, TokenizeInterpolatedString,
            SplitTokens,
-           [&](const Token &Tok, const ParsedTrivia &LeadingTrivia,
-               const ParsedTrivia &TrailingTrivia) { Tokens.push_back(Tok); });
+           [&](const Token &Tok, StringRef LeadingTrivia,
+               StringRef TrailingTrivia) { Tokens.push_back(Tok); });
 
   assert(Tokens.back().is(tok::eof));
   Tokens.pop_back(); // Remove EOF.
@@ -335,24 +334,14 @@ swift::tokenizeWithTrivia(const LangOptions &LangOpts, const SourceManager &SM,
       CommentRetentionMode::AttachToNextToken, TriviaRetentionMode::WithTrivia,
       /*TokenizeInterpolatedString=*/false,
       /*SplitTokens=*/ArrayRef<Token>(),
-      [&](const Token &Tok, const ParsedTrivia &LeadingTrivia,
-          const ParsedTrivia &TrailingTrivia) {
+      [&](const Token &Tok, StringRef LeadingTrivia, StringRef TrailingTrivia) {
         CharSourceRange TokRange = Tok.getRange();
-        SourceLoc LeadingTriviaLoc =
-          TokRange.getStart().getAdvancedLoc(-LeadingTrivia.getLength());
-        SourceLoc TrailingTriviaLoc =
-          TokRange.getStart().getAdvancedLoc(TokRange.getByteLength());
-        Trivia syntaxLeadingTrivia =
-          LeadingTrivia.convertToSyntaxTrivia(LeadingTriviaLoc, SM, BufferID);
-        Trivia syntaxTrailingTrivia =
-          TrailingTrivia.convertToSyntaxTrivia(TrailingTriviaLoc, SM, BufferID);
         auto Text = OwnedString::makeRefCounted(Tok.getRawText());
-        size_t TextLength = LeadingTrivia.getLength() +
-                            TokRange.getByteLength() +
-                            TrailingTrivia.getLength();
-        auto ThisToken = RawSyntax::make(
-            Tok.getKind(), Text, TextLength, syntaxLeadingTrivia.Pieces,
-            syntaxTrailingTrivia.Pieces, SourcePresence::Present);
+        size_t TextLength = LeadingTrivia.size() + TokRange.getByteLength() +
+                            TrailingTrivia.size();
+        auto ThisToken =
+            RawSyntax::make(Tok.getKind(), Text, TextLength, LeadingTrivia,
+                            TrailingTrivia, SourcePresence::Present);
 
         auto ThisTokenPos =
             RunningPos.advancedBy(ThisToken->getLeadingTriviaLength());