carbon-lang/testing/file_test/autoupdate.cpp

// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
// Exceptions. See /LICENSE for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "testing/file_test/autoupdate.h"

#include <fstream>

#include "absl/strings/string_view.h"
#include "common/check.h"
#include "common/ostream.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLFunctionalExtras.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/Support/FormatVariadic.h"
#include "re2/re2.h"

namespace Carbon::Testing {

// Put helper classes in an anonymous namespace.
namespace {

// Converts a matched line number to an int, trimming whitespace.
static auto ParseLineNumber(absl::string_view matched_line_number) -> int {
  llvm::StringRef trimmed = matched_line_number;
  trimmed = trimmed.trim();
  // NOLINTNEXTLINE(google-runtime-int): API requirement.
  long long val;
  CARBON_CHECK(!llvm::getAsSignedInteger(trimmed, 10, val));
  return val;
}

class CheckLine : public FileTestLineBase {
 public:
  // RE2 is passed by a pointer because it doesn't support std::optional.
  explicit CheckLine(int file_number, int line_number,
                     bool line_number_re_has_file, const RE2* line_number_re,
                     std::string line)
      : FileTestLineBase(line_number),
        file_number_(file_number),
        line_number_re_has_file_(line_number_re_has_file),
        line_number_re_(line_number_re),
        line_(std::move(line)) {}

  auto Print(llvm::raw_ostream& out) const -> void override {
    out << indent_ << line_;
  }

  // When the location of the CHECK in output is known, we can set the indent
  // and its line.
  auto SetOutputLine(llvm::StringRef indent, int output_file_number,
                     int output_line_number) -> void {
    indent_ = indent;
    output_file_number_ = output_file_number;
    output_line_number_ = output_line_number;
  }

  // When the location of all lines in a file are known, we can set the line
  // offset based on the target line.
  auto RemapLineNumbers(const std::string& line_formatv,
                        llvm::function_ref<int(int)> line_remap) -> void {
    // Only need to do remappings when there's a regex.
    if (!line_number_re_) {
      return;
    }

    // If the CHECK was written to a different file from the file that it refers
    // to, leave behind an absolute line reference rather than a cross-file
    // offset.
    // TODO: We should also remap cross-file line references so that we don't
    // need multiple runs of autoupdate for the output to stabilize.
    if (output_file_number_ != file_number_) {
      return;
    }

    bool found_one = false;
    while (true) {
      // Look for a line number to replace. There may be multiple, so we
      // repeatedly check.
      absl::string_view matched_line_number;
      if (line_number_re_has_file_) {
        RE2::PartialMatch(line_, *line_number_re_, nullptr,
                          &matched_line_number);
      } else {
        RE2::PartialMatch(line_, *line_number_re_, &matched_line_number);
      }
      if (matched_line_number.empty()) {
        CARBON_CHECK(found_one) << line_;
        return;
      }
      found_one = true;

      // Calculate the offset from the CHECK line to the new line number
      // (possibly with new CHECK lines added, or some removed).
      int new_line_number = line_remap(ParseLineNumber(matched_line_number));
      int offset = new_line_number - output_line_number_;

      // Update the line offset in the CHECK line.
      const char* offset_prefix = offset < 0 ? "" : "+";
      std::string replacement = llvm::formatv(
          line_formatv.c_str(),
          llvm::formatv("[[@LINE{0}{1}]]", offset_prefix, offset));
      line_.replace(matched_line_number.data() - line_.data(),
                    matched_line_number.size(), replacement);
    }
  }

  int file_number() const { return file_number_; }

  auto is_blank() const -> bool override { return false; }

