[clangd] Pass DraftStore to Tweak Apply and Prepare.

[njames93]

Summary

This enables multi file edit tweaks to take effect even if files are unsaved.

[njames93]

Abandoning in favour of VFS approach

clang-tools-extra/clangd/ClangdLSPServer.cpp

auto Action = [this, ApplyEdit, Reply = std::move(Reply),
         (*WE.changes)[URI::createFile(It.first()).toString()] =
             It.second.asTextEdits();
       }
       // ApplyEdit will take care of calling Reply().
       return ApplyEdit(std::move(WE), "Tweak applied.", std::move(Reply));
     };
     Server->applyTweak(Params.tweakArgs->file.file(),
                        Params.tweakArgs->selection, Params.tweakArgs->tweakID,
-                       std::move(Action));
+                       DraftMgr, std::move(Action));
   } else {
     // We should not get here because ExecuteCommandParams would not have
     // parsed in the first place and this handler should not be called. But if
     // more commands are added, this will be here has a safe guard.
     Reply(llvm::make_error<LSPError>(
         llvm::formatv("Unsupported command \"{0}\".", Params.command).str(),
         ErrorCode::InvalidParams));
   }
auto ConsumeActions =
         std::vector<Command> Commands;
         for (const auto &Action : Actions) {
           if (auto Command = asCommand(Action))
             Commands.push_back(std::move(*Command));
         }
         return Reply(llvm::json::Array(Commands));
       };
   Server->enumerateTweaks(
-      File.file(), Params.range,
+      File.file(), Params.range, DraftMgr,
       [this, KindAllowed(std::move(KindAllowed))](const Tweak &T) {
         return Opts.TweakFilter(T) && KindAllowed(T.kind());
       },
       std::move(ConsumeActions));
 }
 
 void ClangdLSPServer::onCompletion(const CompletionParams &Params,
                                    Callback<CompletionList> Reply) {

clang-tools-extra/clangd/ClangdServer.h

public:
 
   struct TweakRef {
     std::string ID;    /// ID to pass for applyTweak.
     std::string Title; /// A single-line message to show in the UI.
     llvm::StringLiteral Kind;
   };
   /// Enumerate the code tweaks available to the user at a specified point.
   /// Tweaks where Filter returns false will not be checked or included.
-  void enumerateTweaks(PathRef File, Range Sel,
+  void enumerateTweaks(PathRef File, Range Sel, const DraftStore &DraftMgr,
                        llvm::unique_function<bool(const Tweak &)> Filter,
                        Callback<std::vector<TweakRef>> CB);
 
   /// Apply the code tweak with a specified \p ID.
   void applyTweak(PathRef File, Range Sel, StringRef ID,
-                  Callback<Tweak::Effect> CB);
+                  const DraftStore &DraftMgr, Callback<Tweak::Effect> CB);
 
   /// Called when an event occurs for a watched file in the workspace.
   void onFileEvent(const DidChangeWatchedFilesParams &Params);
 
   /// Get symbol info for given position.
   /// Clangd extension - not part of official LSP.
   void symbolInfo(PathRef File, Position Pos,
                   Callback<std::vector<SymbolDetails>> CB);

clang-tools-extra/clangd/ClangdServer.cpp

SelectionTree::createEach(
             AST.Inputs.Index, AST.AST, *Begin, *End, std::move(T)));
         return false;
       });
   assert(!Result.empty() && "Expected at least one SelectionTree");
   return std::move(Result);
 }
 
 void ClangdServer::enumerateTweaks(
-    PathRef File, Range Sel, llvm::unique_function<bool(const Tweak &)> Filter,
+    PathRef File, Range Sel, const DraftStore &DraftMgr,
+    llvm::unique_function<bool(const Tweak &)> Filter,
     Callback<std::vector<TweakRef>> CB) {
   // Tracks number of times a tweak has been offered.
   static constexpr trace::Metric TweakAvailable(
       "tweak_available", trace::Metric::Counter, "tweak_id");
-  auto Action = [File = File.str(), Sel, CB = std::move(CB),
+  auto Action = [File = File.str(), Sel, &DraftMgr, CB = std::move(CB),
                  Filter =
                      std::move(Filter)](Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     auto Selections = tweakSelection(Sel, *InpAST);
     if (!Selections)
       return CB(Selections.takeError());
     std::vector<TweakRef> Res;
     // Don't allow a tweak to fire more than once across ambiguous selections.
     llvm::DenseSet<llvm::StringRef> PreparedTweaks;
     auto DeduplicatingFilter = [&](const Tweak &T) {
       return Filter(T) && !PreparedTweaks.count(T.id());
     };
     for (const auto &Sel : *Selections) {
-      for (auto &T : prepareTweaks(*Sel, DeduplicatingFilter)) {
+      for (auto &T : prepareTweaks(*Sel, DraftMgr, DeduplicatingFilter)) {
         Res.push_back({T->id(), T->title(), T->kind()});
         PreparedTweaks.insert(T->id());
         TweakAvailable.record(1, T->id());
       }
     }
 
     CB(std::move(Res));
   };
 
   WorkScheduler.runWithAST("EnumerateTweaks", File, std::move(Action),
                            TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::applyTweak(PathRef File, Range Sel, StringRef TweakID,
+                              const DraftStore &DraftMgr,
                               Callback<Tweak::Effect> CB) {
   // Tracks number of times a tweak has been attempted.
   static constexpr trace::Metric TweakAttempt(
       "tweak_attempt", trace::Metric::Counter, "tweak_id");
   // Tracks number of times a tweak has failed to produce edits.
   static constexpr trace::Metric TweakFailed(
       "tweak_failed", trace::Metric::Counter, "tweak_id");
   TweakAttempt.record(1, TweakID);
   auto Action = [File = File.str(), Sel, TweakID = TweakID.str(),
-                 CB = std::move(CB),
+                 CB = std::move(CB), &DraftMgr,
                  this](Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     auto Selections = tweakSelection(Sel, *InpAST);
     if (!Selections)
       return CB(Selections.takeError());
     llvm::Optional<llvm::Expected<Tweak::Effect>> Effect;
     // Try each selection, take the first one that prepare()s.
     // If they all fail, Effect will hold get the last error.
     for (const auto &Selection : *Selections) {
-      auto T = prepareTweak(TweakID, *Selection);
+      auto T = prepareTweak(TweakID, *Selection, DraftMgr);
       if (T) {
-        Effect = (*T)->apply(*Selection);
+        Effect = (*T)->apply(*Selection, DraftMgr);
         break;
       }
       Effect = T.takeError();
     }
     assert(Effect.hasValue() && "Expected at least one selection");
     if (*Effect) {
       // Tweaks don't apply clang-format, do that centrally here.
       for (auto &It : (*Effect)->ApplyEdits) {

clang-tools-extra/clangd/refactor/Tweak.h

 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include <string>
 
 namespace clang {
 namespace clangd {
 
+class DraftStore;
+
 /// An interface base for small context-sensitive refactoring actions.
 /// To implement a new tweak use the following pattern in a .cpp file:
 ///   class MyTweak : public Tweak {
 ///   public:
 ///     const char* id() const override final; // defined by REGISTER_TWEAK.
 ///     // implement other methods here.
 ///   };
 ///   REGISTER_TWEAK(MyTweak);
public:
   /// REGISTER_TWEAK.
   virtual const char *id() const = 0;
   /// Run the first stage of the action. Returns true indicating that the
   /// action is available and should be shown to the user. Returns false if the
   /// action is not available.
   /// This function should be fast, if the action requires non-trivial work it
   /// should be moved into 'apply'.
   /// Returns true iff the action is available and apply() can be called on it.
-  virtual bool prepare(const Selection &Sel) = 0;
+  virtual bool prepare(const Selection &Sel, const DraftStore &DraftMgr) = 0;
   /// Run the second stage of the action that would produce the actual effect.
   /// EXPECTS: prepare() was called and returned true.
-  virtual Expected<Effect> apply(const Selection &Sel) = 0;
+  virtual Expected<Effect> apply(const Selection &Sel,
+                                 const DraftStore &DraftMgr) = 0;
 
   /// A one-line title of the action that should be shown to the users in the
   /// UI.
   /// EXPECTS: prepare() was called and returned true.
   virtual std::string title() const = 0;
   /// Describes what kind of action this is.
   /// EXPECTS: prepare() was called and returned true.
   virtual llvm::StringLiteral kind() const = 0;
   /// Is this a 'hidden' tweak, which are off by default.
   virtual bool hidden() const { return false; }
 };
 
 // All tweaks must be registered in the .cpp file next to their definition.
 #define REGISTER_TWEAK(Subclass)                                               \
   ::llvm::Registry<::clang::clangd::Tweak>::Add<Subclass>                      \
       TweakRegistrationFor##Subclass(#Subclass, /*Description=*/"");           \
   const char *Subclass::id() const { return #Subclass; }
 
 /// Calls prepare() on all tweaks that satisfy the filter, returning those that
 /// can run on the selection.
 std::vector<std::unique_ptr<Tweak>>
-prepareTweaks(const Tweak::Selection &S,
+prepareTweaks(const Tweak::Selection &S, const DraftStore &DraftMgr,
               llvm::function_ref<bool(const Tweak &)> Filter);
 
 // Calls prepare() on the tweak with a given ID.
 // If prepare() returns false, returns an error.
 // If prepare() returns true, returns the corresponding tweak.
 llvm::Expected<std::unique_ptr<Tweak>> prepareTweak(StringRef TweakID,
-                                                    const Tweak::Selection &S);
+                                                    const Tweak::Selection &S,
+                                                    const DraftStore &DraftMgr);
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/clangd/refactor/Tweak.cpp

Tweak::Selection::Selection(const SymbolIndex *Index, ParsedAST &AST,
     : Index(Index), AST(&AST), SelectionBegin(RangeBegin),
       SelectionEnd(RangeEnd), ASTSelection(std::move(ASTSelection)) {
   auto &SM = AST.getSourceManager();
   Code = SM.getBufferData(SM.getMainFileID());
   Cursor = SM.getComposedLoc(SM.getMainFileID(), RangeBegin);
 }
 
 std::vector<std::unique_ptr<Tweak>>
-prepareTweaks(const Tweak::Selection &S,
+prepareTweaks(const Tweak::Selection &S, const DraftStore &DraftMgr,
               llvm::function_ref<bool(const Tweak &)> Filter) {
   validateRegistry();
 
   std::vector<std::unique_ptr<Tweak>> Available;
   for (const auto &E : TweakRegistry::entries()) {
     std::unique_ptr<Tweak> T = E.instantiate();
-    if (!Filter(*T) || !T->prepare(S))
+    if (!Filter(*T) || !T->prepare(S, DraftMgr))
       continue;
     Available.push_back(std::move(T));
   }
   // Ensure deterministic order of the results.
   llvm::sort(Available,
              [](const std::unique_ptr<Tweak> &L,
                 const std::unique_ptr<Tweak> &R) { return L->id() < R->id(); });
   return Available;
 }
 
-llvm::Expected<std::unique_ptr<Tweak>> prepareTweak(StringRef ID,
-                                                    const Tweak::Selection &S) {
+llvm::Expected<std::unique_ptr<Tweak>>
+prepareTweak(StringRef ID, const Tweak::Selection &S,
+             const DraftStore &DraftMgr) {
   auto It = llvm::find_if(
       TweakRegistry::entries(),
       [ID](const TweakRegistry::entry &E) { return E.getName() == ID; });
   if (It == TweakRegistry::end())
     return error("tweak ID {0} is invalid", ID);
   std::unique_ptr<Tweak> T = It->instantiate();
-  if (!T->prepare(S))
+  if (!T->prepare(S, DraftMgr))
     return error("failed to prepare() tweak {0}", ID);
   return std::move(T);
 }
 
 llvm::Expected<std::pair<Path, Edit>>
 Tweak::Effect::fileEdit(const SourceManager &SM, FileID FID,
                         tooling::Replacements Replacements) {
   Edit Ed(SM.getBufferData(FID), std::move(Replacements));

clang-tools-extra/clangd/refactor/tweaks/AddUsing.cpp

 //
 // Currently this only removes qualifier from under the cursor. In the future,
 // we should improve this to remove qualifier from all occurrences of this
 // symbol.
 class AddUsing : public Tweak {
 public:
   const char *id() const override;
 
-  bool prepare(const Selection &Inputs) override;
-  Expected<Effect> apply(const Selection &Inputs) override;
+  bool prepare(const Selection &Inputs, const DraftStore &) override;
+  Expected<Effect> apply(const Selection &Inputs, const DraftStore &) override;
   std::string title() const override;
   llvm::StringLiteral kind() const override {
     return CodeAction::REFACTOR_KIND;
   }
 
 private:
   // All of the following are set by prepare().
   // The qualifier to remove.

 std::string getNNSLAsString(NestedNameSpecifierLoc &NNSL,
                             const PrintingPolicy &Policy) {
   std::string Out;
   llvm::raw_string_ostream OutStream(Out);
   NNSL.getNestedNameSpecifier()->print(OutStream, Policy);
   return OutStream.str();
 }
 
-bool AddUsing::prepare(const Selection &Inputs) {
+bool AddUsing::prepare(const Selection &Inputs, const DraftStore &) {
   auto &SM = Inputs.AST->getSourceManager();
   const auto &TB = Inputs.AST->getTokens();
 
   // Do not suggest "using" in header files. That way madness lies.
   if (isHeaderFile(SM.getFileEntryForID(SM.getMainFileID())->getName(),
                    Inputs.AST->getLangOpts()))
     return false;
 
if (SM.isMacroBodyExpansion(QualifierToRemove.getBeginLoc()) ||
       !SM.isWrittenInSameFile(QualifierToRemove.getBeginLoc(),
                               QualifierToRemove.getEndLoc())) {
     return false;
   }
 
   return true;
 }
 
-Expected<Tweak::Effect> AddUsing::apply(const Selection &Inputs) {
+Expected<Tweak::Effect> AddUsing::apply(const Selection &Inputs,
+                                        const DraftStore &) {
   auto &SM = Inputs.AST->getSourceManager();
 
   std::string QualifierToRemoveStr = getNNSLAsString(
       QualifierToRemove, Inputs.AST->getASTContext().getPrintingPolicy());
   tooling::Replacements R;
   if (auto Err = R.add(tooling::Replacement(
           SM, SM.getSpellingLoc(QualifierToRemove.getBeginLoc()),
           QualifierToRemoveStr.length(), ""))) {

clang-tools-extra/clangd/refactor/tweaks/AnnotateHighlightings.cpp

 ///   void f() { int abc; }
 ///   ^^^^^^^^^^^^^^^^^^^^^
 /// After:
 ///   void /* entity.name.function.cpp */ f() { int /* variable.cpp */ abc; }
 class AnnotateHighlightings : public Tweak {
 public:
   const char *id() const override final;
 
-  bool prepare(const Selection &Inputs) override { return true; }
-  Expected<Effect> apply(const Selection &Inputs) override;
+  bool prepare(const Selection &Inputs, const DraftStore &) override {
+    return true;
+  }
+  Expected<Effect> apply(const Selection &Inputs, const DraftStore &) override;
 
   std::string title() const override { return "Annotate highlighting tokens"; }
   llvm::StringLiteral kind() const override {
     return CodeAction::REFACTOR_KIND;
   }
   bool hidden() const override { return true; }
 };
 REGISTER_TWEAK(AnnotateHighlightings)
 
-Expected<Tweak::Effect> AnnotateHighlightings::apply(const Selection &Inputs) {
+Expected<Tweak::Effect> AnnotateHighlightings::apply(const Selection &Inputs,
+                                                     const DraftStore &) {
   const Decl *CommonDecl = nullptr;
   for (auto N = Inputs.ASTSelection.commonAncestor(); N && !CommonDecl;
        N = N->Parent)
     CommonDecl = N->ASTNode.get<Decl>();
 
   std::vector<HighlightingToken> HighlightingTokens;
   if (!CommonDecl) {
     // Now we hit the TUDecl case where commonAncestor() returns null

clang-tools-extra/clangd/refactor/tweaks/DefineInline.cpp

llvm::StringLiteral kind() const override {
     return CodeAction::REFACTOR_KIND;
   }
   std::string title() const override {
     return "Move function body to declaration";
   }
 
   // Returns true when selection is on a function definition that does not
   // make use of any internal symbols.
-  bool prepare(const Selection &Sel) override {
+  bool prepare(const Selection &Sel, const DraftStore &) override {
     const SelectionTree::Node *SelNode = Sel.ASTSelection.commonAncestor();
     if (!SelNode)
       return false;
     Source = getSelectedFunction(SelNode);
     if (!Source || !Source->hasBody())
       return false;
     // Only the last level of template parameter locations are not kept in AST,
     // so if we are inlining a method that is in a templated class, there is no
bool prepare(const Selection &Sel, const DraftStore &) override {
     // declaration location.
     if (!checkDeclsAreVisible(getNonLocalDeclRefs(*Sel.AST, Source), Target,
                               Sel.AST->getSourceManager()))
       return false;
 
     return true;
   }
 
-  Expected<Effect> apply(const Selection &Sel) override {
+  Expected<Effect> apply(const Selection &Sel, const DraftStore &) override {
     const auto &AST = Sel.AST->getASTContext();
     const auto &SM = AST.getSourceManager();
 
     auto Semicolon = getSemicolonForDecl(Target);
     if (!Semicolon)
       return error("Couldn't find semicolon for target declaration.");
 
     auto AddInlineIfNecessary = addInlineIfInHeader(Target);

clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp

public:
   bool hidden() const override { return false; }
   llvm::StringLiteral kind() const override {
     return CodeAction::REFACTOR_KIND;
   }
   std::string title() const override {
     return "Move function body to out-of-line";
   }
 
-  bool prepare(const Selection &Sel) override {
+  bool prepare(const Selection &Sel, const DraftStore &) override {
     // Bail out if we are not in a header file.
     // FIXME: We might want to consider moving method definitions below class
     // definition even if we are inside a source file.
     if (!isHeaderFile(Sel.AST->getSourceManager().getFilename(Sel.Cursor),
                       Sel.AST->getLangOpts()))
       return false;
 
     Source = getSelectedFunction(Sel.ASTSelection.commonAncestor());
bool prepare(const Selection &Sel, const DraftStore &) override {
     }
 
     // Note that we don't check whether an implementation file exists or not in
     // the prepare, since performing disk IO on each prepare request might be
     // expensive.
     return true;
   }
 
-  Expected<Effect> apply(const Selection &Sel) override {
+  Expected<Effect> apply(const Selection &Sel, const DraftStore &) override {
     const SourceManager &SM = Sel.AST->getSourceManager();
     auto MainFileName =
         getCanonicalPath(SM.getFileEntryForID(SM.getMainFileID()), SM);
     if (!MainFileName)
       return error("Couldn't get absolute path for main file.");
 
     auto CCFile = getSourceFile(*MainFileName, Sel);
     if (!CCFile)

clang-tools-extra/clangd/refactor/tweaks/DumpAST.cpp

 /// Message:
 ///   CallExpr
 ///   |-DeclRefExpr fcall
 ///   `-StringLiteral "foo"
 class DumpAST : public Tweak {
 public:
   const char *id() const override final;
 
-  bool prepare(const Selection &Inputs) override {
+  bool prepare(const Selection &Inputs, const DraftStore &) override {
     for (auto N = Inputs.ASTSelection.commonAncestor(); N && !Node;
          N = N->Parent)
       if (dumpable(N->ASTNode))
         Node = N->ASTNode;
     return Node.hasValue();
   }
-  Expected<Effect> apply(const Selection &Inputs) override;
+  Expected<Effect> apply(const Selection &Inputs, const DraftStore &) override;
   std::string title() const override {
     return std::string(
         llvm::formatv("Dump {0} AST", Node->getNodeKind().asStringRef()));
   }
   llvm::StringLiteral kind() const override { return CodeAction::INFO_KIND; }
   bool hidden() const override { return true; }
 
 private:
   static bool dumpable(const DynTypedNode &N) {
     // Sadly not all node types can be dumped, and there's no API to check.
     // See DynTypedNode::dump().
     return N.get<Decl>() || N.get<Stmt>() || N.get<Type>();
   }
 
   llvm::Optional<DynTypedNode> Node;
 };
 REGISTER_TWEAK(DumpAST)
 
-llvm::Expected<Tweak::Effect> DumpAST::apply(const Selection &Inputs) {
+llvm::Expected<Tweak::Effect> DumpAST::apply(const Selection &Inputs,
+                                             const DraftStore &) {
   std::string Str;
   llvm::raw_string_ostream OS(Str);
   Node->dump(OS, Inputs.AST->getASTContext());
   return Effect::showMessage(std::move(OS.str()));
 }
 
 /// Dumps the SelectionTree.
 /// Input:

 ///        ImplicitCastExpr fcall
 ///         .DeclRefExpr fcall
 ///        BinaryOperator 2 + 2
 ///          *IntegerLiteral 2
 class ShowSelectionTree : public Tweak {
 public:
   const char *id() const override final;
 
-  bool prepare(const Selection &Inputs) override { return true; }
-  Expected<Effect> apply(const Selection &Inputs) override {
+  bool prepare(const Selection &Inputs, const DraftStore &) override {
+    return true;
+  }
+  Expected<Effect> apply(const Selection &Inputs, const DraftStore &) override {
     return Effect::showMessage(llvm::to_string(Inputs.ASTSelection));
   }
   std::string title() const override { return "Show selection tree"; }
   llvm::StringLiteral kind() const override { return CodeAction::INFO_KIND; }
   bool hidden() const override { return true; }
 };
 REGISTER_TWEAK(ShowSelectionTree)
 
 /// Dumps the symbol under the cursor.
 /// Inputs:
 /// void foo();
 ///      ^^^
 /// Message:
 ///  foo -
 ///  {"containerName":null,"id":"CA2EBE44A1D76D2A","name":"foo","usr":"c:@F@foo#"}
 class DumpSymbol : public Tweak {
   const char *id() const override final;
-  bool prepare(const Selection &Inputs) override { return true; }
-  Expected<Effect> apply(const Selection &Inputs) override {
+  bool prepare(const Selection &Inputs, const DraftStore &) override {
+    return true;
+  }
+  Expected<Effect> apply(const Selection &Inputs, const DraftStore &) override {
     std::string Storage;
     llvm::raw_string_ostream Out(Storage);
 
     for (auto &Sym : getSymbolInfo(
              *Inputs.AST, sourceLocToPosition(Inputs.AST->getSourceManager(),
                                               Inputs.Cursor)))
       Out << Sym;
     return Effect::showMessage(Out.str());

 ///        0 | struct S
 ///        0 |   int foo
 ///          | [sizeof=4, dsize=4, align=4,
 ///          |  nvsize=4, nvalign=4]
 class DumpRecordLayout : public Tweak {
 public:
   const char *id() const override final;
 
-  bool prepare(const Selection &Inputs) override {
+  bool prepare(const Selection &Inputs, const DraftStore &) override {
     if (auto *Node = Inputs.ASTSelection.commonAncestor())
       if (auto *D = Node->ASTNode.get<Decl>())
         Record = dyn_cast<RecordDecl>(D);
     return Record && Record->isThisDeclarationADefinition() &&
            !Record->isDependentType();
   }
-  Expected<Effect> apply(const Selection &Inputs) override {
+  Expected<Effect> apply(const Selection &Inputs, const DraftStore &) override {
     std::string Str;
     llvm::raw_string_ostream OS(Str);
     Inputs.AST->getASTContext().DumpRecordLayout(Record, OS);
     return Effect::showMessage(std::move(OS.str()));
   }
   std::string title() const override {
     return std::string(llvm::formatv(
         "Show {0} layout",

clang-tools-extra/clangd/refactor/tweaks/ExpandAutoType.cpp

 ///    ^^^^^^^
 /// FIXME: Handle decltype as well
 class ExpandAutoType : public Tweak {
 public:
   const char *id() const final;
   llvm::StringLiteral kind() const override {
     return CodeAction::REFACTOR_KIND;
   }
-  bool prepare(const Selection &Inputs) override;
-  Expected<Effect> apply(const Selection &Inputs) override;
+  bool prepare(const Selection &Inputs, const DraftStore &) override;
+  Expected<Effect> apply(const Selection &Inputs, const DraftStore &) override;
   std::string title() const override;
 
 private:
   /// Cache the AutoTypeLoc, so that we do not need to search twice.
   llvm::Optional<clang::AutoTypeLoc> CachedLocation;
 };
 
 REGISTER_TWEAK(ExpandAutoType)
if (const auto *DD = It->ASTNode.get<DeclaratorDecl>()) {
         if (auto *RD = DD->getType()->getAsRecordDecl())
           return RD->isLambda();
       }
     }
   }
   return false;
 }
 
-bool ExpandAutoType::prepare(const Selection& Inputs) {
+bool ExpandAutoType::prepare(const Selection &Inputs, const DraftStore &) {
   CachedLocation = llvm::None;
   if (auto *Node = Inputs.ASTSelection.commonAncestor()) {
     if (auto *TypeNode = Node->ASTNode.get<TypeLoc>()) {
       if (const AutoTypeLoc Result = TypeNode->getAs<AutoTypeLoc>()) {
         // Code in apply() does handle 'decltype(auto)' yet.
         if (!Result.getTypePtr()->isDecltypeAuto() &&
             !isStructuredBindingType(Node) &&
             !isDeducedAsLambda(Node, Result.getBeginLoc()))
           CachedLocation = Result;
       }
     }
   }
 
   return (bool) CachedLocation;
 }
 
-Expected<Tweak::Effect> ExpandAutoType::apply(const Selection& Inputs) {
+Expected<Tweak::Effect> ExpandAutoType::apply(const Selection &Inputs,
+                                              const DraftStore &) {
   auto &SrcMgr = Inputs.AST->getSourceManager();
 
   llvm::Optional<clang::QualType> DeducedType = getDeducedType(
       Inputs.AST->getASTContext(), CachedLocation->getBeginLoc());
 
   // if we can't resolve the type, return an error message
   if (DeducedType == llvm::None || (*DeducedType)->isUndeducedAutoType())
     return error("Could not deduce type for 'auto' type");

clang-tools-extra/clangd/refactor/tweaks/ExpandMacro.cpp

 ///   10*a+10*a
 class ExpandMacro : public Tweak {
 public:
   const char *id() const override final;
   llvm::StringLiteral kind() const override {
     return CodeAction::REFACTOR_KIND;
   }
 
-  bool prepare(const Selection &Inputs) override;
-  Expected<Tweak::Effect> apply(const Selection &Inputs) override;
+  bool prepare(const Selection &Inputs, const DraftStore &) override;
+  Expected<Tweak::Effect> apply(const Selection &Inputs,
+                                const DraftStore &) override;
   std::string title() const override;
 
 private:
   syntax::TokenBuffer::Expansion Expansion;
   std::string MacroName;
 };
 
 REGISTER_TWEAK(ExpandMacro)
findIdentifierUnderCursor(const syntax::TokenBuffer &Tokens,
   if (T == Spelled.begin())
     return nullptr;
   --T;
   if (T->endLocation() != Cursor || T->kind() != tok::identifier)
     return nullptr;
   return T;
 }
 
-bool ExpandMacro::prepare(const Selection &Inputs) {
+bool ExpandMacro::prepare(const Selection &Inputs, const DraftStore &) {
   // FIXME: we currently succeed on selection at the end of the token, e.g.
   //        'FOO[[ ]]BAR'. We should not trigger in that case.
 
   // Find a token under the cursor.
   auto *T = findIdentifierUnderCursor(Inputs.AST->getTokens(), Inputs.Cursor);
   // We are interested only in identifiers, other tokens can't be macro names.
   if (!T)
     return false;
   // If the identifier is a macro we will find the corresponding expansion.
   auto Expansion = Inputs.AST->getTokens().expansionStartingAt(T);
   if (!Expansion)
     return false;
   this->MacroName = std::string(T->text(Inputs.AST->getSourceManager()));
   this->Expansion = *Expansion;
   return true;
 }
 
-Expected<Tweak::Effect> ExpandMacro::apply(const Selection &Inputs) {
+Expected<Tweak::Effect> ExpandMacro::apply(const Selection &Inputs,
+                                           const DraftStore &) {
   auto &SM = Inputs.AST->getSourceManager();
 
   std::string Replacement;
   for (const syntax::Token &T : Expansion.Expanded) {
     Replacement += T.text(SM);
     Replacement += " ";
   }
   if (!Replacement.empty()) {

clang-tools-extra/clangd/refactor/tweaks/ExtractFunction.cpp

if (!createParameters(ExtractedFunc, CapturedInfo) ||
                                 CapturedInfo))
     return error("Too complex to extract.");
   return ExtractedFunc;
 }
 
 class ExtractFunction : public Tweak {
 public:
   const char *id() const override final;
-  bool prepare(const Selection &Inputs) override;
-  Expected<Effect> apply(const Selection &Inputs) override;
+  bool prepare(const Selection &Inputs, const DraftStore &) override;
+  Expected<Effect> apply(const Selection &Inputs, const DraftStore &) override;
   std::string title() const override { return "Extract to function"; }
   llvm::StringLiteral kind() const override {
     return CodeAction::REFACTOR_KIND;
   }
 
 private:
   ExtractionZone ExtZone;
 };
bool hasReturnStmt(const ExtractionZone &ExtZone) {
   for (const Stmt *RootStmt : ExtZone.RootStmts) {
     V.TraverseStmt(const_cast<Stmt *>(RootStmt));
     if (V.Found)
       break;
   }
   return V.Found;
 }
 
-bool ExtractFunction::prepare(const Selection &Inputs) {
+bool ExtractFunction::prepare(const Selection &Inputs, const DraftStore &) {
   const LangOptions &LangOpts = Inputs.AST->getLangOpts();
   if (!LangOpts.CPlusPlus)
     return false;
   const Node *CommonAnc = Inputs.ASTSelection.commonAncestor();
   const SourceManager &SM = Inputs.AST->getSourceManager();
   auto MaybeExtZone = findExtractionZone(CommonAnc, SM, LangOpts);
   if (!MaybeExtZone ||
       (hasReturnStmt(*MaybeExtZone) && !alwaysReturns(*MaybeExtZone)))
     return false;
 
   // FIXME: Get rid of this check once we support hoisting.
   if (MaybeExtZone->requiresHoisting(SM))
     return false;
 
   ExtZone = std::move(*MaybeExtZone);
   return true;
 }
 
-Expected<Tweak::Effect> ExtractFunction::apply(const Selection &Inputs) {
+Expected<Tweak::Effect> ExtractFunction::apply(const Selection &Inputs,
+                                               const DraftStore &) {
   const SourceManager &SM = Inputs.AST->getSourceManager();
   const LangOptions &LangOpts = Inputs.AST->getLangOpts();
   auto ExtractedFunc = getExtractedFunction(ExtZone, SM, LangOpts);
   // FIXME: Add more types of errors.
   if (!ExtractedFunc)
     return ExtractedFunc.takeError();
   tooling::Replacements Result;
   if (auto Err = Result.add(createFunctionDefinition(*ExtractedFunc, SM)))

clang-tools-extra/clangd/refactor/tweaks/ExtractVariable.cpp

 /// int x = 5 + 4 * 3;
 ///         ^^^^^
 /// After:
 /// auto dummy = 5 + 4;
 /// int x = dummy * 3;
 class ExtractVariable : public Tweak {
 public:
   const char *id() const override final;
-  bool prepare(const Selection &Inputs) override;
-  Expected<Effect> apply(const Selection &Inputs) override;
+  bool prepare(const Selection &Inputs, const DraftStore &) override;
+  Expected<Effect> apply(const Selection &Inputs, const DraftStore &) override;
   std::string title() const override {
     return "Extract subexpression to variable";
   }
   llvm::StringLiteral kind() const override {
     return CodeAction::REFACTOR_KIND;
   }
 
 private:
   // the expression to extract
   std::unique_ptr<ExtractionContext> Target;
 };
 REGISTER_TWEAK(ExtractVariable)
-bool ExtractVariable::prepare(const Selection &Inputs) {
+bool ExtractVariable::prepare(const Selection &Inputs, const DraftStore &) {
   // we don't trigger on empty selections for now
   if (Inputs.SelectionBegin == Inputs.SelectionEnd)
     return false;
   const ASTContext &Ctx = Inputs.AST->getASTContext();
   // FIXME: Enable non-C++ cases once we start spelling types explicitly instead
   // of making use of auto.
   if (!Ctx.getLangOpts().CPlusPlus)
     return false;
   const SourceManager &SM = Inputs.AST->getSourceManager();
   if (const SelectionTree::Node *N =
           computeExtractedExpr(Inputs.ASTSelection.commonAncestor()))
     Target = std::make_unique<ExtractionContext>(N, SM, Ctx);
   return Target && Target->isExtractable();
 }
 
-Expected<Tweak::Effect> ExtractVariable::apply(const Selection &Inputs) {
+Expected<Tweak::Effect> ExtractVariable::apply(const Selection &Inputs,
+                                               const DraftStore &) {
   tooling::Replacements Result;
   // FIXME: get variable name from user or suggest based on type
   std::string VarName = "dummy";
   SourceRange Range = Target->getExtractionChars();
   // insert new variable declaration
   if (auto Err = Result.add(Target->insertDeclaration(VarName, Range)))
     return std::move(Err);
   // replace expression with variable name

clang-tools-extra/clangd/refactor/tweaks/ObjCLocalizeStringLiteral.cpp

 ///   NSLocalizedString(@"description", @"")
 class ObjCLocalizeStringLiteral : public Tweak {
 public:
   const char *id() const override final;
   llvm::StringLiteral kind() const override {
     return CodeAction::REFACTOR_KIND;
   }
 
-  bool prepare(const Selection &Inputs) override;
-  Expected<Tweak::Effect> apply(const Selection &Inputs) override;
+  bool prepare(const Selection &Inputs, const DraftStore &) override;
+  Expected<Tweak::Effect> apply(const Selection &Inputs,
+                                const DraftStore &) override;
   std::string title() const override;
 
 private:
   const clang::ObjCStringLiteral *Str = nullptr;
 };
 
 REGISTER_TWEAK(ObjCLocalizeStringLiteral)
 
-bool ObjCLocalizeStringLiteral::prepare(const Selection &Inputs) {
+bool ObjCLocalizeStringLiteral::prepare(const Selection &Inputs,
+                                        const DraftStore &) {
   const SelectionTree::Node *N = Inputs.ASTSelection.commonAncestor();
   if (!N)
     return false;
   // Allow the refactoring even if the user selected only the C string part
   // of the expression.
   if (N->ASTNode.get<StringLiteral>()) {
     if (N->Parent)
       N = N->Parent;
   }
   Str = dyn_cast_or_null<ObjCStringLiteral>(N->ASTNode.get<Stmt>());
   return Str;
 }
 
 Expected<Tweak::Effect>
-ObjCLocalizeStringLiteral::apply(const Selection &Inputs) {
+ObjCLocalizeStringLiteral::apply(const Selection &Inputs, const DraftStore &) {
   auto *AST = Inputs.AST;
   auto &SM = AST->getSourceManager();
   const auto &TB = AST->getTokens();
   auto Toks = TB.spelledForExpanded(TB.expandedTokens(Str->getSourceRange()));
   if (!Toks || Toks->empty())
     return error("Failed to find tokens to replace.");
   // Insert `NSLocalizedString(` before the literal.
   auto Reps = tooling::Replacements(tooling::Replacement(

clang-tools-extra/clangd/refactor/tweaks/PopulateSwitch.cpp

 #include <cassert>
 #include <string>
 
 namespace clang {
 namespace clangd {
 namespace {
 class PopulateSwitch : public Tweak {
   const char *id() const override;
-  bool prepare(const Selection &Sel) override;
-  Expected<Effect> apply(const Selection &Sel) override;
+  bool prepare(const Selection &Sel, const DraftStore &) override;
+  Expected<Effect> apply(const Selection &Sel, const DraftStore &) override;
   std::string title() const override { return "Populate switch"; }
   llvm::StringLiteral kind() const override {
     return CodeAction::QUICKFIX_KIND;
   }
 
 private:
   class ExpectedCase {
   public:
private:
   const EnumDecl *EnumD = nullptr;
   // Maps the Enum values to the EnumConstantDecl and a bool signifying if its
   // covered in the switch.
   llvm::MapVector<llvm::APSInt, ExpectedCase> ExpectedCases;
 };
 
 REGISTER_TWEAK(PopulateSwitch)
 
-bool PopulateSwitch::prepare(const Selection &Sel) {
+bool PopulateSwitch::prepare(const Selection &Sel, const DraftStore &) {
   const SelectionTree::Node *CA = Sel.ASTSelection.commonAncestor();
   if (!CA)
     return false;
 
   const Stmt *CAStmt = CA->ASTNode.get<Stmt>();
   if (!CAStmt)
     return false;
 
for (const SwitchCase *CaseList = Switch->getSwitchCaseList(); CaseList;
     if (Iter != ExpectedCases.end())
       Iter->second.setCovered();
   }
 
   return !llvm::all_of(ExpectedCases,
                        [](auto &Pair) { return Pair.second.isCovered(); });
 }
 
-Expected<Tweak::Effect> PopulateSwitch::apply(const Selection &Sel) {
+Expected<Tweak::Effect> PopulateSwitch::apply(const Selection &Sel,
+                                              const DraftStore &) {
   ASTContext &Ctx = Sel.AST->getASTContext();
 
   SourceLocation Loc = Body->getRBracLoc();
   ASTContext &DeclASTCtx = DeclCtx->getParentASTContext();
 
   llvm::SmallString<256> Text;
   for (auto &EnumConstant : ExpectedCases) {
     // Skip any enum constants already covered

clang-tools-extra/clangd/refactor/tweaks/RawStringLiteral.cpp

 ///          ^^^^^^^^^
 /// After:
 ///   printf(R"("a"
 /// b)");
 class RawStringLiteral : public Tweak {
 public:
   const char *id() const override final;
 
-  bool prepare(const Selection &Inputs) override;
-  Expected<Effect> apply(const Selection &Inputs) override;
+  bool prepare(const Selection &Inputs, const DraftStore &) override;
+  Expected<Effect> apply(const Selection &Inputs, const DraftStore &) override;
   std::string title() const override { return "Convert to raw string"; }
   llvm::StringLiteral kind() const override {
     return CodeAction::REFACTOR_KIND;
   }
 
 private:
   const clang::StringLiteral *Str = nullptr;
 };

 
 static bool canBeRaw(llvm::StringRef Content) {
   for (char C : Content)
     if (!llvm::isPrint(C) && C != '\n' && C != '\t')
       return false;
   return !Content.contains(")\"");
 }
 
-bool RawStringLiteral::prepare(const Selection &Inputs) {
+bool RawStringLiteral::prepare(const Selection &Inputs, const DraftStore &) {
   const SelectionTree::Node *N = Inputs.ASTSelection.commonAncestor();
   if (!N)
     return false;
   Str = dyn_cast_or_null<StringLiteral>(N->ASTNode.get<Stmt>());
   return Str &&
          isNormalString(*Str, Inputs.Cursor, Inputs.AST->getSourceManager()) &&
          needsRaw(Str->getBytes()) && canBeRaw(Str->getBytes());
 }
 
-Expected<Tweak::Effect> RawStringLiteral::apply(const Selection &Inputs) {
+Expected<Tweak::Effect> RawStringLiteral::apply(const Selection &Inputs,
+                                                const DraftStore &) {
   auto &SM = Inputs.AST->getSourceManager();
   auto Reps = tooling::Replacements(
       tooling::Replacement(SM, Str, ("R\"(" + Str->getBytes() + ")\"").str(),
                            Inputs.AST->getLangOpts()));
   return Effect::mainFileEdit(SM, std::move(Reps));
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/refactor/tweaks/RemoveUsingNamespace.cpp

 ///   std::vector<int> foo(std::map<int, int>);
 /// Currently limited to using namespace directives inside global namespace to
 /// simplify implementation. Also the namespace must not contain using
 /// directives.
 class RemoveUsingNamespace : public Tweak {
 public:
   const char *id() const override;
 
-  bool prepare(const Selection &Inputs) override;
-  Expected<Effect> apply(const Selection &Inputs) override;
+  bool prepare(const Selection &Inputs, const DraftStore &) override;
+  Expected<Effect> apply(const Selection &Inputs, const DraftStore &) override;
   std::string title() const override {
     return "Remove using namespace, re-qualify names instead";
   }
   llvm::StringLiteral kind() const override {
     return CodeAction::REFACTOR_KIND;
   }
 
 private:

 // enum E { a, b, c, d };
 // }
 const DeclContext *visibleContext(const DeclContext *D) {
   while (D->isInlineNamespace() || D->isTransparentContext())
     D = D->getParent();
   return D;
 }
 
-bool RemoveUsingNamespace::prepare(const Selection &Inputs) {
+bool RemoveUsingNamespace::prepare(const Selection &Inputs,
+                                   const DraftStore &) {
   // Find the 'using namespace' directive under the cursor.
   auto *CA = Inputs.ASTSelection.commonAncestor();
   if (!CA)
     return false;
   TargetDirective = CA->ASTNode.get<UsingDirectiveDecl>();
   if (!TargetDirective)
     return false;
   if (!dyn_cast<Decl>(TargetDirective->getDeclContext()))
     return false;
   // FIXME: Unavailable for namespaces containing using-namespace decl.
   // It is non-trivial to deal with cases where identifiers come from the inner
   // namespace. For example map has to be changed to aa::map.
   // namespace aa {
   //   namespace bb { struct map {}; }
   //   using namespace bb;
   // }
   // using namespace a^a;
   // int main() { map m; }
   // We need to make this aware of the transitive using-namespace decls.
   if (!TargetDirective->getNominatedNamespace()->using_directives().empty())
     return false;
   return isTopLevelDecl(CA);
 }
 
-Expected<Tweak::Effect> RemoveUsingNamespace::apply(const Selection &Inputs) {
+Expected<Tweak::Effect> RemoveUsingNamespace::apply(const Selection &Inputs,
+                                                    const DraftStore &) {
   auto &Ctx = Inputs.AST->getASTContext();
   auto &SM = Ctx.getSourceManager();
   // First, collect *all* using namespace directives that redeclare the same
   // namespace.
   std::vector<const UsingDirectiveDecl *> AllDirectives;
   FindSameUsings(*TargetDirective, AllDirectives).TraverseAST(Ctx);
 
   SourceLocation FirstUsingDirectiveLoc;

clang-tools-extra/clangd/refactor/tweaks/SwapIfBranches.cpp

 ///   if (foo) { return 10; } else { continue; }
 ///   ^^^^^^^                 ^^^^
 /// After:
 ///   if (foo) { continue; } else { return 10; }
 class SwapIfBranches : public Tweak {
 public:
   const char *id() const override final;
 
-  bool prepare(const Selection &Inputs) override;
-  Expected<Effect> apply(const Selection &Inputs) override;
+  bool prepare(const Selection &Inputs, const DraftStore &) override;
+  Expected<Effect> apply(const Selection &Inputs, const DraftStore &) override;
   std::string title() const override { return "Swap if branches"; }
   llvm::StringLiteral kind() const override {
     return CodeAction::REFACTOR_KIND;
   }
   bool hidden() const override { return true; }
 
 private:
   const IfStmt *If = nullptr;
 };
 
 REGISTER_TWEAK(SwapIfBranches)
 
-bool SwapIfBranches::prepare(const Selection &Inputs) {
+bool SwapIfBranches::prepare(const Selection &Inputs, const DraftStore &) {
   for (const SelectionTree::Node *N = Inputs.ASTSelection.commonAncestor();
        N && !If; N = N->Parent) {
     // Stop once we hit a block, e.g. a lambda in the if condition.
     if (dyn_cast_or_null<CompoundStmt>(N->ASTNode.get<Stmt>()))
       return false;
     If = dyn_cast_or_null<IfStmt>(N->ASTNode.get<Stmt>());
   }
   // avoid dealing with single-statement brances, they require careful handling
   // to avoid changing semantics of the code (i.e. dangling else).
   return If && dyn_cast_or_null<CompoundStmt>(If->getThen()) &&
          dyn_cast_or_null<CompoundStmt>(If->getElse());
 }
 
-Expected<Tweak::Effect> SwapIfBranches::apply(const Selection &Inputs) {
+Expected<Tweak::Effect> SwapIfBranches::apply(const Selection &Inputs,
+                                              const DraftStore &) {
   auto &Ctx = Inputs.AST->getASTContext();
   auto &SrcMgr = Inputs.AST->getSourceManager();
 
   auto ThenRng = toHalfOpenFileRange(SrcMgr, Ctx.getLangOpts(),
                                      If->getThen()->getSourceRange());
   if (!ThenRng)
     return error("Could not obtain range of the 'then' branch. Macros?");
   auto ElseRng = toHalfOpenFileRange(SrcMgr, Ctx.getLangOpts(),

clang-tools-extra/clangd/tool/Check.cpp

bool buildAST() {
     return true;
   }
 
   // Run AST-based features at each token in the file.
   void testLocationFeatures() {
     log("Testing features at each token (may be slow in large files)");
     auto SpelledTokens =
         AST->getTokens().spelledTokens(AST->getSourceManager().getMainFileID());
+
+    DraftStore EmptyDrafts;
     for (const auto &Tok : SpelledTokens) {
       unsigned Start = AST->getSourceManager().getFileOffset(Tok.location());
       unsigned End = Start + Tok.length();
       Position Pos = offsetToPosition(Inputs.Contents, Start);
       // FIXME: dumping the tokens may leak sensitive code into bug reports.
       // Add an option to turn this off, once we decide how options work.
       vlog("  {0} {1}", Pos, Tok.text(AST->getSourceManager()));
       auto Tree = SelectionTree::createRight(AST->getASTContext(),
                                              AST->getTokens(), Start, End);
       Tweak::Selection Selection(&Index, *AST, Start, End, std::move(Tree));
-      for (const auto &T : prepareTweaks(Selection, Opts.TweakFilter)) {
-        auto Result = T->apply(Selection);
+      for (const auto &T :
+           prepareTweaks(Selection, EmptyDrafts, Opts.TweakFilter)) {
+        auto Result = T->apply(Selection, EmptyDrafts);
         if (!Result) {
           elog("    tweak: {0} ==> FAIL: {1}", T->id(), Result.takeError());
           ++ErrCount;
         } else {
           vlog("    tweak: {0}", T->id());
         }
       }
       unsigned Definitions = locateSymbolAt(*AST, Pos, &Index).size();

clang-tools-extra/clangd/unittests/tweaks/TweakTesting.cpp

 //===-- TweakTesting.cpp ------------------------------------------------*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "TweakTesting.h"
 
 #include "Annotations.h"
+#include "DraftStore.h"
 #include "SourceCode.h"
 #include "TestFS.h"
 #include "refactor/Tweak.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/Support/Error.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <string>

 }
 
 // Prepare and apply the specified tweak based on the selection in Input.
 // Returns None if and only if prepare() failed.
 llvm::Optional<llvm::Expected<Tweak::Effect>>
 applyTweak(ParsedAST &AST, const Annotations &Input, StringRef TweakID,
            const SymbolIndex *Index) {
   auto Range = rangeOrPoint(Input);
+  DraftStore EmptyDrafts;
   llvm::Optional<llvm::Expected<Tweak::Effect>> Result;
   SelectionTree::createEach(AST.getASTContext(), AST.getTokens(), Range.first,
                             Range.second, [&](SelectionTree ST) {
                               Tweak::Selection S(Index, AST, Range.first,
                                                  Range.second, std::move(ST));
-                              if (auto T = prepareTweak(TweakID, S)) {
-                                Result = (*T)->apply(S);
+                              if (auto T =
+                                      prepareTweak(TweakID, S, EmptyDrafts)) {
+                                Result = (*T)->apply(S, EmptyDrafts);
                                 return true;
                               } else {
                                 llvm::consumeError(T.takeError());
                                 return false;
                               }
                             });
   return Result;
 }