Add SourceManagerForFile helper which sets up SourceManager and dependencies for a single file with code snippet

Summary: This can be used to create a virtual environment (incl. VFS, source manager) for code snippets.

Reviewers: sammccall, klimek

Reviewed By: sammccall

Subscribers: klimek, mgorny, cfe-commits

Differential Revision: https://reviews.llvm.org/D46176

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@331923 91177308-0d34-0410-b5e6-96231b3b80d8

6 files changed, 90 insertions, 100 deletions

diff --git a/include/clang/Basic/SourceManager.h b/include/clang/Basic/SourceManager.h
index 25acdd2600..33e63d9d40 100644
--- a/include/clang/Basic/SourceManager.h
+++ b/include/clang/Basic/SourceManager.h
@@ -35,6 +35,7 @@
 #ifndef LLVM_CLANG_BASIC_SOURCEMANAGER_H
 #define LLVM_CLANG_BASIC_SOURCEMANAGER_H
 
+#include "clang/Basic/Diagnostic.h"
 #include "clang/Basic/FileManager.h"
 #include "clang/Basic/SourceLocation.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -60,7 +61,6 @@ namespace clang {
 
 class ASTReader;
 class ASTWriter;
-class DiagnosticsEngine;
 class LineTableInfo;
 class SourceManager;
 
@@ -1815,6 +1815,28 @@ public:
   }
 };
 
+/// SourceManager and necessary depdencies (e.g. VFS, FileManager) for a single
+/// in-memorty file.
+class SourceManagerForFile {
+public:
+  /// Creates SourceManager and necessary depdencies (e.g. VFS, FileManager).
+  /// The main file in the SourceManager will be \p FileName with \p Content.
+  SourceManagerForFile(StringRef FileName, StringRef Content);
+
+  SourceManager &get() {
+    assert(SourceMgr);
+    return *SourceMgr;
+  }
+
+private:
+  // The order of these fields are important - they should be in the same order
+  // as they are created in `createSourceManagerForFile` so that they can be
+  // deleted in the reverse order as they are created.
+  std::unique_ptr<FileManager> FileMgr;
+  std::unique_ptr<DiagnosticsEngine> Diagnostics;
+  std::unique_ptr<SourceManager> SourceMgr;
+};
+
 } // namespace clang
 
 #endif // LLVM_CLANG_BASIC_SOURCEMANAGER_H
diff --git a/lib/Basic/SourceManager.cpp b/lib/Basic/SourceManager.cpp
index c05dfeaa8a..ae76817826 100644
--- a/lib/Basic/SourceManager.cpp
+++ b/lib/Basic/SourceManager.cpp
@@ -2258,3 +2258,29 @@ size_t SourceManager::getDataStructureSizes() const {
 
   return size;
 }