 private:
  int file_number_;
  bool line_number_re_has_file_;
  const RE2* line_number_re_;
  std::string line_;
  llvm::StringRef indent_;
  int output_file_number_ = -1;
  int output_line_number_ = -1;
};

}  // namespace

// Builds CheckLine lists for autoupdate.
static auto BuildCheckLines(
    llvm::StringRef output, const char* label,
    const llvm::SmallVector<llvm::StringRef>& filenames,
    bool line_number_re_has_file, const RE2& line_number_re,
    std::function<void(std::string&)> do_extra_check_replacements)
    -> llvm::SmallVector<CheckLine> {
  llvm::SmallVector<CheckLine> check_lines;
  if (output.empty()) {
    return check_lines;
  }

  // Prepare to look for filenames in lines.
  llvm::DenseMap<llvm::StringRef, int> file_to_number_map;
  for (auto [number, name] : llvm::enumerate(filenames)) {
    file_to_number_map.insert({name, number});
  }

  // %t substitution means we may see TEST_TMPDIR in output.
  char* tmpdir_env = getenv("TEST_TMPDIR");
  CARBON_CHECK(tmpdir_env != nullptr);
  llvm::StringRef tmpdir = tmpdir_env;

  llvm::SmallVector<llvm::StringRef> lines(llvm::split(output, '\n'));
  // It's typical that output ends with a newline, but we don't want to add a
  // blank CHECK for it.
  if (lines.back().empty()) {
    lines.pop_back();
  }

  // `{{` and `[[` are escaped as a regex matcher.
  RE2 double_brace_re(R"(\{\{)");
  RE2 double_square_bracket_re(R"(\[\[)");
  // End-of-line whitespace is replaced with a regex matcher to make it visible.
  RE2 end_of_line_whitespace_re(R"((\s+)$)");

  for (const auto& line : lines) {
    std::string check_line = llvm::formatv("// CHECK:{0}:{1}{2}", label,
                                           line.empty() ? "" : " ", line);
    RE2::Replace(&check_line, double_brace_re, R"({{\\{\\{}})");
    RE2::Replace(&check_line, double_square_bracket_re, R"({{\\[\\[}})");
    RE2::Replace(&check_line, end_of_line_whitespace_re, R"({{\1}})");

    // Ignore TEST_TMPDIR in output.
    if (auto pos = check_line.find(tmpdir); pos != std::string::npos) {
      check_line.replace(pos, tmpdir.size(), "{{.+}}");
    }

    do_extra_check_replacements(check_line);

    // Look for line information in the output. use_line_number is only set if
    // the match is correct.
    std::optional<llvm::StringRef> use_line_number;
    absl::string_view match_line_number;
    int file_number = 0;
    if (line_number_re_has_file) {
      absl::string_view match_filename;
      if (RE2::PartialMatch(check_line, line_number_re, &match_filename,
                            &match_line_number)) {
        if (auto it = file_to_number_map.find(match_filename);
            it != file_to_number_map.end()) {
          file_number = it->second;
          use_line_number = match_line_number;
        }
      }
    } else {
      // There's no file association, so we only look at the line, and assume it
      // refers to the main file.
      if (RE2::PartialMatch(check_line, line_number_re, &match_line_number)) {
        use_line_number = match_line_number;
      }
    }
    int line_number = use_line_number ? ParseLineNumber(*use_line_number) : -1;
    check_lines.push_back(
        CheckLine(file_number, line_number, line_number_re_has_file,
                  use_line_number ? &line_number_re : nullptr, check_line));
  }

  return check_lines;
}

auto AutoupdateFileTest(
    const std::filesystem::path& file_test_path, llvm::StringRef input_content,
    const llvm::SmallVector<llvm::StringRef>& filenames,
    int autoupdate_line_number,
    llvm::SmallVector<llvm::SmallVector<FileTestLine>>& non_check_lines,
    llvm::StringRef stdout, llvm::StringRef stderr,
    FileTestLineNumberReplacement line_number_replacement,
    std::function<void(std::string&)> do_extra_check_replacements) -> bool {
  RE2 line_number_re(line_number_replacement.pattern);
  CARBON_CHECK(line_number_re.ok()) << "Invalid line replacement RE2: `"
                                    << line_number_replacement.pattern << "`";

  // Prepare CHECK lines.
  llvm::SmallVector<CheckLine> stdout_check_lines = BuildCheckLines(
      stdout, "STDOUT", filenames, line_number_replacement.has_file,
      line_number_re, do_extra_check_replacements);
  llvm::SmallVector<CheckLine> stderr_check_lines = BuildCheckLines(
      stderr, "STDERR", filenames, line_number_replacement.has_file,
      line_number_re, do_extra_check_replacements);
  auto* stdout_check_line = stdout_check_lines.begin();
  auto* stderr_check_line = stderr_check_lines.begin();

  bool any_attached_stdout_lines = std::any_of(
      stdout_check_lines.begin(), stdout_check_lines.end(),
      [&](const CheckLine& line) { return line.line_number() != -1; });

  // All CHECK lines are suppressed until we reach AUTOUPDATE.
  bool reached_autoupdate = false;

  const FileTestLine blank_line(-1, "");

  // Stitch together content.
  llvm::SmallVector<const FileTestLineBase*> new_lines;
  for (auto [file_number_as_size_t, filename, non_check_file] :
       llvm::enumerate(filenames, non_check_lines)) {
    auto file_number = static_cast<int>(file_number_as_size_t);
    llvm::DenseMap<int, int> output_line_remap;
    llvm::SmallVector<CheckLine*> check_lines_this_file;
    int output_line_number = 0;

    // Add all check lines from the given vector until we reach a check line
    // attached to a line later than `to_line_number`.
    auto add_check_lines = [&](const llvm::SmallVector<CheckLine>& lines,
                               CheckLine*& line, int to_line_number,
                               llvm::StringRef indent) {
      for (; line != lines.end() && (line->file_number() < file_number ||
                                     (line->file_number() == file_number &&
                                      line->line_number() <= to_line_number));
           ++line) {
        new_lines.push_back(line);
        line->SetOutputLine(indent, file_number, ++output_line_number);
        check_lines_this_file.push_back(line);
      }
    };

    // Looping through the original file, print check lines preceding each
    // original line.
    for (const auto& non_check_line : non_check_file) {
      // If there are any non-check lines with an invalid line_number, it's
      // something like a split directive which shouldn't increment
      // output_line_number.
      if (non_check_line.line_number() < 1) {
        new_lines.push_back(&non_check_line);
        continue;
      }

      // STDERR check lines are placed before the line they refer to, or as
      // early as possible if they don't refer to a line. Include all STDERR
      // lines until we find one that wants to go later in the file.
      if (reached_autoupdate) {
        add_check_lines(stderr_check_lines, stderr_check_line,
                        non_check_line.line_number(), non_check_line.indent());
      } else if (autoupdate_line_number == non_check_line.line_number()) {
        // This is the AUTOUPDATE line, so we'll print it, then start printing
        // CHECK lines.
        reached_autoupdate = true;
      }

      new_lines.push_back(&non_check_line);
      CARBON_CHECK(
          output_line_remap
              .insert({non_check_line.line_number(), ++output_line_number})
              .second);

      // STDOUT check lines are placed after the line they refer to, or at the
      // end of the file if none of them refers to a line.
      if (reached_autoupdate && any_attached_stdout_lines) {
        // Include any early STDERR lines now, so that the initial batch of
        // CHECK lines have STDERR before STDOUT.
        if (autoupdate_line_number == non_check_line.line_number()) {
          add_check_lines(stderr_check_lines, stderr_check_line,
                          non_check_line.line_number(),
                          non_check_line.indent());
        }

        add_check_lines(stdout_check_lines, stdout_check_line,
                        non_check_line.line_number(), non_check_line.indent());
      }
    }

    // This should always be true after the first file is processed.
    CARBON_CHECK(reached_autoupdate);

    // At the end of the last file, print remaining check lines which -- for
    // whatever reason -- come after all original lines.
    if (file_number == static_cast<int>(filenames.size()) - 1 &&
        (stderr_check_line != stderr_check_lines.end() ||
         stdout_check_line != stdout_check_lines.end())) {
      // Ensure there's a blank line before any trailing CHECKs.
      if (!new_lines.empty() && !new_lines.back()->is_blank()) {
        new_lines.push_back(&blank_line);
        ++output_line_number;
      }

      add_check_lines(stderr_check_lines, stderr_check_line, INT_MAX, "");
      add_check_lines(stdout_check_lines, stdout_check_line, INT_MAX, "");
    }

    // Update all remapped lines in CHECK output.
    for (auto* offset_check_line : check_lines_this_file) {
      int last_non_check_line = non_check_file.back().line_number();
      offset_check_line->RemapLineNumbers(
          line_number_replacement.line_formatv, [&](int old_line_number) {
            // Map old non-check lines to their new line numbers.
            auto remapped = output_line_remap.find(old_line_number);
            if (remapped != output_line_remap.end()) {
              return remapped->second;
            }

            // Map any reference to a line past the final non-check line to
            // the new end-of-file. We assume that any such reference is
            // referring to the end of file, not to some specific CHECK
            // comment.
            if (old_line_number > last_non_check_line) {
              return output_line_number;
            }

            // Line didn't get remapped; maybe it refers to a CHECK line.
            // We can't express that as an offset, just leave it as-is.
            return old_line_number;
          });
    }
  }

  // Generate the autoupdated file.
  std::string new_content;
  llvm::raw_string_ostream new_content_stream(new_content);
  for (const auto& line : new_lines) {
    line->Print(new_content_stream);
    new_content_stream << '\n';
  }

  // Update the file on disk if needed.
  if (new_content == input_content) {
    return false;
  }
  std::ofstream out(file_test_path);
  out << new_content;
  return true;
}

}  // namespace Carbon::Testing