Inserting an include of the new header during the rewrite.

Bug: 1069567
Change-Id: I85fa41b304db2f4902b0c2f46290f49279316f3d
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2151338
Reviewed-by: Daniel Cheng <dcheng@chromium.org>
Reviewed-by: Bartek Nowierski <bartekn@chromium.org>
Commit-Queue: Łukasz Anforowicz <lukasza@chromium.org>
Cr-Commit-Position: refs/heads/master@{#766144}

docs/clang_tool_refactoring.md

@@ -54,8 +54,10 @@ Other useful references when writing the tool:
 ### Edit serialization format
 ```
 ==== BEGIN EDITS ====
-r:::path/to/file1:::offset1:::length1:::replacement text
-r:::path/to/file2:::offset2:::length2:::replacement text
+r:::path/to/file/to/edit:::offset1:::length1:::replacement text
+r:::path/to/file/to/edit:::offset2:::length2:::replacement text
+r:::path/to/file2/to/edit:::offset3:::length3:::replacement text
+include-user-header:::path/to/file2/to/edit:::-1:::-1:::header/file/to/include.h
 
        ...
 
@@ -64,8 +66,8 @@ r:::path/to/file2:::offset2:::length2:::replacement text
 
 The header and footer are required. Each line between the header and footer
 represents one edit. Fields are separated by `:::`, and the first field must
-be `r` (for replacement). In the future, this may be extended to handle header
-insertion/removal. A deletion is an edit with no replacement text.
+be `r` (for replacement) or `include-user-header`.
+A deletion is an edit with no replacement text.
 
 The edits are applied by [`apply_edits.py`](#Running), which understands certain
 conventions:

tools/clang/rewrite_raw_ptr_fields/RewriteRawPtrFields.cpp

@@ -41,15 +41,48 @@
 #include "llvm/Support/TargetSelect.h"
 
 using namespace clang::ast_matchers;
-using clang::tooling::CommonOptionsParser;
-using clang::tooling::Replacement;
 
 namespace {
 
+// Include path that needs to be added to all the files where CheckedPtr<...>
+// replaces a raw pointer.
+const char kIncludePath[] = "base/memory/checked_ptr.h";
+
+// Output format is documented in //docs/clang_tool_refactoring.md
+class ReplacementsPrinter {
+ public:
+  ReplacementsPrinter() { llvm::outs() << "==== BEGIN EDITS ====\n"; }
+
+  ~ReplacementsPrinter() { llvm::outs() << "==== END EDITS ====\n"; }
+
+  void PrintReplacement(const clang::SourceManager& source_manager,
+                        const clang::SourceRange& replacement_range,
+                        std::string replacement_text) {
+    clang::tooling::Replacement replacement(
+        source_manager, clang::CharSourceRange::getCharRange(replacement_range),
+        replacement_text);
+    llvm::StringRef file_path = replacement.getFilePath();
+    std::replace(replacement_text.begin(), replacement_text.end(), '\n', '\0');
+    llvm::outs() << "r:::" << file_path << ":::" << replacement.getOffset()
+                 << ":::" << replacement.getLength()
+                 << ":::" << replacement_text << "\n";
+
+    bool was_inserted = false;
+    std::tie(std::ignore, was_inserted) =
+        files_with_already_added_includes_.insert(file_path.str());
+    if (was_inserted)
+      llvm::outs() << "include-user-header:::" << file_path
+                   << ":::-1:::-1:::" << kIncludePath << "\n";
+  }
+
+ private:
+  std::set<std::string> files_with_already_added_includes_;
+};
+
 class FieldDeclRewriter : public MatchFinder::MatchCallback {
  public:
-  explicit FieldDeclRewriter(std::vector<Replacement>* replacements)
-      : replacements_(replacements) {}
+  explicit FieldDeclRewriter(ReplacementsPrinter* replacements_printer)
+      : replacements_printer_(replacements_printer) {}
 
   void run(const MatchFinder::MatchResult& result) override {
     const clang::SourceManager& source_manager = *result.SourceManager;
@@ -88,10 +121,9 @@ class FieldDeclRewriter : public MatchFinder::MatchCallback {
     if (field_decl->isMutable())
       replacement_text.insert(0, "mutable ");
 
-    // Generate and add a replacement.
-    replacements_->emplace_back(
-        source_manager, clang::CharSourceRange::getCharRange(replacement_range),
-        replacement_text);
+    // Generate and print a replacement.
+    replacements_printer_->PrintReplacement(source_manager, replacement_range,
+                                            replacement_text);
   }
 
  private:
@@ -120,7 +152,7 @@ class FieldDeclRewriter : public MatchFinder::MatchCallback {
     return result;
   }
 
-  std::vector<Replacement>* const replacements_;
+  ReplacementsPrinter* const replacements_printer_;
 };
 
 }  // namespace
@@ -132,12 +164,12 @@ int main(int argc, const char* argv[]) {
   llvm::InitializeNativeTargetAsmParser();
   llvm::cl::OptionCategory category(
       "rewrite_raw_ptr_fields: changes |T* field_| to |CheckedPtr<T> field_|.");
-  CommonOptionsParser options(argc, argv, category);
+  clang::tooling::CommonOptionsParser options(argc, argv, category);
   clang::tooling::ClangTool tool(options.getCompilations(),
                                  options.getSourcePathList());
 
   MatchFinder match_finder;
-  std::vector<Replacement> replacements;
+  ReplacementsPrinter replacements_printer;
 
   // Field declarations =========
   // Given
@@ -146,7 +178,7 @@ int main(int argc, const char* argv[]) {
   //   };
   // matches |int* y|.
   auto field_decl_matcher = fieldDecl(hasType(pointerType())).bind("fieldDecl");
-  FieldDeclRewriter field_decl_rewriter(&replacements);
+  FieldDeclRewriter field_decl_rewriter(&replacements_printer);
   match_finder.addMatcher(field_decl_matcher, &field_decl_rewriter);
 
   // Prepare and run the tool.
@@ -156,15 +188,5 @@ int main(int argc, const char* argv[]) {
   if (result != 0)
     return result;
 
-  // Serialization format is documented in tools/clang/scripts/run_tool.py
-  llvm::outs() << "==== BEGIN EDITS ====\n";
-  for (const auto& r : replacements) {
-    std::string replacement_text = r.getReplacementText().str();
-    std::replace(replacement_text.begin(), replacement_text.end(), '\n', '\0');
-    llvm::outs() << "r:::" << r.getFilePath() << ":::" << r.getOffset()
-                 << ":::" << r.getLength() << ":::" << replacement_text << "\n";
-  }
-  llvm::outs() << "==== END EDITS ====\n";
-
   return 0;
 }

tools/clang/rewrite_raw_ptr_fields/tests/attributes-expected.cc

@@ -2,6 +2,8 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+#include "base/memory/checked_ptr.h"
+
 class SomeClass;
 
 // Based on Chromium's //base/thread_annotations.h

tools/clang/rewrite_raw_ptr_fields/tests/basics-expected.cc

@@ -2,6 +2,8 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+#include "base/memory/checked_ptr.h"
+
 class SomeClass;
 
 class MyClass {

tools/clang/rewrite_raw_ptr_fields/tests/initializers-expected.cc

@@ -2,6 +2,8 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+#include "base/memory/checked_ptr.h"
+
 class SomeClass;
 
 SomeClass* GetPointer();

tools/clang/rewrite_raw_ptr_fields/tests/qualifiers-expected.cc

@@ -2,6 +2,8 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+#include "base/memory/checked_ptr.h"
+
 class SomeClass;
 
 class MyClass {

tools/clang/rewrite_raw_ptr_fields/tests/typedefs-expected.cc

@@ -2,6 +2,8 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+#include "base/memory/checked_ptr.h"
+
 class SomeClass;
 
 // Expected rewrite: typedef CheckedPtr<SomeClass> SomeClassPtrTypedef.

tools/clang/rewrite_raw_ptr_fields/tests/various-types-expected.cc

@@ -2,6 +2,8 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+#include "base/memory/checked_ptr.h"
+
 class SomeClass;
 
 struct MyStruct {

tools/clang/scripts/apply_edits.py

@@ -22,6 +22,7 @@
 import multiprocessing
 import os
 import os.path
+import re
 import subprocess
 import sys
 
@@ -98,38 +99,164 @@ def _ResolvePath(path):
   return edits
 
 
-def _ApplyEditsToSingleFile(filename, edits):
+_PLATFORM_SUFFIX = \
+    r'(?:_(?:android|aura|chromeos|ios|linux|mac|ozone|posix|win|x11))?'
+_TEST_SUFFIX = \
+    r'(?:_(?:browser|interactive_ui|ui|unit)?test)?'
+_suffix_regex = re.compile(_PLATFORM_SUFFIX + _TEST_SUFFIX)
+
+
+def _FindPrimaryHeaderBasename(filepath):
+  """ Translates bar/foo.cc -> foo
+                 bar/foo_posix.cc -> foo
+                 bar/foo_unittest.cc -> foo
+                 bar/foo.h -> None
+  """
+  dirname, filename = os.path.split(filepath)
+  basename, extension = os.path.splitext(filename)
+  if extension == '.h':
+    return None
+  basename = _suffix_regex.sub('', basename)
+  return basename
+
+
+_INCLUDE_INSERTION_POINT_REGEX_TEMPLATE = r'''
+   ^(?!         # Match the start of the first line that is
+                # not one of the following:
+
+      \s+       # 1. Line starting with whitespace
+                #    (this includes blank lines and continuations of
+                #     C comments that start with whitespace/indentation)
+
+    | //        # 2a. A C++ comment
+    | /\*       # 2b. A C comment
+    | \*        # 2c. A continuation of a C comment (see also rule 1. above)
+
+      # 3. Include guards (Chromium-style)
+    | \#ifndef \s+ [A-Z0-9_]+_H ( | _ | __ ) \b \s* $
+    | \#define \s+ [A-Z0-9_]+_H ( | _ | __ ) \b \s* $
+
+      # 3b. Include guards (anything that repeats):
+      #     - the same <guard> has to repeat in both the #ifndef and the #define
+      #     - #define has to be "simple" - either:
+      #         - either: #define GUARD
+      #         - or    : #define GUARD 1
+    | \#ifndef \s+ (?P<guard> [A-Za-z0-9_]* ) \s* $ ( \n | \r )* ^
+      \#define \s+ (?P=guard) \s* ( | 1 \s* ) $
+    | \#define \s+ (?P=guard) \s* ( | 1 \s* ) $  # Skipping previous line.
+
+      # 4. A C/C++ system include
+    | \#include \s* < .* >
+
+      # 5. A primary header include
+      #    (%%s should be the basename returned by _FindPrimaryHeaderBasename).
+      #
+      # TODO(lukasza): Do not allow any directory below - require the top-level
+      # directory to be the same and at least one itermediate dirname to be the
+      # same.
+    | \#include \s*   "
+          [^"]* \b       # Allowing any directory
+          %s[^"/]*\.h "  # Matching both basename.h and basename_posix.h
+    )
+'''
+
+
+def _InsertNonSystemIncludeHeader(filepath, header_line_to_add, contents):
+  """ Mutates |contents| (contents of |filepath|) to #include
+      the |header_to_add
+  """
+  # Don't add the header if it is already present.
+  replacement_text = header_line_to_add
+  if replacement_text in contents:
+    return
+  replacement_text += '\n'
+
+  # Find the right insertion point.
+  #
+  # Note that we depend on a follow-up |git cl format| for the right order of
+  # headers.  Therefore we just need to find the right header group (e.g. skip
+  # system headers and the primary header).
+  primary_header_basename = _FindPrimaryHeaderBasename(filepath)
+  if primary_header_basename is None:
+    primary_header_basename = ':this:should:never:match:'
+  regex_text = _INCLUDE_INSERTION_POINT_REGEX_TEMPLATE % primary_header_basename
+  match = re.search(regex_text, contents, re.MULTILINE | re.VERBOSE)
+  assert (match is not None)
+  insertion_point = match.start()
+
+  # Extra empty line is required if the addition is not adjacent to other
+  # includes.
+  if not contents[insertion_point:].startswith('#include'):
+    replacement_text += '\n'
+
+  # Make the edit.
+  contents[insertion_point:insertion_point] = replacement_text
+
+
+def _ApplyReplacement(filepath, contents, edit, last_edit):
+  if (last_edit is not None and edit.edit_type == last_edit.edit_type
+      and edit.offset == last_edit.offset and edit.length == last_edit.length):
+    raise ValueError(('Conflicting replacement text: ' +
+                      '%s at offset %d, length %d: "%s" != "%s"\n') %
+                     (filepath, edit.offset, edit.length, edit.replacement,
+                      last_edit.replacement))
+
+  contents[edit.offset:edit.offset + edit.length] = edit.replacement
+  if not edit.replacement:
+    _ExtendDeletionIfElementIsInList(contents, edit.offset)
+
+
+def _ApplyIncludeHeader(filepath, contents, edit, last_edit):
+  header_line_to_add = '#include "%s"' % edit.replacement
+  _InsertNonSystemIncludeHeader(filepath, header_line_to_add, contents)
+
+
+def _ApplySingleEdit(filepath, contents, edit, last_edit):
+  if edit.edit_type == 'r':
+    _ApplyReplacement(filepath, contents, edit, last_edit)
+  elif edit.edit_type == 'include-user-header':
+    _ApplyIncludeHeader(filepath, contents, edit, last_edit)
+  else:
+    raise ValueError('Unrecognized edit directive "%s": %s\n' %
+                     (edit.edit_type, filepath))
+
+
+def _ApplyEditsToSingleFileContents(filepath, contents, edits):
   # Sort the edits and iterate through them in reverse order. Sorting allows
   # duplicate edits to be quickly skipped, while reversing means that
   # subsequent edits don't need to have their offsets updated with each edit
   # applied.
+  #
+  # Note that after sorting in reverse, the 'i' directives will come after 'r'
+  # directives.
+  edits.sort(reverse=True)
+
   edit_count = 0
   error_count = 0
-  edits.sort()
   last_edit = None
-  with open(filename, 'rb+') as f:
-    contents = bytearray(f.read())
-    for edit in reversed(edits):
-      if edit == last_edit:
-        continue
-      if (last_edit is not None and edit.edit_type == last_edit.edit_type and
-          edit.offset == last_edit.offset and edit.length == last_edit.length):
-        sys.stderr.write(
-            'Conflicting edit: %s at offset %d, length %d: "%s" != "%s"\n' %
-            (filename, edit.offset, edit.length, edit.replacement,
-             last_edit.replacement))
-        error_count += 1
-        continue
-
+  for edit in edits:
+    if edit == last_edit:
+      continue
+    try:
+      _ApplySingleEdit(filepath, contents, edit, last_edit)
       last_edit = edit
-      contents[edit.offset:edit.offset + edit.length] = edit.replacement
-      if not edit.replacement:
-        _ExtendDeletionIfElementIsInList(contents, edit.offset)
       edit_count += 1
+    except ValueError as err:
+      sys.stderr.write(str(err) + '\n')
+      error_count += 1
+
+  return (edit_count, error_count)
+
+
+def _ApplyEditsToSingleFile(filepath, edits):
+  with open(filepath, 'rb+') as f:
+    contents = bytearray(f.read())
+    edit_and_error_counts = _ApplyEditsToSingleFileContents(
+        filepath, contents, edits)
     f.seek(0)
     f.truncate()
     f.write(contents)
-  return (edit_count, error_count)
+  return edit_and_error_counts
 
 
 def _ApplyEdits(edits):