carbon-lang/toolchain/lex/tokenized_buffer.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 "toolchain/lex/tokenized_buffer.h"

#include <algorithm>
#include <cmath>
#include <iterator>
#include <optional>
#include <utility>

#include "common/check.h"
#include "common/string_helpers.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Format.h"
#include "llvm/Support/FormatVariadic.h"
#include "toolchain/base/shared_value_stores.h"
#include "toolchain/base/value_store_impl.h"
#include "toolchain/diagnostics/emitter.h"
#include "toolchain/lex/character_set.h"
#include "toolchain/lex/numeric_literal.h"
#include "toolchain/lex/string_literal.h"

namespace Carbon::Lex {

auto TokenizedBuffer::GetLine(TokenIndex token) const -> LineIndex {
  return FindLineIndex(token_infos_.Get(token).byte_offset());
}

auto TokenizedBuffer::GetLineNumber(TokenIndex token) const -> int {
  return GetLine(token).index + 1;
}

auto TokenizedBuffer::GetLine(CommentIndex comment) const -> LineIndex {
  return FindLineIndex(comments_.Get(comment).start);
}

auto TokenizedBuffer::GetColumnNumber(TokenIndex token) const -> int {
  const auto& token_info = token_infos_.Get(token);
  const auto& line_info =
      line_infos_.Get(FindLineIndex(token_info.byte_offset()));
  return token_info.byte_offset() - line_info.start + 1;
}

auto TokenizedBuffer::GetEndLoc(TokenIndex token) const
    -> std::pair<LineIndex, int> {
  auto line = GetLine(token);
  int column = GetColumnNumber(token);
  auto token_text = GetTokenText(token);

  if (auto [before_newline, after_newline] = token_text.rsplit('\n');
      before_newline.size() == token_text.size()) {
    // Token fits on one line, advance the column number.
    column += before_newline.size();
  } else {
    // Token contains newlines.
    line.index += before_newline.count('\n') + 1;
    column = 1 + after_newline.size();
  }

  return {line, column};
}

auto TokenizedBuffer::GetTokenText(TokenIndex token) const -> llvm::StringRef {
  const auto& token_info = token_infos_.Get(token);
  llvm::StringRef fixed_spelling = token_info.kind().fixed_spelling();
  if (!fixed_spelling.empty()) {
    return fixed_spelling;
  }

  if (token_info.kind() == TokenKind::Error) {
    return source_->text().substr(token_info.byte_offset(),
                                  token_info.error_length());
  }

  // Refer back to the source text to preserve oddities like radix or digit
  // separators the author included.
  if (token_info.kind() == TokenKind::IntLiteral ||
      token_info.kind() == TokenKind::RealLiteral) {
    std::optional<NumericLiteral> relexed_token =
        NumericLiteral::Lex(source_->text().substr(token_info.byte_offset()),
                            token_info.kind() == TokenKind::RealLiteral);
    CARBON_CHECK(relexed_token, "Could not reform numeric literal token.");
    return relexed_token->text();
  }

  // Refer back to the source text to find the original spelling, including
  // escape sequences etc.
  if (token_info.kind() == TokenKind::StringLiteral ||
      token_info.kind() == TokenKind::CharLiteral) {
    std::optional<StringLiteral> relexed_token =
        StringLiteral::Lex(source_->text().substr(token_info.byte_offset()));
    CARBON_CHECK(relexed_token, "Could not reform string literal token.");
    return relexed_token->text();
  }

  // Refer back to the source text to avoid needing to reconstruct the
  // spelling from the size.
  if (token_info.kind().is_sized_type_literal() ||
      token_info.kind().is_dollar_int_literal()) {
    llvm::StringRef suffix = source_->text()
                                 .substr(token_info.byte_offset() + 1)
                                 .take_while(IsDecimalDigit);
    return llvm::StringRef(suffix.data() - 1, suffix.size() + 1);
  }

  if (token_info.kind() == TokenKind::FileStart ||
      token_info.kind() == TokenKind::FileEnd) {
    return llvm::StringRef();
  }

  CARBON_CHECK(token_info.kind() == TokenKind::Identifier, "{0}",
               token_info.kind());

  // If this is a raw identifier, obtain its spelling from the source text.
  auto ident = value_stores_->identifiers().Get(token_info.ident_id());
  if (IsRawIdentifier(token)) {
    llvm::StringRef raw_ident =
        source_->text().substr(token_info.byte_offset(), ident.size() + 2);
    CARBON_CHECK(raw_ident[0] == 'r' && raw_ident[1] == '#' &&
                     raw_ident.substr(2) == ident,
                 "`{0}` != `r#` + `{1}`", raw_ident, ident);
    return raw_ident;
  }
  return ident;
}

auto TokenizedBuffer::GetIdentifier(TokenIndex token) const -> IdentifierId {
  const auto& token_info = token_infos_.Get(token);
  CARBON_CHECK(token_info.kind() == TokenKind::Identifier, "{0}",
               token_info.kind());
  return token_info.ident_id();
}

auto TokenizedBuffer::GetIntLiteral(TokenIndex token) const -> IntId {
  const auto& token_info = token_infos_.Get(token);
  CARBON_CHECK(token_info.kind() == TokenKind::IntLiteral, "{0}",
               token_info.kind());
  return token_info.int_id();
}

auto TokenizedBuffer::GetRealLiteral(TokenIndex token) const -> RealId {
  const auto& token_info = token_infos_.Get(token);
  CARBON_CHECK(token_info.kind() == TokenKind::RealLiteral, "{0}",
               token_info.kind());
  return token_info.real_id();
}

auto TokenizedBuffer::GetStringLiteralValue(TokenIndex token) const
    -> StringLiteralValueId {
  const auto& token_info = token_infos_.Get(token);
  CARBON_CHECK(token_info.kind() == TokenKind::StringLiteral, "{0}",
               token_info.kind());
  return token_info.string_literal_id();
}

auto TokenizedBuffer::GetCharLiteralValue(TokenIndex token) const
    -> CharLiteralValue {
  const auto& token_info = token_infos_.Get(token);
  CARBON_CHECK(token_info.kind() == TokenKind::CharLiteral, "{0}",
               token_info.kind());
  return token_info.char_literal();
}

auto TokenizedBuffer::GetTypeLiteralSize(TokenIndex token) const -> IntId {
  const auto& token_info = token_infos_.Get(token);
  CARBON_CHECK(token_info.kind().is_sized_type_literal(), "{0}",
               token_info.kind());
  return token_info.int_id();
}

auto TokenizedBuffer::GetMatchedClosingToken(TokenIndex opening_token) const
    -> TokenIndex {
  const auto& opening_token_info = token_infos_.Get(opening_token);
  CARBON_CHECK(opening_token_info.kind().is_opening_symbol(), "{0}",
               opening_token_info.kind());
  return opening_token_info.closing_token_index();
}

auto TokenizedBuffer::GetMatchedOpeningToken(TokenIndex closing_token) const
    -> TokenIndex {
  const auto& closing_token_info = token_infos_.Get(closing_token);
  CARBON_CHECK(closing_token_info.kind().is_closing_symbol(), "{0}",
               closing_token_info.kind());
  return closing_token_info.opening_token_index();
}

auto TokenizedBuffer::IsRawIdentifier(TokenIndex token) const -> bool {
  const auto& token_info = token_infos_.Get(token);
  if (token_info.kind() != TokenKind::Identifier) {
    return false;
  }
  // Check if the spelling starts `r#`. A small nuance here: we also need to
  // check the character after the `#` is an identifier start character, because
  // `r#<non-identifier-character>` is lexed as an `r` token followed by a token
  // starting with `#`. It suffices to check that character is the first
  // character of the identifier.
  auto token_text = source_->text().substr(token_info.byte_offset());
  return token_text.size() > 2 && token_text.starts_with("r#") &&
         token_text[2] ==
             value_stores_->identifiers().Get(token_info.ident_id()).front();
}

auto TokenizedBuffer::IsRecoveryToken(TokenIndex token) const -> bool {
  if (recovery_tokens_.empty()) {
    return false;
  }
  return recovery_tokens_[token.index];
}

auto TokenizedBuffer::AddPostLexingRecoveryTokenAsIdentifier(TokenIndex token)
    -> TokenIndex {
  auto kind = GetKind(token);
  CARBON_CHECK(kind.is_word(),
               "Invalid token kind {0} for recovery as identifier", kind);
  CARBON_CHECK(kind != TokenKind::Identifier, "Recovery not required");

  auto identifier_id = value_stores_->identifiers().Add(GetTokenText(token));
  return AddPostLexingRecoveryToken(
      token_infos_.Get(token).AsErrorRecoveryIdentifier(identifier_id));
}

auto TokenizedBuffer::AddPostLexingRecoveryToken(TokenInfo info) -> TokenIndex {
  // Only resize once to avoid quadratic behavior if lots of recovery tokens are
  // added.
  if (recovery_tokens_.empty()) {
    recovery_tokens_.resize(token_infos_.size());
  }

  auto token = token_infos_.Add(info);
  // Overwrite the first `FileEnd` padding entry and append a new one at the end
  // to keep `KindDataPadding` trailing entries intact.
  token_kinds_[token.index] = info.kind();
  token_kinds_.push_back(TokenKind::FileEnd);
  token_has_leading_space_.push_back(info.has_leading_space());
  recovery_tokens_.push_back(true);
  ++post_lexing_recovery_tokens_;
  return token;
}

auto TokenizedBuffer::BuildTokenArrays() -> void {
  int32_t num_tokens = token_infos_.size();
  token_kinds_.resize_for_overwrite(num_tokens + KindDataPadding);
  token_has_leading_space_.resize_for_overwrite(num_tokens);

  TokenKind* kinds = token_kinds_.data();
  bool* has_leading_space = token_has_leading_space_.data();
  for (const auto& chunk : token_infos_.chunks()) {
    // Building these flat arrays is only a performance benefit when we can do
    // so using especially efficient vectorized processing of the underlying
    // `TokenInfo` structs that are created during lexing. We rely on Clang to
    // autovectorize the inner loop below, but have to take some care to ensure
    // that works.
    //
    // Because `TokenInfo`'s bitfield loads as 32-bit words, LLVM's default x86
    // cost model caps the vector width at `256 / 32 = 8` lanes and emits
    // narrower shuffles; requesting a width of 16 packs each iteration's
    // extracted bytes into full 128-bit stores via `vpackusdw`/`vpackuswb`.
    //
    // We guard the pragma on `NDEBUG` (rather than `__OPTIMIZE__`, which is
    // also set under `fastbuild`'s `-Og` where the loop vectorizer is disabled)
    // and skip it under UBSan (whose per-iteration checks block vectorization),
    // as Clang otherwise diagnoses `-Wpass-failed` under `-Werror`.
    //
    // Last but not least, it is important to iterate over `TokenInfo` by value
    // (8 bytes) so `kind()` and `has_leading_space()` extract from a single SSA
    // value. When using a reference, Clang lowers them as two narrower memory
    // loads that AArch64 vectorizes as duplicate strided `ld2` passes over the
    // chunk instead of single-pass `ldp` loads.
#if defined(NDEBUG) && !__has_feature(undefined_behavior_sanitizer)
#pragma clang loop vectorize_width(16)
#endif
    for (TokenInfo info : chunk) {
      *kinds++ = info.kind();
      *has_leading_space++ = info.has_leading_space();
    }
  }
  std::fill_n(kinds, KindDataPadding, TokenKind::FileEnd);
}

auto TokenizedBuffer::GetIndentColumnNumber(LineIndex line) const -> int {
  return line_infos_.Get(line).indent + 1;
}

auto TokenizedBuffer::PrintWidths::Widen(const PrintWidths& widths) -> void {
  index = std::max(widths.index, index);
  kind = std::max(widths.kind, kind);
  column = std::max(widths.column, column);
  line = std::max(widths.line, line);
  indent = std::max(widths.indent, indent);
}

// Compute the printed width of a number. When numbers are printed in decimal,
// the number of digits needed is one more than the log-base-10 of the
// value. We handle a value of `zero` explicitly.
//
// This routine requires its argument to be *non-negative*.
static auto ComputeDecimalPrintedWidth(int number) -> int {
  CARBON_CHECK(number >= 0, "Negative numbers are not supported.");
  if (number == 0) {
    return 1;
  }

  return static_cast<int>(std::log10(number)) + 1;
}

auto TokenizedBuffer::GetTokenPrintWidths(TokenIndex token) const
    -> PrintWidths {
  PrintWidths widths = {};
  widths.index = ComputeDecimalPrintedWidth(token_infos_.size());
  widths.kind = GetKind(token).name().size();
  widths.line = ComputeDecimalPrintedWidth(GetLineNumber(token));
  widths.column = ComputeDecimalPrintedWidth(GetColumnNumber(token));
  widths.indent =
      ComputeDecimalPrintedWidth(GetIndentColumnNumber(GetLine(token)));
  return widths;
}

auto TokenizedBuffer::Print(llvm::raw_ostream& output_stream,
                            bool omit_file_boundary_tokens) const -> void {
  output_stream << "- filename: " << source_->filename() << "\n"
                << "  tokens:\n";

  PrintWidths widths = {};
  widths.index = ComputeDecimalPrintedWidth((token_infos_.size()));
  for (TokenIndex token : tokens()) {
    widths.Widen(GetTokenPrintWidths(token));
  }

  for (TokenIndex token : tokens()) {
    if (omit_file_boundary_tokens) {
      auto kind = GetKind(token);
      if (kind == TokenKind::FileStart || kind == TokenKind::FileEnd) {
        continue;
      }
    }
    PrintToken(output_stream, token, widths);
    output_stream << "\n";
  }

  if (has_include_in_dumps_) {
    output_stream << "  has_include_in_dumps: true\n";
  }

  if (!dump_sem_ir_ranges_.empty()) {
    output_stream << "  dump_sem_ir_ranges:\n";
    for (auto range : dump_sem_ir_ranges_) {
      output_stream << "  - {begin: " << range.begin.index
                    << ", end: " << range.end.index << "}\n";
    }
  }
}

auto TokenizedBuffer::PrintToken(llvm::raw_ostream& output_stream,
                                 TokenIndex token) const -> void {
  PrintToken(output_stream, token, {});
}

auto TokenizedBuffer::PrintToken(llvm::raw_ostream& output_stream,
                                 TokenIndex token, PrintWidths widths) const
    -> void {
  widths.Widen(GetTokenPrintWidths(token));
  int token_index = token.index;
  const auto& token_info = token_infos_.Get(token);
  LineIndex line_index = FindLineIndex(token_info.byte_offset());
  llvm::StringRef token_text = GetTokenText(token);

  // Output the main chunk using one format string. We have to do the
  // justification manually in order to use the dynamically computed widths
  // and get the quotes included.
  output_stream << llvm::formatv(
      "  - { index: {0}, kind: {1}, line: {2}, column: {3}, indent: {4}, "
      "spelling: \"{5}\"",
      llvm::format_decimal(token_index, widths.index),
      llvm::right_justify(
          llvm::formatv("\"{0}\"", token_info.kind().name()).str(),
          widths.kind + 2),
      llvm::format_decimal(GetLineNumber(token), widths.line),
      llvm::format_decimal(GetColumnNumber(token), widths.column),
      llvm::format_decimal(GetIndentColumnNumber(line_index), widths.indent),
      FormatEscaped(token_text, /*use_hex_escapes=*/true));

  switch (token_info.kind()) {
    case TokenKind::Identifier:
      output_stream << ", identifier: " << GetIdentifier(token).index;
      break;
    case TokenKind::IntLiteral:
      output_stream << ", value: \"";
      value_stores_->ints()
          .Get(GetIntLiteral(token))
          .print(output_stream, /*isSigned=*/false);
      output_stream << "\"";
      break;
    case TokenKind::RealLiteral:
      output_stream << ", value: \""
                    << value_stores_->reals().Get(GetRealLiteral(token))
                    << "\"";
      break;
    case TokenKind::StringLiteral:
      output_stream << ", value: \""
                    << FormatEscaped(value_stores_->string_literal_values().Get(
                                         GetStringLiteralValue(token)),
                                     /*use_hex_escapes=*/true)
                    << "\"";
      break;
    default:
      if (token_info.kind().is_opening_symbol()) {
        output_stream << ", closing_token: "
                      << GetMatchedClosingToken(token).index;
      } else if (token_info.kind().is_closing_symbol()) {
        output_stream << ", opening_token: "
                      << GetMatchedOpeningToken(token).index;
      }
      break;
  }

  if (token_info.has_leading_space()) {
    output_stream << ", has_leading_space: true";
  }
  if (IsRecoveryToken(token)) {
    output_stream << ", recovery: true";
  }

  output_stream << " }";
}

// Find the line index corresponding to a specific byte offset within the source
// text for this tokenized buffer.
//
// This takes advantage of the lines being sorted by their starting byte offsets
// to do a binary search for the line that contains the provided offset.
auto TokenizedBuffer::FindLineIndex(int32_t byte_offset) const -> LineIndex {
  CARBON_DCHECK(line_infos_.size() > 0);

  auto line_range = line_infos_.values();
  auto line_it =
      llvm::partition_point(line_range, [byte_offset](LineInfo line_info) {
        return line_info.start <= byte_offset;
      });
  --line_it;

  // If this isn't the first line but it starts past the end of the source, then
  // this is a synthetic line added for simplicity of lexing. Step back one
  // further to find the last non-synthetic line.
  if (line_it != line_range.begin() &&
      line_it->start == static_cast<int32_t>(source_->text().size())) {
    --line_it;
  }
  CARBON_DCHECK(line_it->start <= byte_offset);
  return LineIndex(line_it - line_range.begin());
}

auto TokenizedBuffer::IsAfterComment(TokenIndex token,
                                     CommentIndex comment_index) const -> bool {
  const auto& comment_data = comments_.Get(comment_index);
  return token_infos_.Get(token).byte_offset() > comment_data.start;
}

auto TokenizedBuffer::GetCommentText(CommentIndex comment_index) const
    -> llvm::StringRef {
  const auto& comment_data = comments_.Get(comment_index);
  return source_->text().substr(comment_data.start, comment_data.length);
}

auto TokenizedBuffer::IsTrailingComment(CommentIndex comment_index) const
    -> bool {
  return comments_.Get(comment_index).is_trailing;
}

namespace {
// The category of a full-line comment, determined by the byte after the `//`
// introducer. Adjacent full-line comments coalesce only within a category, so
// a comment's category is well defined by its first line.
enum class CommentCategory : uint8_t {
  // Whitespace, or nothing before the end of the line or file.
  Ordinary,
  // A `//@...` tooling directive.
  Directive,
  // Every other byte; the invalid spellings are lumped into one category.
  Invalid,
};
}  // namespace

// Returns the comment's category; see `CommentCategory`.
static auto GetCommentCategory(llvm::StringRef source, int32_t comment_start)
    -> CommentCategory {
  if (comment_start + 2 >= static_cast<int32_t>(source.size()) ||
      IsSpace(source[comment_start + 2])) {
    return CommentCategory::Ordinary;
  }
  if (source[comment_start + 2] == '@') {
    return CommentCategory::Directive;
  }
  return CommentCategory::Invalid;
}

auto TokenizedBuffer::AddComment(int32_t indent, int32_t start, int32_t end,
                                 bool is_trailing) -> void {
  // A comment runs forward from its start, and its length is stored in 31 bits
  // (the high bit holds `is_trailing`). A non-negative byte offset always fits
  // in 31 bits because the source size is bounded by `INT32_MAX`.
  CARBON_DCHECK(start <= end);

  // A block of adjacent full-line comments in the same category is coalesced
  // into a single comment; transitioning between ordinary comments, `//@...`
  // directives, and invalid introducers starts a new one. A trailing comment
  // is always standalone: it never extends a preceding comment, nor is it
  // extended by a following one.
  if (!is_trailing && comments_.size() > 0) {
    auto& comment = comments_.Get(CommentIndex(comments_.size() - 1));
    if (!comment.is_trailing &&
        comment.start + comment.length + indent == start &&
        GetCommentCategory(source_->text(), comment.start) ==
            GetCommentCategory(source_->text(), start)) {
      CARBON_DCHECK(comment.start <= end);
      comment.length = end - comment.start;
      return;
    }
  }
  comments_.Add({.start = start,
                 .length = static_cast<uint32_t>(end - start),
                 .is_trailing = is_trailing});
}

auto TokenizedBuffer::CollectMemUsage(MemUsage& mem_usage,
                                      llvm::StringRef label) const -> void {
  mem_usage.Collect(MemUsage::ConcatLabel(label, "allocator_"), allocator_);
  mem_usage.Collect(MemUsage::ConcatLabel(label, "token_infos_"), token_infos_);
  mem_usage.Collect(MemUsage::ConcatLabel(label, "token_kinds_"), token_kinds_);
  mem_usage.Collect(MemUsage::ConcatLabel(label, "token_has_leading_space_"),
                    token_has_leading_space_);
  mem_usage.Collect(MemUsage::ConcatLabel(label, "line_infos_"), line_infos_);
  mem_usage.Collect(MemUsage::ConcatLabel(label, "comments_"), comments_);
}

auto TokenizedBuffer::SourcePointerToDiagnosticLoc(const char* loc) const
    -> Diagnostics::ConvertedLoc {
  CARBON_CHECK(StringRefContainsPointer(source_->text(), loc),
               "location not within buffer");
  int32_t offset = loc - source_->text().begin();

  auto line_range = line_infos_.values();

  // Find the first line starting after the given location.
  const auto next_line_it = llvm::partition_point(
      line_range,
      [offset](const LineInfo& line) { return line.start <= offset; });

  // Step back one line to find the line containing the given position.
  CARBON_CHECK(next_line_it != line_range.begin(),
               "location precedes the start of the first line");
  const auto line_it = std::prev(next_line_it);
  int line_number = line_it - line_range.begin();
  int column_number = offset - line_it->start;

  // Grab the line from the buffer by slicing from this line to the next
  // minus the newline. When on the last line, instead use the start to the end
  // of the buffer.
  llvm::StringRef text = source_->text();
  llvm::StringRef line = next_line_it != line_range.end()
                             ? text.slice(line_it->start, next_line_it->start)
                             : text.substr(line_it->start);

  // Remove a newline at the end of the line if present.
  // TODO: This should expand to remove all vertical whitespace bytes at the
  // tail of the line such as CR+LF, etc.
  line.consume_back("\n");

  return {.loc = {.filename = source_->filename(),
                  .line = line,
                  .file_text = text,
                  .line_number = line_number + 1,
                  .column_number = column_number + 1},
          .last_byte_offset = offset};
}

auto TokenizedBuffer::TokenToDiagnosticLoc(TokenIndex token) const
    -> Diagnostics::ConvertedLoc {
  // Map the token location into a position within the source buffer.
  const char* token_start =
      source_->text().begin() + token_infos_.Get(token).byte_offset();

  // Find the corresponding file location.
  // TODO: Should we somehow indicate in the diagnostic location if this token
  // is a recovery token that doesn't correspond to the original source?
  auto converted = SourcePointerToDiagnosticLoc(token_start);
  converted.loc.length = GetTokenText(token).size();
  return converted;
}

auto TokenizedBuffer::OverlapsWithDumpSemIRRange(
    Lex::InclusiveTokenRange range) const -> bool {
  CARBON_CHECK(!dump_sem_ir_ranges_.empty());

  // Ranges are ordered, so we can decide overlap as soon as we find a range
  // that ends after `begin`.
  for (auto dump_range : dump_sem_ir_ranges_) {
    if (dump_range.end >= range.begin) {
      return dump_range.begin <= range.end;
    }
  }
  return false;
}

}  // namespace Carbon::Lex

namespace Carbon {
template class ValueStore<Lex::TokenIndex, Lex::TokenInfo>;
template class ValueStore<Lex::LineIndex, Lex::LineInfo>;
template class ValueStore<Lex::CommentIndex, Lex::CommentData>;
}  // namespace Carbon