[pseudo] Share the underly payload when stripping comments for a token stream

[hokein]

Summary

stripComments(cook(...)) is a common pattern being written.
Without this patch, this has a use-after-free issue (cook returns a temporary
TokenStream object which has its own payload, but the payload is not
shared with the one returned by stripComments).

[hokein]

Example: https://github.com/llvm/llvm-project/blob/main/clang-tools-extra/pseudo/tool/ClangPseudo.cpp#L88

[sammccall]

This looks good, but if that's the idiom we should add it to cook() also.

[hokein]

In D125311#3503452, @sammccall wrote:

This looks good, but if that's the idiom we should add it to cook() also.

cook is more tricky, it takes a TokenStream, and returns a new one with an allocator payload:

• the input TokenStream doesn't have a payload, this is OK
• the input TokenStream has a payload, but what if the type of payload is not an allocator?

Possible solution: 1) always assert the input TokenStream doesn't have a payload or the payload is allocator type 2) merge the lex and cook function, it looks like we always use the cook one.

[sammccall]

In D125311#3523388, @hokein wrote:

In D125311#3503452, @sammccall wrote:

This looks good, but if that's the idiom we should add it to cook() also.

cook is more tricky, it takes a TokenStream, and returns a new one with an allocator payload:

• the input TokenStream doesn't have a payload, this is OK
• the input TokenStream has a payload, but what if the type of payload is not an allocator?

Payload is opaque, it can be a pair of {new payload, inner payload}. e.g.

void TokenStream::addPayload(shared_ptr<void> P) {
  Payload = make_shared<pair<shared_ptr<void>, shared_ptr<void>>>(std::move(P), std::move(Payload));
}

[hokein]

add addPayload method, and share payload in cook as well

[hokein]

In D125311#3524190, @sammccall wrote:

In D125311#3523388, @hokein wrote:

In D125311#3503452, @sammccall wrote:

This looks good, but if that's the idiom we should add it to cook() also.

cook is more tricky, it takes a TokenStream, and returns a new one with an allocator payload:

• the input TokenStream doesn't have a payload, this is OK
• the input TokenStream has a payload, but what if the type of payload is not an allocator?

Payload is opaque, it can be a pair of {new payload, inner payload}. e.g.

void TokenStream::addPayload(shared_ptr<void> P) {
  Payload = make_shared<pair<shared_ptr<void>, shared_ptr<void>>>(std::move(P), std::move(Payload));
}

Ohh, right, this is clever!

clang-tools-extra/pseudo/include/clang-pseudo/Token.h

public:
   }
 
   /// May return the end sentinel if the stream is empty.
   const Token &front() const {
     assert(isFinalized());
     return Storage[1];
   }
 
+  /// Returns the shared payload.
+  std::shared_ptr<void> getPayload() const { return Payload; }
+  /// Adds the given payload to the stream.
+  void addPayload(std::shared_ptr<void> P) {
+    if (!Payload)
+      Payload = std::move(P);
+    else
+      Payload = std::make_shared<
+          std::pair<std::shared_ptr<void>, std::shared_ptr<void>>>(
+          std::move(P), std::move(Payload));
+  }
+
   /// Print the tokens in this stream to the output stream.
   ///
   /// The presence of newlines/spaces is preserved, but not the quantity.
   void print(llvm::raw_ostream &) const;
 
 private:
   std::shared_ptr<void> Payload;
 

clang-tools-extra/pseudo/lib/Lex.cpp

TokenStream lex(const std::string &Code, const clang::LangOptions &LangOpts) {
   Result.finalize();
   return Result;
 }
 
 TokenStream cook(const TokenStream &Code, const LangOptions &LangOpts) {
   auto CleanedStorage = std::make_shared<llvm::BumpPtrAllocator>();
   clang::IdentifierTable Identifiers(LangOpts);
   TokenStream Result(CleanedStorage);
-
+  Result.addPayload(Code.getPayload());
   for (auto Tok : Code.tokens()) {
     if (Tok.flag(LexFlags::NeedsCleaning)) {
       // Remove escaped newlines and trigraphs.
       llvm::SmallString<64> CleanBuffer;
       const char *Pos = Tok.text().begin();
       while (Pos < Tok.text().end()) {
         unsigned CharSize = 0;
         CleanBuffer.push_back(

clang-tools-extra/pseudo/lib/Token.cpp

if (Opts.CPlusPlus)
     Opts.MicrosoftExt = true;  // kw__try, kw__finally
   Opts.DeclSpecKeyword = true; // __declspec
   Opts.WChar = true;
 
   return Opts;
 }
 
 TokenStream stripComments(const TokenStream &Input) {
-  TokenStream Out;
+  TokenStream Out(Input.getPayload());
   for (const Token &T : Input.tokens()) {
     if (T.Kind == tok::comment)
       continue;
     Out.push(T);
   }
   Out.finalize();
   return Out;
 }
 
 } // namespace pseudo
 } // namespace clang