+
+SourceManagerForFile::SourceManagerForFile(StringRef FileName,
+                                           StringRef Content) {
+  // This is referenced by `FileMgr` and will be released by `FileMgr` when it
+  // is deleted.
+  IntrusiveRefCntPtr<vfs::InMemoryFileSystem> InMemoryFileSystem(
+      new vfs::InMemoryFileSystem);
+  InMemoryFileSystem->addFile(
+      FileName, 0,
+      llvm::MemoryBuffer::getMemBuffer(Content, FileName,
+                                       /*RequiresNullTerminator=*/false));
+  // This is passed to `SM` as reference, so the pointer has to be referenced
+  // in `Environment` so that `FileMgr` can out-live this function scope.
+  FileMgr =
+      llvm::make_unique<FileManager>(FileSystemOptions(), InMemoryFileSystem);
+  // This is passed to `SM` as reference, so the pointer has to be referenced
+  // by `Environment` due to the same reason above.
+  Diagnostics = llvm::make_unique<DiagnosticsEngine>(
+      IntrusiveRefCntPtr<DiagnosticIDs>(new DiagnosticIDs),
+      new DiagnosticOptions);
+  SourceMgr = llvm::make_unique<SourceManager>(*Diagnostics, *FileMgr);
+  FileID ID = SourceMgr->createFileID(FileMgr->getFile(FileName),
+                                      SourceLocation(), clang::SrcMgr::C_User);
+  assert(ID.isValid());
+  SourceMgr->setMainFileID(ID);
+}
diff --git a/lib/Format/Format.cpp b/lib/Format/Format.cpp
index a2a1c7adde..d2748dbe77 100644
--- a/lib/Format/Format.cpp
+++ b/lib/Format/Format.cpp
@@ -1905,10 +1905,9 @@ unsigned getOffsetAfterTokenSequence(
     StringRef FileName, StringRef Code, const FormatStyle &Style,
     llvm::function_ref<unsigned(const SourceManager &, Lexer &, Token &)>
         GetOffsetAfterSequence) {
-  std::unique_ptr<Environment> Env =
-      Environment::CreateVirtualEnvironment(Code, FileName, /*Ranges=*/{});
-  const SourceManager &SourceMgr = Env->getSourceManager();
-  Lexer Lex(Env->getFileID(), SourceMgr.getBuffer(Env->getFileID()), SourceMgr,
+  Environment Env(Code, FileName, /*Ranges=*/{});
+  const SourceManager &SourceMgr = Env.getSourceManager();
+  Lexer Lex(Env.getFileID(), SourceMgr.getBuffer(Env.getFileID()), SourceMgr,
             getFormattingLangOpts(Style));
   Token Tok;
   // Get the first token.
@@ -2361,9 +2360,9 @@ reformat(const FormatStyle &Style, StringRef Code,
     return Formatter(Env, Expanded, Status).process();
   });
 
-  std::unique_ptr<Environment> Env = Environment::CreateVirtualEnvironment(
-      Code, FileName, Ranges, FirstStartColumn, NextStartColumn,
-      LastStartColumn);
+  auto Env =
+      llvm::make_unique<Environment>(Code, FileName, Ranges, FirstStartColumn,
+                                     NextStartColumn, LastStartColumn);
   llvm::Optional<std::string> CurrentCode = None;
   tooling::Replacements Fixes;
   unsigned Penalty = 0;
@@ -2376,7 +2375,7 @@ reformat(const FormatStyle &Style, StringRef Code,
       Penalty += PassFixes.second;
       if (I + 1 < E) {
         CurrentCode = std::move(*NewCode);
-        Env = Environment::CreateVirtualEnvironment(
+        Env = llvm::make_unique<Environment>(
             *CurrentCode, FileName,
             tooling::calculateRangesAfterReplacements(Fixes, Ranges),
             FirstStartColumn, NextStartColumn, LastStartColumn);
@@ -2405,10 +2404,7 @@ tooling::Replacements cleanup(const FormatStyle &Style, StringRef Code,
   // cleanups only apply to C++ (they mostly concern ctor commas etc.)
   if (Style.Language != FormatStyle::LK_Cpp)
     return tooling::Replacements();
-  std::unique_ptr<Environment> Env =
-      Environment::CreateVirtualEnvironment(Code, FileName, Ranges);
-  Cleaner Clean(*Env, Style);
-  return Clean.process().first;
+  return Cleaner(Environment(Code, FileName, Ranges), Style).process().first;
 }
 
 tooling::Replacements reformat(const FormatStyle &Style, StringRef Code,
@@ -2425,20 +2421,18 @@ tooling::Replacements fixNamespaceEndComments(const FormatStyle &Style,
                                               StringRef Code,
                                               ArrayRef<tooling::Range> Ranges,
                                               StringRef FileName) {
-  std::unique_ptr<Environment> Env =
-      Environment::CreateVirtualEnvironment(Code, FileName, Ranges);
-  NamespaceEndCommentsFixer Fix(*Env, Style);
-  return Fix.process().first;
+  return NamespaceEndCommentsFixer(Environment(Code, FileName, Ranges), Style)
+      .process()
+      .first;
 }
 
 tooling::Replacements sortUsingDeclarations(const FormatStyle &Style,
                                             StringRef Code,
                                             ArrayRef<tooling::Range> Ranges,
                                             StringRef FileName) {
-  std::unique_ptr<Environment> Env =
-      Environment::CreateVirtualEnvironment(Code, FileName, Ranges);
-  UsingDeclarationsSorter Sorter(*Env, Style);
-  return Sorter.process().first;
+  return UsingDeclarationsSorter(Environment(Code, FileName, Ranges), Style)
+      .process()
+      .first;
 }
 
 LangOptions getFormattingLangOpts(const FormatStyle &Style) {
@@ -2498,9 +2492,8 @@ FormatStyle::LanguageKind guessLanguage(StringRef FileName, StringRef Code) {
     // of the code to see if it contains Objective-C.
     if (Extension.empty() || Extension == ".h") {
       auto NonEmptyFileName = FileName.empty() ? "guess.h" : FileName;
-      std::unique_ptr<Environment> Env =
-          Environment::CreateVirtualEnvironment(Code, NonEmptyFileName, /*Ranges=*/{});
-      ObjCHeaderStyleGuesser Guesser(*Env, getLLVMStyle());
+      Environment Env(Code, NonEmptyFileName, /*Ranges=*/{});
+      ObjCHeaderStyleGuesser Guesser(Env, getLLVMStyle());
       Guesser.process();
       if (Guesser.isObjC())
         return FormatStyle::LK_ObjC;
diff --git a/lib/Format/SortJavaScriptImports.cpp b/lib/Format/SortJavaScriptImports.cpp
index 0ca2c4433c..e8df04c8bd 100644
--- a/lib/Format/SortJavaScriptImports.cpp
+++ b/lib/Format/SortJavaScriptImports.cpp
@@ -445,10 +445,9 @@ tooling::Replacements sortJavaScriptImports(const FormatStyle &Style,
                                             ArrayRef<tooling::Range> Ranges,
                                             StringRef FileName) {
   // FIXME: Cursor support.
-  std::unique_ptr<Environment> Env =
-      Environment::CreateVirtualEnvironment(Code, FileName, Ranges);
-  JavaScriptImportSorter Sorter(*Env, Style);
-  return Sorter.process().first;
+  return JavaScriptImportSorter(Environment(Code, FileName, Ranges), Style)
+      .process()
+      .first;
 }
 
 } // end namespace format
diff --git a/lib/Format/TokenAnalyzer.cpp b/lib/Format/TokenAnalyzer.cpp
index 806a596179..ac31f403d5 100644
--- a/lib/Format/TokenAnalyzer.cpp
+++ b/lib/Format/TokenAnalyzer.cpp
@@ -34,48 +34,19 @@
 namespace clang {
 namespace format {
 
-// This sets up an virtual file system with file \p FileName containing \p
-// Code.
-std::unique_ptr<Environment>
-Environment::CreateVirtualEnvironment(StringRef Code, StringRef FileName,
-                                      ArrayRef<tooling::Range> Ranges,
-                                      unsigned FirstStartColumn,
-                                      unsigned NextStartColumn,
-                                      unsigned LastStartColumn) {
-  // This is referenced by `FileMgr` and will be released by `FileMgr` when it
-  // is deleted.
-  IntrusiveRefCntPtr<vfs::InMemoryFileSystem> InMemoryFileSystem(
-      new vfs::InMemoryFileSystem);
-  // This is passed to `SM` as reference, so the pointer has to be referenced
-  // in `Environment` so that `FileMgr` can out-live this function scope.
-  std::unique_ptr<FileManager> FileMgr(
-      new FileManager(FileSystemOptions(), InMemoryFileSystem));
-  // This is passed to `SM` as reference, so the pointer has to be referenced
-  // by `Environment` due to the same reason above.
-  std::unique_ptr<DiagnosticsEngine> Diagnostics(new DiagnosticsEngine(
-      IntrusiveRefCntPtr<DiagnosticIDs>(new DiagnosticIDs),
-      new DiagnosticOptions));
-  // This will be stored as reference, so the pointer has to be stored in
-  // due to the same reason above.
-  std::unique_ptr<SourceManager> VirtualSM(
-      new SourceManager(*Diagnostics, *FileMgr));
-  InMemoryFileSystem->addFile(
-      FileName, 0,
-      llvm::MemoryBuffer::getMemBuffer(Code, FileName,
-                                       /*RequiresNullTerminator=*/false));
-  FileID ID = VirtualSM->createFileID(FileMgr->getFile(FileName),
-                                      SourceLocation(), clang::SrcMgr::C_User);
-  assert(ID.isValid());
-  SourceLocation StartOfFile = VirtualSM->getLocForStartOfFile(ID);
-  std::vector<CharSourceRange> CharRanges;
+Environment::Environment(StringRef Code, StringRef FileName,
+                         ArrayRef<tooling::Range> Ranges,
+                         unsigned FirstStartColumn, unsigned NextStartColumn,
+                         unsigned LastStartColumn)
+    : VirtualSM(new SourceManagerForFile(FileName, Code)), SM(VirtualSM->get()),
+      ID(VirtualSM->get().getMainFileID()), FirstStartColumn(FirstStartColumn),
+      NextStartColumn(NextStartColumn), LastStartColumn(LastStartColumn) {
+  SourceLocation StartOfFile = SM.getLocForStartOfFile(ID);
   for (const tooling::Range &Range : Ranges) {
     SourceLocation Start = StartOfFile.getLocWithOffset(Range.getOffset());
     SourceLocation End = Start.getLocWithOffset(Range.getLength());
     CharRanges.push_back(CharSourceRange::getCharRange(Start, End));
   }
-  return llvm::make_unique<Environment>(
-      ID, std::move(FileMgr), std::move(VirtualSM), std::move(Diagnostics),
-      CharRanges, FirstStartColumn, NextStartColumn, LastStartColumn);
 }
 
 TokenAnalyzer::TokenAnalyzer(const Environment &Env, const FormatStyle &Style)
diff --git a/lib/Format/TokenAnalyzer.h b/lib/Format/TokenAnalyzer.h
index a8978b56cb..e43a860e46 100644
--- a/lib/Format/TokenAnalyzer.h
+++ b/lib/Format/TokenAnalyzer.h
@@ -37,44 +37,24 @@ namespace format {
 class Environment {
 public:
   Environment(SourceManager &SM, FileID ID, ArrayRef<CharSourceRange> Ranges)
-      : ID(ID), CharRanges(Ranges.begin(), Ranges.end()), SM(SM),
-      FirstStartColumn(0),
-      NextStartColumn(0),
-      LastStartColumn(0) {}
-
-  Environment(FileID ID, std::unique_ptr<FileManager> FileMgr,
-              std::unique_ptr<SourceManager> VirtualSM,
-              std::unique_ptr<DiagnosticsEngine> Diagnostics,
-              const std::vector<CharSourceRange> &CharRanges,
-              unsigned FirstStartColumn,
-              unsigned NextStartColumn,
-              unsigned LastStartColumn)
-      : ID(ID), CharRanges(CharRanges.begin(), CharRanges.end()),
-        SM(*VirtualSM), 
-        FirstStartColumn(FirstStartColumn),
-        NextStartColumn(NextStartColumn),
-        LastStartColumn(LastStartColumn),
-        FileMgr(std::move(FileMgr)),
-        VirtualSM(std::move(VirtualSM)), Diagnostics(std::move(Diagnostics)) {}
+      : SM(SM), ID(ID), CharRanges(Ranges.begin(), Ranges.end()),
+        FirstStartColumn(0), NextStartColumn(0), LastStartColumn(0) {}
 
   // This sets up an virtual file system with file \p FileName containing the
   // fragment \p Code. Assumes that \p Code starts at \p FirstStartColumn,
   // that the next lines of \p Code should start at \p NextStartColumn, and
   // that \p Code should end at \p LastStartColumn if it ends in newline.
   // See also the documentation of clang::format::internal::reformat.
-  static std::unique_ptr<Environment>
-  CreateVirtualEnvironment(StringRef Code, StringRef FileName,
-                           ArrayRef<tooling::Range> Ranges,
-                           unsigned FirstStartColumn = 0,
-                           unsigned NextStartColumn = 0,
-                           unsigned LastStartColumn = 0);
+  Environment(StringRef Code, StringRef FileName,
+              ArrayRef<tooling::Range> Ranges, unsigned FirstStartColumn = 0,
+              unsigned NextStartColumn = 0, unsigned LastStartColumn = 0);
 
   FileID getFileID() const { return ID; }
 
-  ArrayRef<CharSourceRange> getCharRanges() const { return CharRanges; }
-
   const SourceManager &getSourceManager() const { return SM; }
 
+  ArrayRef<CharSourceRange> getCharRanges() const { return CharRanges; }
+
   // Returns the column at which the fragment of code managed by this
   // environment starts.
   unsigned getFirstStartColumn() const { return FirstStartColumn; }
@@ -88,19 +68,18 @@ public:
   unsigned getLastStartColumn() const { return LastStartColumn; }
 
 private:
+  // This is only set if constructed from string.
+  std::unique_ptr<SourceManagerForFile> VirtualSM;
+
+  // This refers to either a SourceManager provided by users or VirtualSM
+  // created for a single file.
+  SourceManager &SM;
   FileID ID;
+
   SmallVector<CharSourceRange, 8> CharRanges;
-  SourceManager &SM;
   unsigned FirstStartColumn;
   unsigned NextStartColumn;
   unsigned LastStartColumn;
-
-  // The order of these fields are important - they should be in the same order
-  // as they are created in `CreateVirtualEnvironment` so that they can be
-  // deleted in the reverse order as they are created.
-  std::unique_ptr<FileManager> FileMgr;
-  std::unique_ptr<SourceManager> VirtualSM;
-  std::unique_ptr<DiagnosticsEngine> Diagnostics;
 };
 
 class TokenAnalyzer : public UnwrappedLineConsumer {