diff --git a/clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp b/clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp --- a/clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp +++ b/clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp @@ -7,6 +7,7 @@ //===----------------------------------------------------------------------===// #include "AST.h" +#include "DraftStore.h" #include "FindTarget.h" #include "HeaderSourceSwitch.h" #include "ParsedAST.h" @@ -36,6 +37,7 @@ #include "llvm/ADT/StringRef.h" #include "llvm/Support/Casting.h" #include "llvm/Support/Error.h" +#include "llvm/Support/MemoryBuffer.h" #include #include @@ -397,7 +399,8 @@ return true; } - Expected apply(const Selection &Sel, const DraftStore &) override { + Expected apply(const Selection &Sel, + const DraftStore &DraftMgr) override { const SourceManager &SM = Sel.AST->getSourceManager(); auto MainFileName = getCanonicalPath(SM.getFileEntryForID(SM.getMainFileID()), SM); @@ -408,14 +411,24 @@ if (!CCFile) return error("Couldn't find a suitable implementation file."); - auto &FS = - Sel.AST->getSourceManager().getFileManager().getVirtualFileSystem(); - auto Buffer = FS.getBufferForFile(*CCFile); - // FIXME: Maybe we should consider creating the implementation file if it - // doesn't exist? - if (!Buffer) - return llvm::errorCodeToError(Buffer.getError()); - auto Contents = Buffer->get()->getBuffer(); + StringRef Contents; + // We need to keep both these alive as Contents can reference the storage + // contained in either of them. + auto Draft = DraftMgr.getDraft(*CCFile); + std::unique_ptr Buffer; + if (Draft) { + Contents = Draft->Contents; + } else { + auto &FS = + Sel.AST->getSourceManager().getFileManager().getVirtualFileSystem(); + auto BufferOr = FS.getBufferForFile(*CCFile); + // FIXME: Maybe we should consider creating the implementation file if it + // doesn't exist? + if (!BufferOr) + return llvm::errorCodeToError(BufferOr.getError()); + Buffer = std::move(*BufferOr); + Contents = Buffer->getBuffer(); + } auto InsertionPoint = getInsertionPoint( Contents, Source->getQualifiedNameAsString(), Sel.AST->getLangOpts()); if (!InsertionPoint)