// 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/lex.h"

#include <algorithm>
#include <array>
#include <limits>
#include <optional>
#include <utility>

#include "common/check.h"
#include "common/vlog.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/StringSwitch.h"
#include "llvm/Support/Compiler.h"
#include "toolchain/base/kind_switch.h"
#include "toolchain/base/shared_value_stores.h"
#include "toolchain/diagnostics/format_providers.h"
#include "toolchain/lex/character_set.h"
#include "toolchain/lex/helpers.h"
#include "toolchain/lex/mismatched_brackets.h"
#include "toolchain/lex/numeric_literal.h"
#include "toolchain/lex/string_literal.h"
#include "toolchain/lex/token_index.h"
#include "toolchain/lex/token_info.h"
#include "toolchain/lex/token_kind.h"
#include "toolchain/lex/tokenized_buffer.h"

#if __ARM_NEON
#include <arm_neon.h>
#define CARBON_USE_SIMD 1
#elif __x86_64__
#include <x86intrin.h>
#define CARBON_USE_SIMD 1
#else
#define CARBON_USE_SIMD 0
#endif

namespace Carbon::Lex {

// Implementation of the lexer logic itself.
//
// The design is that lexing can loop over the source buffer, consuming it into
// tokens by calling into this API. This class handles the state and breaks down
// the different lexing steps that may be used. It directly updates the provided
// tokenized buffer with the lexed tokens.
//
// We'd typically put this in an anonymous namespace, but it is `friend`-ed by
// the `TokenizedBuffer`. One of the important benefits of being in an anonymous
// namespace is having internal linkage. That allows the optimizer to much more
// aggressively inline away functions that are called in only one place. We keep
// that benefit for now by using the `internal_linkage` attribute.
//
// TODO: Investigate ways to refactor the code that allow moving this into an
// anonymous namespace without overly exposing implementation details of the
// `TokenizedBuffer` or undermining the performance constraints of the lexer.
class [[clang::internal_linkage]] Lexer {
 public:
  // Symbolic result of a lexing action. This indicates whether we successfully
  // lexed a token, or whether other lexing actions should be attempted.
  //
  // While it wraps a simple boolean state, its API both helps make the failures
  // more self documenting, and by consuming the actual token constructively
  // when one is produced, it helps ensure the correct result is returned.
  class LexResult {
   public:
    // Consumes (and discard) a valid token to construct a result
    // indicating a token has been produced. Relies on implicit conversions.
    explicit(false) LexResult(TokenIndex /*discarded_token*/)
        : LexResult(true) {}

    // Returns a result indicating no token was produced.
    static auto NoMatch() -> LexResult { return LexResult(false); }

    // Tests whether a token was produced by the lexing routine, and
    // the lexer can continue forming tokens.
    explicit operator bool() const { return formed_token_; }

   private:
    explicit LexResult(bool formed_token) : formed_token_(formed_token) {}

    bool formed_token_;
  };

  Lexer(const LexOptions& options, SharedValueStores& value_stores,
        SourceBuffer& source, Diagnostics::Consumer& consumer)
      : options_(options),
        buffer_(value_stores, source),
        consumer_(consumer),
        emitter_(&consumer_, &buffer_),
        token_emitter_(&consumer_, &buffer_) {}

  // Find all line endings and create the line data structures.
  //
  // Explicitly kept out-of-line because this is a significant loop that is
  // useful to have in the profile and it doesn't simplify by inlining at all.
  // But because it can, the compiler will flatten this otherwise.
  [[gnu::noinline]] auto MakeLines(llvm::StringRef source_text) -> void;

  auto current_line() -> LineIndex { return line_index_; }

  auto current_line_info() -> LineInfo& {
    return buffer_.line_infos_.Get(line_index_);
  }

  auto next_line() -> LineIndex { return LineIndex(line_index_.index + 1); }

  auto next_line_info() -> LineInfo& {
    return buffer_.line_infos_.Get(next_line());
  }

  // Note when the lexer has encountered whitespace, and the next lexed token
  // should reflect that it was preceded by some amount of whitespace.
  auto NoteWhitespace() -> void { has_leading_space_ = true; }

  // Add a lexed token to the tokenized buffer, and reset any token-specific
  // state tracked in the lexer for the next token.
  auto AddLexedToken(TokenInfo info) -> TokenIndex {
    has_leading_space_ = false;
    return buffer_.AddToken(info);
  }

  // Lexes a token with no payload: builds the correctly encoded token info,
  // adds it to the tokenized buffer and returns the token index.
  auto LexToken(TokenKind kind, int32_t byte_offset) -> TokenIndex {
    // Check that we don't accidentally call this for one of the token kinds
    // that *always* has a payload up front.
    CARBON_DCHECK(!kind.IsOneOf(
        {TokenKind::Identifier, TokenKind::StringLiteral, TokenKind::IntLiteral,
         TokenKind::IntTypeLiteral, TokenKind::UnsignedIntTypeLiteral,
         TokenKind::FloatTypeLiteral, TokenKind::RealLiteral,
         TokenKind::Error}));
    return AddLexedToken(TokenInfo(kind, has_leading_space_, byte_offset));
  }

  // Lexes a token with a payload: builds the correctly encoded token info,
  // adds it to the tokenized buffer and returns the token index.
  auto LexTokenWithPayload(TokenKind kind, int token_payload,
                           int32_t byte_offset) -> TokenIndex {
    return AddLexedToken(
        TokenInfo(kind, has_leading_space_, token_payload, byte_offset));
  }

  auto SkipHorizontalWhitespace(llvm::StringRef source_text, ssize_t& position)
      -> void;

  // Starts a new line, skipping whitespace and setting the indent.
  auto AdvanceToLine(llvm::StringRef source_text, ssize_t& position,
                     LineIndex to_line_index) -> void;

  auto LexHorizontalWhitespace(llvm::StringRef source_text, ssize_t& position)
      -> void;

  auto LexVerticalWhitespace(llvm::StringRef source_text, ssize_t& position)
      -> void;

  auto LexCR(llvm::StringRef source_text, ssize_t& position) -> void;

  auto LexCommentOrSlash(llvm::StringRef source_text, ssize_t& position)
      -> void;

  auto LexComment(llvm::StringRef source_text, ssize_t& position) -> void;

  // Determines whether a real literal can be formed at the current location.
  // This is the case unless the preceding token is `.` or `->` and there is no
  // intervening whitespace.
  auto CanFormRealLiteral() -> bool;

  auto LexNumericLiteral(llvm::StringRef source_text, ssize_t& position)
      -> LexResult;

  auto LexStringLiteral(llvm::StringRef source_text, ssize_t& position)
      -> LexResult;

  auto LexOneCharSymbolToken(llvm::StringRef source_text, TokenKind kind,
                             ssize_t& position) -> TokenIndex;

  auto LexOpeningSymbolToken(llvm::StringRef source_text, TokenKind kind,
                             ssize_t& position) -> LexResult;

  auto LexClosingSymbolToken(llvm::StringRef source_text, TokenKind kind,
                             ssize_t& position) -> LexResult;

  auto LexSymbolToken(llvm::StringRef source_text, ssize_t& position)
      -> LexResult;

  // Given a word that has already been lexed, determine whether it is a type
  // literal and if so form the corresponding token.
  auto LexWordAsTypeLiteralToken(llvm::StringRef word, int32_t byte_offset)
      -> LexResult;

  // Given a lexed word, determine whether it is a dollar int literal and if so
  // form the corresponding token,
  auto LexWordAsDollarIntLiteralToken(llvm::StringRef word, int32_t byte_offset)
      -> LexResult;

  auto LexKeywordOrIdentifier(llvm::StringRef source_text, ssize_t& position)
      -> LexResult;

  auto LexHash(llvm::StringRef source_text, ssize_t& position) -> LexResult;

  auto LexError(llvm::StringRef source_text, ssize_t& position) -> LexResult;

  auto LexFileStart(llvm::StringRef source_text, ssize_t& position) -> void;

  auto LexFileEnd(llvm::StringRef source_text, ssize_t position) -> void;

  // Perform final checking and cleanup that should be done once we have
  // finished lexing the whole file, and before we consider the tokenized buffer
  // to be complete.
  auto Finalize() -> void;

  auto DiagnoseAndFixMismatchedBrackets() -> void;

  // The main entry point for dispatching through the lexer's table. This method
  // should always fully consume the source text.
  auto Lex() && -> TokenizedBuffer;

  // Checks for an ends a `DumpSemIRRange` that's missing an explicit end
  // marker.
  auto EndDumpSemIRRangeIfIncomplete(const char* diag_loc) -> void;

  auto has_dump_sem_ir_ranges() -> bool {
    return buffer_.has_dump_sem_ir_ranges();
  }

 private:
  class ErrorRecoveryBuffer;

  // Handles `//@dump-sem-ir-begin` for a `DumpSemIRRange`.
  auto BeginDumpSemIRRange(const char* diag_loc) -> void;

  // Handles `//@dump-sem-ir-end` for a `DumpSemIRRange`.
  auto EndDumpSemIRRange(const char* diag_loc) -> void;

  LexOptions options_;

  TokenizedBuffer buffer_;

  LineIndex line_index_ = LineIndex::None;

  // Tracks whether the lexer has encountered whitespace that will be leading
  // whitespace for the next lexed token. Reset after each token lexed.
  bool has_leading_space_ = false;

  llvm::SmallVector<TokenIndex> open_groups_;
  bool has_mismatched_brackets_ = false;

  Diagnostics::ErrorTrackingConsumer consumer_;

  TokenizedBuffer::SourcePointerDiagnosticEmitter emitter_;

  TokenizedBuffer::TokenDiagnosticEmitter token_emitter_;
};

#if CARBON_USE_SIMD
namespace {
#if __ARM_NEON
using SimdMaskT = uint8x16_t;
#elif __x86_64__
using SimdMaskT = __m128i;
#else
#error "Unsupported SIMD architecture!"
#endif
using SimdMaskArrayT = std::array<SimdMaskT, sizeof(SimdMaskT) + 1>;
}  // namespace
// A table of masks to include 0-16 bytes of an SSE register.
static constexpr SimdMaskArrayT PrefixMasks = []() constexpr {
  SimdMaskArrayT masks = {};
  for (int i = 1; i < static_cast<int>(masks.size()); ++i) {
    masks[i] =
        // The SIMD types and constexpr require a C-style cast.
        // NOLINTNEXTLINE(google-readability-casting)
        (SimdMaskT)(std::numeric_limits<unsigned __int128>::max() >>
                    ((sizeof(SimdMaskT) - i) * 8));
  }
  return masks;
}();
#endif  // CARBON_USE_SIMD

// A table of booleans that we can use to classify bytes as being valid
// identifier start. This is used by raw identifier detection.
static constexpr std::array<bool, 256> IsIdStartByteTable = [] {
  std::array<bool, 256> table = {};
  for (char c = 'A'; c <= 'Z'; ++c) {
    table[c] = true;
  }
  for (char c = 'a'; c <= 'z'; ++c) {
    table[c] = true;
  }
  table['_'] = true;
  table['$'] = true;
  return table;
}();

// A table of booleans that we can use to classify bytes as being valid
// identifier (or keyword) characters. This is used in the generic,
// non-vectorized fallback code to scan for length of an identifier.
static constexpr std::array<bool, 256> IsIdByteTable = [] {
  std::array<bool, 256> table = IsIdStartByteTable;
  for (char c = '0'; c <= '9'; ++c) {
    table[c] = true;
  }
  // Identifiers can only have `$` in start.
  table['$'] = false;
  return table;
}();

// Baseline scalar version, also available for scalar-fallback in SIMD code.
// Uses `ssize_t` for performance when indexing in the loop.
//
// TODO: This assumes all Unicode characters are non-identifiers.
static auto ScanForIdentifierPrefixScalar(llvm::StringRef text, ssize_t i)
    -> llvm::StringRef {
  const ssize_t size = text.size();
  if (i == 0 && !text.empty()) {
    if (!IsIdStartByteTable[static_cast<unsigned char>(text[i])]) {
      return {};
    }
    ++i;
  }
  while (i < size && IsIdByteTable[static_cast<unsigned char>(text[i])]) {
    ++i;
  }

  return text.substr(0, i);
}

#if CARBON_USE_SIMD && __x86_64__
// The SIMD code paths uses a scheme derived from the techniques in Geoff
// Langdale and Daniel Lemire's work on parsing JSON[1]. Specifically, that
// paper outlines a technique of using two 4-bit indexed in-register look-up
// tables (LUTs) to classify bytes in a branchless SIMD code sequence.
//
// [1]: https://arxiv.org/pdf/1902.08318.pdf
//
// The goal is to get a bit mask classifying different sets of bytes. For each
// input byte, we first test for a high bit indicating a UTF-8 encoded Unicode
// character. Otherwise, we want the mask bits to be set with the following
// logic derived by inspecting the high nibble and low nibble of the input:
// bit0 = 1 for `_`: high `0x5` and low `0xF`
// bit1 = 1 for `0-9`: high `0x3` and low `0x0` - `0x9`
// bit2 = 1 for `A-O` and `a-o`: high `0x4` or `0x6` and low `0x1` - `0xF`
// bit3 = 1 for `P-Z` and 'p-z': high `0x5` or `0x7` and low `0x0` - `0xA`
// bit4 = unused
// bit5 = unused
// bit6 = unused
// bit7 = unused
//
// No bits set means definitively non-ID ASCII character.
//
// Bits 4-7 remain unused if we need to classify more characters.
namespace {
// Struct used to implement the nibble LUT for SIMD implementations.
//
// Forced to 16-byte alignment to ensure we can load it easily in SIMD code.
struct alignas(16) NibbleLUT {
  auto Load() const -> __m128i {
    return _mm_load_si128(reinterpret_cast<const __m128i*>(this));
  }

  uint8_t nibble_0;
  uint8_t nibble_1;
  uint8_t nibble_2;
  uint8_t nibble_3;
  uint8_t nibble_4;
  uint8_t nibble_5;
  uint8_t nibble_6;
  uint8_t nibble_7;
  uint8_t nibble_8;
  uint8_t nibble_9;
  uint8_t nibble_a;
  uint8_t nibble_b;
  uint8_t nibble_c;
  uint8_t nibble_d;
  uint8_t nibble_e;
  uint8_t nibble_f;
};
}  // namespace

static constexpr NibbleLUT HighLUT = {
    .nibble_0 = 0b0000'0000,
    .nibble_1 = 0b0000'0000,
    .nibble_2 = 0b0000'0000,
    .nibble_3 = 0b0000'0010,
    .nibble_4 = 0b0000'0100,
    .nibble_5 = 0b0000'1001,
    .nibble_6 = 0b0000'0100,
    .nibble_7 = 0b0000'1000,
    .nibble_8 = 0b1000'0000,
    .nibble_9 = 0b1000'0000,
    .nibble_a = 0b1000'0000,
    .nibble_b = 0b1000'0000,
    .nibble_c = 0b1000'0000,
    .nibble_d = 0b1000'0000,
    .nibble_e = 0b1000'0000,
    .nibble_f = 0b1000'0000,
};
static constexpr NibbleLUT LowLUT = {
    .nibble_0 = 0b1000'1010,
    .nibble_1 = 0b1000'1110,
    .nibble_2 = 0b1000'1110,
    .nibble_3 = 0b1000'1110,
    .nibble_4 = 0b1000'1110,
    .nibble_5 = 0b1000'1110,
    .nibble_6 = 0b1000'1110,
    .nibble_7 = 0b1000'1110,
    .nibble_8 = 0b1000'1110,
    .nibble_9 = 0b1000'1110,
    .nibble_a = 0b1000'1100,
    .nibble_b = 0b1000'0100,
    .nibble_c = 0b1000'0100,
    .nibble_d = 0b1000'0100,
    .nibble_e = 0b1000'0100,
    .nibble_f = 0b1000'0101,
};

static auto ScanForIdentifierPrefixX86(llvm::StringRef text)
    -> llvm::StringRef {
  const auto high_lut = HighLUT.Load();
  const auto low_lut = LowLUT.Load();

  // Use `ssize_t` for performance here as we index memory in a tight loop.
  ssize_t i = 0;
  if (!text.empty()) {
    if (!IsIdStartByteTable[static_cast<unsigned char>(text[i])]) {
      return {};
    }
    ++i;
  }

  const ssize_t size = text.size();
  while ((i + 16) <= size) {
    __m128i input =
        _mm_loadu_si128(reinterpret_cast<const __m128i*>(text.data() + i));

    // The high bits of each byte indicate a non-ASCII character encoded using
    // UTF-8. Test those and fall back to the scalar code if present. These
    // bytes will also cause spurious zeros in the LUT results, but we can
    // ignore that because we track them independently here.
#if __SSE4_1__
    if (!_mm_test_all_zeros(_mm_set1_epi8(0x80), input)) {
      break;
    }
#else
    if (_mm_movemask_epi8(input) != 0) {
      break;
    }
#endif

    // Do two LUT lookups and mask the results together to get the results for
    // both low and high nibbles. Note that we don't need to mask out the high
    // bit of input here because we track that above for UTF-8 handling.
    __m128i low_mask = _mm_shuffle_epi8(low_lut, input);
    // Note that the input needs to be masked to only include the high nibble or
    // we could end up with bit7 set forcing the result to a zero byte.
    __m128i input_high =
        _mm_and_si128(_mm_srli_epi32(input, 4), _mm_set1_epi8(0x0f));
    __m128i high_mask = _mm_shuffle_epi8(high_lut, input_high);
    __m128i mask = _mm_and_si128(low_mask, high_mask);

    // Now compare to find the completely zero bytes.
    __m128i id_byte_mask_vec = _mm_cmpeq_epi8(mask, _mm_setzero_si128());
    int tail_ascii_mask = _mm_movemask_epi8(id_byte_mask_vec);

    // Check if there are bits in the tail mask, which means zero bytes and the
    // end of the identifier. We could do this without materializing the scalar
    // mask on more recent CPUs, but we generally expect the median length we
    // encounter to be <16 characters and so we avoid the extra instruction in
    // that case and predict this branch to succeed so it is laid out in a
    // reasonable way.
    if (LLVM_LIKELY(tail_ascii_mask != 0)) {
      // Move past the definitively classified bytes that are part of the
      // identifier, and return the complete identifier text.
      i += __builtin_ctz(tail_ascii_mask);
      return text.substr(0, i);
    }
    i += 16;
  }

  return ScanForIdentifierPrefixScalar(text, i);
}

#endif  // CARBON_USE_SIMD && __x86_64__

// Scans the provided text and returns the prefix `StringRef` of contiguous
// identifier characters.
//
// This is a performance sensitive function and where profitable uses vectorized
// code sequences to optimize its scanning. When modifying, the identifier
// lexing benchmarks should be checked for regressions.
//
// Identifier characters here are currently the ASCII characters `[0-9A-Za-z_]`.
//
// TODO: Currently, this code does not implement Carbon's design for Unicode
// characters in identifiers. It does work on UTF-8 code unit sequences, but
// currently considers non-ASCII characters to be non-identifier characters.
// Some work has been done to ensure the hot loop, while optimized, retains
// enough information to add Unicode handling without completely destroying the
// relevant optimizations.
static auto ScanForIdentifierPrefix(llvm::StringRef text) -> llvm::StringRef {
  // Dispatch to an optimized architecture optimized routine.
#if CARBON_USE_SIMD && __x86_64__
  return ScanForIdentifierPrefixX86(text);
#elif CARBON_USE_SIMD && __ARM_NEON
  // Somewhat surprisingly, there is basically nothing worth doing in SIMD on
  // Arm to optimize this scan. The Neon SIMD operations end up requiring you to
  // move from the SIMD unit to the scalar unit in the critical path of finding
  // the offset of the end of an identifier. Current ARM cores make the code
  // sequences here (quite) unpleasant. For example, on Apple M1 and similar
  // cores, the latency is as much as 10 cycles just to extract from the vector.
  // SIMD might be more interesting on Neoverse cores, but it'd be nice to avoid
  // core-specific tunings at this point.
  //
  // If this proves problematic and critical to optimize, the current leading
  // theory is to have the newline searching code also create a bitmask for the
  // entire source file of identifier and non-identifier bytes, and then use the
  // bit-counting instructions here to do a fast scan of that bitmask. However,
  // crossing that bridge will add substantial complexity to the newline
  // scanner, and so currently we just use a boring scalar loop that pipelines
  // well.
#endif
  return ScanForIdentifierPrefixScalar(text, 0);
}

using DispatchFunctionT = auto(Lexer& lexer, llvm::StringRef source_text,
                               ssize_t position) -> void;
using DispatchTableT = std::array<DispatchFunctionT*, 256>;

static constexpr std::array<TokenKind, 256> OneCharTokenKindTable = [] {
  std::array<TokenKind, 256> table = {};
#define CARBON_ONE_CHAR_SYMBOL_TOKEN(TokenName, Spelling) \
  table[(Spelling)[0]] = TokenKind::TokenName;
#define CARBON_OPENING_GROUP_SYMBOL_TOKEN(TokenName, Spelling, ClosingName) \
  table[(Spelling)[0]] = TokenKind::TokenName;
#define CARBON_CLOSING_GROUP_SYMBOL_TOKEN(TokenName, Spelling, OpeningName) \
  table[(Spelling)[0]] = TokenKind::TokenName;
#include "toolchain/lex/token_kind.def"
  return table;
}();

// We use a collection of static member functions for table-based dispatch to
// lexer methods. These are named static member functions so that they show up
// helpfully in profiles and backtraces, but they tend to not contain the
// interesting logic and simply delegate to the relevant methods. All of their
// signatures need to be exactly the same however in order to ensure we can
// build efficient dispatch tables out of them. All of them end by doing a
// must-tail return call to this routine. It handles continuing the dispatch
// chain.
static auto DispatchNext(Lexer& lexer, llvm::StringRef source_text,
                         ssize_t position) -> void;

// Define a set of dispatch functions that simply forward to a method that
// lexes a token. This includes validating that an actual token was produced,
// and continuing the dispatch.
#define CARBON_DISPATCH_LEX_TOKEN(LexMethod)                                 \
  static auto Dispatch##LexMethod(Lexer& lexer, llvm::StringRef source_text, \
                                  ssize_t position) -> void {                \
    Lexer::LexResult result = lexer.LexMethod(source_text, position);        \
    CARBON_CHECK(result, "Failed to form a token!");                         \
    [[clang::musttail]] return DispatchNext(lexer, source_text, position);   \
  }
CARBON_DISPATCH_LEX_TOKEN(LexError)
CARBON_DISPATCH_LEX_TOKEN(LexSymbolToken)
CARBON_DISPATCH_LEX_TOKEN(LexKeywordOrIdentifier)
CARBON_DISPATCH_LEX_TOKEN(LexHash)
CARBON_DISPATCH_LEX_TOKEN(LexNumericLiteral)
CARBON_DISPATCH_LEX_TOKEN(LexStringLiteral)

// A set of custom dispatch functions that preselect the symbol token to lex.
#define CARBON_DISPATCH_LEX_SYMBOL_TOKEN(LexMethod)                          \
  static auto Dispatch##LexMethod##SymbolToken(                              \
      Lexer& lexer, llvm::StringRef source_text, ssize_t position) -> void { \
    Lexer::LexResult result = lexer.LexMethod##SymbolToken(                  \
        source_text,                                                         \
        OneCharTokenKindTable[static_cast<unsigned char>(                    \
            source_text[position])],                                         \
        position);                                                           \
    CARBON_CHECK(result, "Failed to form a token!");                         \
    [[clang::musttail]] return DispatchNext(lexer, source_text, position);   \
  }
CARBON_DISPATCH_LEX_SYMBOL_TOKEN(LexOneChar)
CARBON_DISPATCH_LEX_SYMBOL_TOKEN(LexOpening)
CARBON_DISPATCH_LEX_SYMBOL_TOKEN(LexClosing)

// Define a set of non-token dispatch functions that handle things like
// whitespace and comments.
#define CARBON_DISPATCH_LEX_NON_TOKEN(LexMethod)                             \
  static auto Dispatch##LexMethod(Lexer& lexer, llvm::StringRef source_text, \
                                  ssize_t position) -> void {                \
    lexer.LexMethod(source_text, position);                                  \
    [[clang::musttail]] return DispatchNext(lexer, source_text, position);   \
  }
CARBON_DISPATCH_LEX_NON_TOKEN(LexHorizontalWhitespace)
CARBON_DISPATCH_LEX_NON_TOKEN(LexVerticalWhitespace)
CARBON_DISPATCH_LEX_NON_TOKEN(LexCR)
CARBON_DISPATCH_LEX_NON_TOKEN(LexCommentOrSlash)

// Build a table of function pointers that we can use to dispatch to the
// correct lexer routine based on the first byte of source text.
//
// While it is tempting to simply use a `switch` on the first byte and
// dispatch with cases into this, in practice that doesn't produce great code.
// There seem to be two issues that are the root cause.
//
// First, there are lots of different values of bytes that dispatch to a
// fairly small set of routines, and then some byte values that dispatch
// differently for each byte. This pattern isn't one that the compiler-based
// lowering of switches works well with -- it tries to balance all the cases,
// and in doing so emits several compares and other control flow rather than a
// simple jump table.
//
// Second, with a `case`, it isn't as obvious how to create a single, uniform
// interface that is effective for *every* byte value, and thus makes for a
// single consistent table-based dispatch. By forcing these to be function
// pointers, we also coerce the code to use a strictly homogeneous structure
// that can form a single dispatch table.
//
// These two actually interact -- the second issue is part of what makes the
// non-table lowering in the first one desirable for many switches and cases.
//
// Ultimately, when table-based dispatch is such an important technique, we
// get better results by taking full control and manually creating the
// dispatch structures.
//
// The functions in this table also use tail-recursion to implement the loop
// of the lexer. This is based on the technique described more fully for any
// kind of byte-stream loop structure here:
// https://blog.reverberate.org/2021/04/21/musttail-efficient-interpreters.html
static constexpr auto MakeDispatchTable() -> DispatchTableT {
  DispatchTableT table = {};
  // First set the table entries to dispatch to our error token handler as the
  // base case. Everything valid comes from an override below.
  for (int i = 0; i < 256; ++i) {
    table[i] = &DispatchLexError;
  }

  // Symbols have some special dispatching. First, set the first character of
  // each symbol token spelling to dispatch to the symbol lexer. We don't
  // provide a pre-computed token here, so the symbol lexer will compute the
  // exact symbol token kind. We'll override this with more specific dispatch
  // below.
#define CARBON_SYMBOL_TOKEN(TokenName, Spelling) \
  table[(Spelling)[0]] = &DispatchLexSymbolToken;
#include "toolchain/lex/token_kind.def"

  // Now special cased single-character symbols that are guaranteed to not
  // join with another symbol. These are grouping symbols, terminators,
  // or separators in the grammar and have a good reason to be
  // orthogonal to any other punctuation. We do this separately because this
  // needs to override some of the generic handling above, and provide a
  // custom token.
#define CARBON_ONE_CHAR_SYMBOL_TOKEN(TokenName, Spelling) \
  table[(Spelling)[0]] = &DispatchLexOneCharSymbolToken;
#define CARBON_OPENING_GROUP_SYMBOL_TOKEN(TokenName, Spelling, ClosingName) \
  table[(Spelling)[0]] = &DispatchLexOpeningSymbolToken;
#define CARBON_CLOSING_GROUP_SYMBOL_TOKEN(TokenName, Spelling, OpeningName) \
  table[(Spelling)[0]] = &DispatchLexClosingSymbolToken;
#include "toolchain/lex/token_kind.def"

  // Override the handling for `/` to consider comments as well as a `/`
  // symbol.
  table['/'] = &DispatchLexCommentOrSlash;

  table['_'] = &DispatchLexKeywordOrIdentifier;
  table['$'] = &DispatchLexKeywordOrIdentifier;
  // Note that we don't use `llvm::seq` because this needs to be `constexpr`
  // evaluated.
  for (unsigned char c = 'a'; c <= 'z'; ++c) {
    table[c] = &DispatchLexKeywordOrIdentifier;
  }
  for (unsigned char c = 'A'; c <= 'Z'; ++c) {
    table[c] = &DispatchLexKeywordOrIdentifier;
  }
  // We dispatch all non-ASCII UTF-8 characters to the identifier lexing
  // as whitespace characters should already have been skipped and the
  // only remaining valid Unicode characters would be part of an
  // identifier. That code can either accept or reject.
  for (int i = 0x80; i < 0x100; ++i) {
    table[i] = &DispatchLexKeywordOrIdentifier;
  }

  for (unsigned char c = '0'; c <= '9'; ++c) {
    table[c] = &DispatchLexNumericLiteral;
  }

  table['\''] = &DispatchLexStringLiteral;
  table['"'] = &DispatchLexStringLiteral;
  table['#'] = &DispatchLexHash;

  table[' '] = &DispatchLexHorizontalWhitespace;
  table['\t'] = &DispatchLexHorizontalWhitespace;
  table['\n'] = &DispatchLexVerticalWhitespace;
  table['\r'] = &DispatchLexCR;

  return table;
}

static constexpr DispatchTableT DispatchTable = MakeDispatchTable();

static auto DispatchNext(Lexer& lexer, llvm::StringRef source_text,
                         ssize_t position) -> void {
  if (LLVM_LIKELY(position < static_cast<ssize_t>(source_text.size()))) {
    // The common case is to tail recurse based on the next character. Note
    // that because this is a must-tail return, this cannot fail to tail-call
    // and will not grow the stack. This is in essence a loop with dynamic
    // tail dispatch to the next stage of the loop.
    // NOLINTNEXTLINE(readability-avoid-return-with-void-value): For musttail.
    [[clang::musttail]] return DispatchTable[static_cast<unsigned char>(
        source_text[position])](lexer, source_text, position);
  }

  if (lexer.has_dump_sem_ir_ranges()) {
    // Incomplete ranges will use the next token for their end; we want that to
    // be `FileEnd` in this case, so check before adding `FileEnd`. The argument
    // is just the final character for diagnostic locations.
    // TODO: This offset may not be needed if `file_test` handled diagnostics
    // pointing at `.end()`.
    lexer.EndDumpSemIRRangeIfIncomplete(source_text.end() - 1);
  }

  // When we finish the source text, stop recursing. We also hint this so that
  // the tail-dispatch is optimized as that's essentially the loop back-edge
  // and this is the loop exit.
  lexer.LexFileEnd(source_text, position);
}

// Estimate an upper bound on the number of identifiers we will need to lex.
//
// When analyzing both Carbon and LLVM's C++ code, we have found a roughly
// normal distribution of unique identifiers in the file centered at 0.5 *
// lines, and in the vast majority of cases bounded below 1.0 * lines. For
// example, here is LLVM's distribution computed with `scripts/source_stats.py`
// and rendered in an ASCII-art histogram:
//
//   ## Unique IDs per 10 lines ## (median: 5, p90: 8, p95: 9, p99: 14)
//   1 ids   [  29]  ▍
//   2 ids   [ 282]  ███▊
//   3 ids   [1492]  ███████████████████▉
//   4 ids   [2674]  ███████████████████████████████████▌
//   5 ids   [3011]  ████████████████████████████████████████
//   6 ids   [2267]  ██████████████████████████████▏
//   7 ids   [1549]  ████████████████████▋
//   8 ids   [ 817]  ██████████▉
//   9 ids   [ 301]  ████
//   10 ids  [  98]  █▎
//
//   (Trimmed to only cover 1 - 10 unique IDs per 10 lines of code, 272 files
//    with more unique IDs in the tail.)
//
// We have checked this distribution with several large codebases (currently
// those at Google, happy to cross check with others) that use a similar coding
// style, and it appears to be very consistent. However, we suspect it may be
// dependent on the column width style. Currently, Carbon's toolchain style
// specifies 80-columns, but if we expect the lexer to routinely see files in
// different styles we should re-compute this estimate.
static auto EstimateUpperBoundOnNumIdentifiers(int line_count) -> int {
  return line_count;
}

auto Lexer::Lex() && -> TokenizedBuffer {
  llvm::StringRef source_text = buffer_.source_->text();

  // Enforced by the source buffer, but something we heavily rely on throughout
  // the lexer.
  CARBON_CHECK(source_text.size() < std::numeric_limits<int32_t>::max());

  // First build up our line data structures.
  MakeLines(source_text);

  // Use the line count (and any other info needed from this scan) to make rough
  // estimated reservations of memory in the hot data structures used by the
  // lexer. In practice, scanning for lines is one of the easiest parts of the
  // lexer to accelerate, and we can use its results to minimize the cost of
  // incrementally growing data structures during the hot path of the lexer.
  //
  // Note that for hashtables we want estimates near the upper bound to minimize
  // growth across the vast majority of inputs. They will also typically reserve
  // more memory than we request due to load factor and rounding to power-of-two
  // size. This overshoot is usually fine for hot parts of the lexer where
  // latency is expected to be more important than minimizing memory usage.
  buffer_.value_stores_->identifiers().Reserve(
      EstimateUpperBoundOnNumIdentifiers(buffer_.line_infos_.size()));

  ssize_t position = 0;
  LexFileStart(source_text, position);

  // Manually enter the dispatch loop. This call will tail-recurse through the
  // dispatch table until everything from source_text is consumed.
  DispatchNext(*this, source_text, position);

  Finalize();

  if (consumer_.seen_error()) {
    buffer_.has_errors_ = true;
  }

  return std::move(buffer_);
}

auto Lexer::MakeLines(llvm::StringRef source_text) -> void {
  if (source_text.empty()) {
    // Construct a single line for empty input.
    buffer_.line_infos_.Add(LineInfo(0));
    line_index_ = LineIndex(0);
    return;
  }

  // We currently use `memchr` here which typically is well optimized to use
  // SIMD or other significantly faster than byte-wise scanning. We also use
  // carefully selected variables and the `ssize_t` type for performance and
  // code size of this hot loop.
  //
  // Note that the `memchr` approach here works equally well for LF and CR+LF
  // line endings. Either way, it finds the end of the line and the start of the
  // next line. The lexer below will find the CR byte and peek to see the
  // following LF and jump to the next line correctly. However, this approach
  // does *not* support plain CR or LF+CR line endings. Nor does it support
  // vertical tab or other vertical whitespace.
  //
  // TODO: Eventually, we should extend this to have correct fallback support
  // for handling CR, LF+CR, vertical tab, and other esoteric vertical
  // whitespace as line endings. Notably, including *mixtures* of them. This
  // will likely be somewhat tricky as even detecting their absence without
  // performance overhead and without a custom scanner here rather than memchr
  // is likely to be difficult.
  const char* const text = source_text.data();
  const ssize_t size = source_text.size();
  ssize_t start = 0;
  while (const char* nl = reinterpret_cast<const char*>(
             memchr(&text[start], '\n', size - start))) {
    ssize_t nl_index = nl - text;
    buffer_.line_infos_.Add(LineInfo(start));
    start = nl_index + 1;
  }
  // The last line ends at the end of the file.
  buffer_.line_infos_.Add(LineInfo(start));

  // If the last line wasn't empty, the file ends with an unterminated line.
  // Add an extra blank line so that we never need to handle the special case
  // of being on the last line inside the lexer and needing to not increment
  // to the next line.
  if (start != size) {
    buffer_.line_infos_.Add(LineInfo(size));
  }

  // Now that all the infos are allocated, get a fresh pointer to the first
  // info for use while lexing.
  line_index_ = LineIndex(0);
}

auto Lexer::SkipHorizontalWhitespace(llvm::StringRef source_text,
                                     ssize_t& position) -> void {
  // Handle adjacent whitespace quickly. This comes up frequently for example
  // due to indentation. We don't expect *huge* runs, so just use a scalar
  // loop. While still scalar, this avoids repeated table dispatch and marking
  // whitespace.
  while (position < static_cast<ssize_t>(source_text.size()) &&
         (source_text[position] == ' ' || source_text[position] == '\t')) {
    ++position;
  }
}

auto Lexer::AdvanceToLine(llvm::StringRef source_text, ssize_t& position,
                          LineIndex to_line_index) -> void {
  CARBON_DCHECK(to_line_index >= line_index_);
  line_index_ = to_line_index;
  auto& line_info = current_line_info();
  ssize_t line_start = line_info.start;
  position = line_start;
  SkipHorizontalWhitespace(source_text, position);
  line_info.indent = position - line_start;
}

auto Lexer::LexHorizontalWhitespace(llvm::StringRef source_text,
                                    ssize_t& position) -> void {
  CARBON_DCHECK(source_text[position] == ' ' || source_text[position] == '\t');
  NoteWhitespace();
  // Skip runs using an optimized code path.
  SkipHorizontalWhitespace(source_text, position);
}

auto Lexer::LexVerticalWhitespace(llvm::StringRef source_text,
                                  ssize_t& position) -> void {
  NoteWhitespace();
  AdvanceToLine(source_text, position, next_line());
}

auto Lexer::LexCR(llvm::StringRef source_text, ssize_t& position) -> void {
  if (LLVM_LIKELY((position + 1) < static_cast<ssize_t>(source_text.size())) &&
      LLVM_LIKELY(source_text[position + 1] == '\n')) {
    // Skip to the vertical whitespace path, it will skip over both CR and LF.
    LexVerticalWhitespace(source_text, position);
    return;
  }

  CARBON_DIAGNOSTIC(UnsupportedLfCrLineEnding, Error,
                    "the LF+CR line ending is not supported, only LF and CR+LF "
                    "are supported");
  CARBON_DIAGNOSTIC(UnsupportedCrLineEnding, Error,
                    "a raw CR line ending is not supported, only LF and CR+LF "
                    "are supported");
  bool is_lfcr = position > 0 && source_text[position - 1] == '\n';
  // TODO: This diagnostic has an unfortunate snippet -- we should tweak the
  // snippet rendering to gracefully handle CRs.
  emitter_.Emit(source_text.begin() + position,
                is_lfcr ? UnsupportedLfCrLineEnding : UnsupportedCrLineEnding);

  // Recover by treating the CR as a horizontal whitespace. This should make our
  // whitespace rules largely work and parse cleanly without disrupting the line
  // tracking data structures that were pre-built.
  NoteWhitespace();
  ++position;
}

auto Lexer::LexCommentOrSlash(llvm::StringRef source_text, ssize_t& position)
    -> void {
  CARBON_DCHECK(source_text[position] == '/');

  // Both comments and slash symbols start with a `/`. We disambiguate with a
  // max-munch rule -- if the next character is another `/` then we lex it as
  // a comment start. If it isn't, then we lex as a slash. We also optimize
  // for the comment case as we expect that to be much more important for
  // overall lexer performance.
  if (LLVM_LIKELY(position + 1 < static_cast<ssize_t>(source_text.size()) &&
                  source_text[position + 1] == '/')) {
    LexComment(source_text, position);
    return;
  }

  // This code path should produce a token, make sure that happens.
  LexResult result = LexSymbolToken(source_text, position);
  CARBON_CHECK(result, "Failed to form a token!");
}

auto Lexer::BeginDumpSemIRRange(const char* diag_loc) -> void {
  EndDumpSemIRRangeIfIncomplete(diag_loc);

  // The begin here will be the next token, which may be dump-sem-ir-begin. The
  // end will be assigned by either AddDumpSemIREnd or, if invalid,
  // EndDumpSemIRRangeIfIncomplete.
  buffer_.dump_sem_ir_ranges_.push_back(
      {.begin = TokenIndex(buffer_.size()), .end = TokenIndex::None});
}

auto Lexer::EndDumpSemIRRange(const char* diag_loc) -> void {
  if (buffer_.dump_sem_ir_ranges_.empty() ||
      buffer_.dump_sem_ir_ranges_.back().end != TokenIndex::None) {
    CARBON_DIAGNOSTIC(
        DumpSemIRRangeMissingBegin, Error,
        "missing `//@dump-sem-ir-begin` to match `//@dump-sem-ir-end`");
    emitter_.Emit(diag_loc, DumpSemIRRangeMissingBegin);
    return;
  }

  buffer_.dump_sem_ir_ranges_.back().end = TokenIndex(buffer_.size() - 1);
}

auto Lexer::EndDumpSemIRRangeIfIncomplete(const char* diag_loc) -> void {
  if (buffer_.dump_sem_ir_ranges_.empty() ||
      buffer_.dump_sem_ir_ranges_.back().end != TokenIndex::None) {
    return;
  }

  // The location here won't be closely associated with the start location.
  // However, this is a developer feature and not worth complexity to diagnose
  // better.
  CARBON_DIAGNOSTIC(
      DumpSemIRRangeMissingEnd, Error,
      "missing `//@dump-sem-ir-end` to match `//@dump-sem-ir-begin`");
  emitter_.Emit(diag_loc, DumpSemIRRangeMissingEnd);

  EndDumpSemIRRange(diag_loc);
}

auto Lexer::LexComment(llvm::StringRef source_text, ssize_t& position) -> void {
  CARBON_DCHECK(source_text.substr(position).starts_with("//"));
  int32_t comment_start = position;
  const auto line_info = current_line_info();

  // A comment is _trailing_ when it follows other content on its line, rather
  // than being the only non-whitespace on the line. Both kinds of comment are
  // lexed identically -- they run from `//` to the end of the line -- but we
  // record the distinction so that tooling can tell a comment annotating the
  // code on its line apart from one introducing the code below it.
  //
  // `line_info.indent` is the width of the line's leading whitespace, so
  // `line_info.start + line_info.indent` is the line's first non-whitespace
  // byte.
  const bool is_trailing = position != line_info.start + line_info.indent;

  // Check whether the `//` introducer is followed by something valid.
  bool is_valid_after_slashes = true;
  if (position + 2 < static_cast<ssize_t>(source_text.size()) &&
      LLVM_UNLIKELY(!IsSpace(source_text[position + 2]))) {
    llvm::StringRef comment_text = source_text.substr(position);
    // The `//@...` directives are tooling markers that are only meaningful as
    // full-line comments, so we only recognize them when not trailing.
    if (!is_trailing) {
      // A directive is also recorded as a comment: the tokens and comments
      // together reconstruct the source, so tooling such as the formatter
      // would otherwise silently drop the directive line.
      auto add_directive_comment_line = [&] {
        buffer_.AddComment(line_info.indent, comment_start,
                           buffer_.line_infos_.Get(next_line()).start,
                           /*is_trailing=*/false);
        AdvanceToLine(source_text, position, next_line());
      };
      if (comment_text.starts_with("//@include-in-dumps\n")) {
        buffer_.has_include_in_dumps_ = true;
        add_directive_comment_line();
        return;
      }
      if (comment_text.starts_with("//@dump-sem-ir-begin\n")) {
        BeginDumpSemIRRange(comment_text.begin());
        add_directive_comment_line();
        return;
      }
      if (comment_text.starts_with("//@dump-sem-ir-end\n")) {
        EndDumpSemIRRange(comment_text.begin());
        add_directive_comment_line();
        return;
      }
    }
    CARBON_DIAGNOSTIC(NoWhitespaceAfterCommentIntroducer, Error,
                      "whitespace is required after '//'");
    emitter_.Emit(comment_text.begin() + 2, NoWhitespaceAfterCommentIntroducer);

    // We use this to tweak the lexing of blocks below.
    is_valid_after_slashes = false;
  }

  // Skip over this line.
  LineIndex line_index = next_line();
  position = buffer_.line_infos_.Get(line_index).start;

  // A trailing comment runs to the end of its line. Unlike a full-line comment,
  // it can never be part of a block of identical comment lines, so we skip the
  // block-skipping optimization below and simply advance past this one line. We
  // also don't optimize for the case of a trailing comment as we expect them to
  // be relatively rare compared to other comment structures.
  if (LLVM_UNLIKELY(is_trailing)) {
    buffer_.AddComment(line_info.indent, comment_start, position,
                       /*is_trailing=*/true);
    // Unlike a full-line comment, a trailing comment can directly follow a
    // token that cleared the leading-whitespace flag (`x;// y`), so restore it
    // here for the next line's first token.
    NoteWhitespace();
    AdvanceToLine(source_text, position, line_index);
    return;
  }

  // A very common pattern is a long block of comment lines all with the same
  // indent and comment start. We skip these comment blocks in bulk for speed,
  // and with SIMD support short indents can be scanned extremely quickly; we
  // expect these to be the dominant cases.
  //
  // An invalid comment start was already diagnosed above, so its block is
  // instead skipped line by line below: a run of invalid comment lines lumps
  // into one block regardless of which invalid byte follows each `//`, keeping
  // the diagnostic noise to one per run, while a line whose introducer is
  // valid (whitespace, or a `//@...` directive) ends the run and is lexed on
  // its own.
  //
  // TODO: We should extend this to 32-byte SIMD on platforms with support.
  constexpr int MaxIndent = 13;
  const int indent = line_info.indent;
  const ssize_t first_line_start = line_info.start;
  ssize_t prefix_size = indent + (is_valid_after_slashes ? 3 : 2);
  auto skip_to_next_line = [this, indent, &line_index, &position] {
    // We're guaranteed to have a line here even on a comment on the last line
    // as we ensure there is an empty line structure at the end of every file.
    ++line_index.index;
    auto& next_line_info = buffer_.line_infos_.Get(line_index);
    next_line_info.indent = indent;
    position = next_line_info.start;
  };
  if (CARBON_USE_SIMD && is_valid_after_slashes &&
      position + 16 < static_cast<ssize_t>(source_text.size()) &&
      indent <= MaxIndent) {
    // Load a mask based on the amount of text we want to compare.
    auto mask = PrefixMasks[prefix_size];
#if __ARM_NEON
    // Load and mask the prefix of the current line.
    auto prefix = vld1q_u8(reinterpret_cast<const uint8_t*>(source_text.data() +
                                                            first_line_start));
    prefix = vandq_u8(mask, prefix);
    do {
      // Load and mask the next line to consider's prefix.
      auto next_prefix = vld1q_u8(
          reinterpret_cast<const uint8_t*>(source_text.data() + position));
      next_prefix = vandq_u8(mask, next_prefix);
      // Compare the two prefixes and if any lanes differ, break.
      auto compare = vceqq_u8(prefix, next_prefix);
      if (vminvq_u8(compare) == 0) {
        break;
      }

      skip_to_next_line();
    } while (position + 16 < static_cast<ssize_t>(source_text.size()));
#elif __x86_64__
    // Use the current line's prefix as the exemplar to compare against.
    // We don't mask here as we will mask when doing the comparison.
    auto prefix = _mm_loadu_si128(reinterpret_cast<const __m128i*>(
        source_text.data() + first_line_start));
    do {
      // Load the next line to consider's prefix.
      auto next_prefix = _mm_loadu_si128(
          reinterpret_cast<const __m128i*>(source_text.data() + position));
      // Compute the difference between the next line and our exemplar. Again,
      // we don't mask the difference because the comparison below will be
      // masked.
      auto prefix_diff = _mm_xor_si128(prefix, next_prefix);
      // If we have any differences (non-zero bits) within the mask, we can't
      // skip the next line too.
      if (!_mm_test_all_zeros(mask, prefix_diff)) {
        break;
      }

      skip_to_next_line();
    } while (position + 16 < static_cast<ssize_t>(source_text.size()));
#else
#error "Unsupported SIMD architecture!"
#endif
  } else {
    auto continues_block = [&](ssize_t position) -> bool {
      // Make sure the source text extends far enough for us to continue the
      // block.
      if (position + prefix_size > static_cast<ssize_t>(source_text.size())) {
        return false;
      }

      // Check that the prefix matches. Otherwise, the block is done.
      if (memcmp(source_text.data() + first_line_start,
                 source_text.data() + position, prefix_size) != 0) {
        return false;
      }

      // For something valid after `//`, we're done as we've ensured it was the
      // _same_ valid suffix in the `memcmp`.
      if (LLVM_LIKELY(is_valid_after_slashes)) {
        return true;
      }

      // Past here, the block is a run of invalid comment lines and the
      // matched prefix only covers this line's `//`, so examine what follows
      // those slashes: only another invalid comment line continues the block,
      // while a valid comment or directive ends it and is lexed on its own.

      // A `//` that ends the source is a valid comment, so it doesn't
      // continue the block.
      if (position + prefix_size == static_cast<ssize_t>(source_text.size())) {
        return false;
      }

      char after_slashes = source_text[position + prefix_size];

      // Whitespace after the `//` makes this line a valid comment, so it
      // doesn't continue the block.
      if (IsSpace(after_slashes)) {
        return false;
      }

      // An `@` makes this line a `//@...` directive that must be lexed on its
      // own to recognize its side effects, so it doesn't continue the block.
      if (after_slashes == '@') {
        return false;
      }

      // Anything else is another invalid comment line continuing the block.
      return true;
    };

    // Skip lines that are combined into a comment block.
    while (continues_block(position)) {
      skip_to_next_line();
    }
  }

  buffer_.AddComment(indent, comment_start, position, /*is_trailing=*/false);
  AdvanceToLine(source_text, position, line_index);
}

auto Lexer::CanFormRealLiteral() -> bool {
  // When a numeric literal immediately follows a `.` or `->` token, with no
  // intervening whitespace, a real literal is never formed.
  if (has_leading_space_) {
    return true;
  }
  auto kind = buffer_.GetKind(buffer_.tokens().end()[-1]);
  return kind != TokenKind::Period && kind != TokenKind::MinusGreater;
}

auto Lexer::LexNumericLiteral(llvm::StringRef source_text, ssize_t& position)
    -> LexResult {
  std::optional<NumericLiteral> literal =
      NumericLiteral::Lex(source_text.substr(position), CanFormRealLiteral());
  if (!literal) {
    return LexError(source_text, position);
  }

  // Capture the position before we step past the token.
  int32_t byte_offset = position;
  int token_size = literal->text().size();
  position += token_size;

  CARBON_KIND_SWITCH(literal->ComputeValue(emitter_)) {
    case CARBON_KIND(NumericLiteral::IntValue&& value): {
      return LexTokenWithPayload(TokenKind::IntLiteral,
                                 buffer_.value_stores_->ints()
                                     .AddUnsigned(std::move(value.value))
                                     .AsTokenPayload(),
                                 byte_offset);
    }
    case CARBON_KIND(NumericLiteral::RealValue&& value): {
      auto real_id = buffer_.value_stores_->reals().Add(
          Real{.mantissa = value.mantissa,
               .exponent = value.exponent,
               .is_decimal = (value.radix == NumericLiteral::Radix::Decimal)});
      return LexTokenWithPayload(TokenKind::RealLiteral, real_id.index,
                                 byte_offset);
    }
    case CARBON_KIND(NumericLiteral::UnrecoverableError _): {
      return LexTokenWithPayload(TokenKind::Error, token_size, byte_offset);
    }
  }
}

static auto DiagnoseUnterminatedString(
    Diagnostics::Emitter<const char*>& emitter, const StringLiteral& literal,
    bool is_char) -> void {
  CARBON_DIAGNOSTIC(UnterminatedString, Error,
                    "{0:character|string} literal is missing a terminator",
                    Diagnostics::BoolAsSelect);
  emitter.Emit(literal.text().begin(), UnterminatedString, is_char);
}

auto Lexer::LexStringLiteral(llvm::StringRef source_text, ssize_t& position)
    -> LexResult {
  std::optional<StringLiteral> literal =
      StringLiteral::Lex(source_text.substr(position));
  if (!literal) {
    return LexError(source_text, position);
  }

  // Capture the position before we step past the token.
  int32_t byte_offset = position;
  position += literal->text().size();

  // Helper for error paths.
  auto lex_as_error = [&]() {
    return LexTokenWithPayload(TokenKind::Error, literal->text().size(),
                               byte_offset);
  };

  if (literal->kind() == StringLiteral::Kind::Char) {
    if (!literal->is_terminated()) {
      DiagnoseUnterminatedString(emitter_, *literal, /*is_char=*/true);
      return lex_as_error();
    }
    if (auto value = literal->ComputeCharLiteralValue(emitter_)) {
      return LexTokenWithPayload(TokenKind::CharLiteral, value->value,
                                 byte_offset);
    }
    return lex_as_error();
  }

  if (literal->has_invalid_introducer()) {
    // The literal covers only the malformed introducer line, so it spans no
    // lines and needs no line updates.
    CARBON_DIAGNOSTIC(MultiLineStringInvalidIntroducer, Error,
                      "invalid multi-line string literal introducer; a file "
                      "type indicator may not contain `'`, `#`, or `\"`, and "
                      "the content must begin on a new line");
    emitter_.Emit(literal->text().begin(), MultiLineStringInvalidIntroducer);
    return lex_as_error();
  }

  // Update line and column information.
  if (literal->kind() != StringLiteral::Kind::SingleLine) {
    // A block string literal's content is indented to match its closing
    // delimiter: leading whitespace up to the delimiter's column is
    // indentation, and anything past it is part of the content. Each line the
    // literal spans is given the closing delimiter's column as its indentation.
    // The closing line's indentation must be correct because tokens and
    // comments can follow the closing delimiter and rely on it, for example to
    // detect a trailing comment. Multi-line literals are rare, so this cold
    // path need not be fast.
    LineIndex first_spanned_line(line_index_.index + 1);
    while (next_line_info().start < position) {
      ++line_index_.index;
    }
    // The closing delimiter is the first non-whitespace on the closing line, so
    // its column is that line's leading-whitespace width.
    LineInfo& closing_line_info = current_line_info();
    ssize_t indent_end = closing_line_info.start;
    SkipHorizontalWhitespace(source_text, indent_end);
    int32_t indent = indent_end - closing_line_info.start;
    for (int32_t i = first_spanned_line.index; i <= line_index_.index; ++i) {
      buffer_.line_infos_.Get(LineIndex(i)).indent = indent;
    }
  }

  if (!literal->is_terminated()) {
    DiagnoseUnterminatedString(emitter_, *literal, /*is_char=*/false);
    return lex_as_error();
  }
  auto string_id = buffer_.value_stores_->string_literal_values().Add(
      literal->ComputeStringValue(buffer_.allocator_, emitter_));
  return LexTokenWithPayload(TokenKind::StringLiteral, string_id.index,
                             byte_offset);
}

auto Lexer::LexOneCharSymbolToken(llvm::StringRef source_text, TokenKind kind,
                                  ssize_t& position) -> TokenIndex {
  // Verify in a debug build that the incoming token kind is correct.
  CARBON_DCHECK(kind != TokenKind::Error);
  CARBON_DCHECK(kind.fixed_spelling().size() == 1);
  CARBON_DCHECK(source_text[position] == kind.fixed_spelling().front(),
                "Source text starts with '{0}' instead of the spelling '{1}' "
                "of the incoming token kind '{2}'",
                source_text[position], kind.fixed_spelling(), kind);

  TokenIndex token = LexToken(kind, position);
  ++position;
  return token;
}

auto Lexer::LexOpeningSymbolToken(llvm::StringRef source_text, TokenKind kind,
                                  ssize_t& position) -> LexResult {
  CARBON_DCHECK(kind.is_opening_symbol());
  CARBON_DCHECK(kind.fixed_spelling().size() == 1);
  CARBON_DCHECK(source_text[position] == kind.fixed_spelling().front(),
                "Source text starts with '{0}' instead of the spelling '{1}' "
                "of the incoming token kind '{2}'",
                source_text[position], kind.fixed_spelling(), kind);

  int32_t byte_offset = position;
  ++position;

  // Lex the opening symbol with a zero closing index. We'll add a payload later
  // when we match a closing symbol or in recovery.
  TokenIndex token = LexToken(kind, byte_offset);
  open_groups_.push_back(token);
  return token;
}

auto Lexer::LexClosingSymbolToken(llvm::StringRef source_text, TokenKind kind,
                                  ssize_t& position) -> LexResult {
  CARBON_DCHECK(kind.is_closing_symbol());
  CARBON_DCHECK(kind.fixed_spelling().size() == 1);
  CARBON_DCHECK(source_text[position] == kind.fixed_spelling().front(),
                "Source text starts with '{0}' instead of the spelling '{1}' "
                "of the incoming token kind '{2}'",
                source_text[position], kind.fixed_spelling(), kind);

  int32_t byte_offset = position;
  ++position;

  // If there's not a matching opening symbol, just track that we had an error.
  // We will diagnose and recover when we reach the end of the file. See
  // `DiagnoseAndFixMismatchedBrackets` for details.
  if (LLVM_UNLIKELY(open_groups_.empty())) {
    has_mismatched_brackets_ = true;
    // Lex without a matching index payload -- we'll add one during recovery.
    return LexToken(kind, byte_offset);
  }

  TokenIndex opening_token = open_groups_.pop_back_val();
  TokenIndex token =
      LexTokenWithPayload(kind, opening_token.index, byte_offset);

  auto& opening_token_info = buffer_.token_infos_.Get(opening_token);
  if (LLVM_UNLIKELY(opening_token_info.kind() != kind.opening_symbol())) {
    has_mismatched_brackets_ = true;
    buffer_.token_infos_.Get(token).set_opening_token_index(TokenIndex::None);
    return token;
  }

  opening_token_info.set_closing_token_index(token);
  return token;
}

auto Lexer::LexSymbolToken(llvm::StringRef source_text, ssize_t& position)
    -> LexResult {
  // One character symbols and grouping symbols are handled with dedicated
  // dispatch. We only lex the multi-character tokens here.
  TokenKind kind = llvm::StringSwitch<TokenKind>(source_text.substr(position))
#define CARBON_SYMBOL_TOKEN(Name, Spelling) \
  .StartsWith(Spelling, TokenKind::Name)
#define CARBON_ONE_CHAR_SYMBOL_TOKEN(TokenName, Spelling)
#define CARBON_OPENING_GROUP_SYMBOL_TOKEN(TokenName, Spelling, ClosingName)
#define CARBON_CLOSING_GROUP_SYMBOL_TOKEN(TokenName, Spelling, OpeningName)
#include "toolchain/lex/token_kind.def"
                       .Default(TokenKind::Error);
  if (kind == TokenKind::Error) {
    return LexError(source_text, position);
  }

  TokenIndex token = LexToken(kind, position);
  position += kind.fixed_spelling().size();
  return token;
}

auto Lexer::LexWordAsTypeLiteralToken(llvm::StringRef word, int32_t byte_offset)
    -> LexResult {
  if (word.size() < 2) {
    // Too short to form one of these tokens.
    return LexResult::NoMatch();
  }

  TokenKind kind;
  switch (word.front()) {
    case 'i':
      kind = TokenKind::IntTypeLiteral;
      break;
    case 'u':
      kind = TokenKind::UnsignedIntTypeLiteral;
      break;
    case 'f':
      kind = TokenKind::FloatTypeLiteral;
      break;
    default:
      return LexResult::NoMatch();
  };

  // No leading zeros allowed.
  if ('1' > word[1] || word[1] > '9') {
    return LexResult::NoMatch();
  }

  llvm::StringRef suffix = word.substr(1);

  // Type bit-widths can't usefully be large integers so we restrict to small
  // ones that are especially easy to parse into a normal integer variable by
  // restricting the number of digits to round trip.
  int64_t suffix_value;
  constexpr ssize_t DigitLimit =
      std::numeric_limits<decltype(suffix_value)>::digits10;
  if (suffix.size() > DigitLimit) {
    // See if this is not actually a type literal.
    if (!llvm::all_of(suffix, IsDecimalDigit)) {
      return LexResult::NoMatch();
    }

    // Otherwise, diagnose and produce an error token.
    CARBON_DIAGNOSTIC(TooManyTypeBitWidthDigits, Error,
                      "found a type literal with a bit width using {0} digits, "
                      "which is greater than the limit of {1}",
                      size_t, size_t);
    emitter_.Emit(word.begin() + 1, TooManyTypeBitWidthDigits, suffix.size(),
                  DigitLimit);
    return LexTokenWithPayload(TokenKind::Error, word.size(), byte_offset);
  }

  // It's tempting to do something more clever because we know the length ahead
  // of time, but we expect these to be short (1-3 digits) and profiling doesn't
  // show the loop as hot in the short cases.
  suffix_value = suffix[0] - '0';
  for (char c : suffix.drop_front()) {
    if (!IsDecimalDigit(c)) {
      return LexResult::NoMatch();
    }
    suffix_value = suffix_value * 10 + (c - '0');
  }

  // Add the bit width to our integer store and get its index. We treat it as
  // unsigned as that's less expensive and it can't be negative.
  CARBON_CHECK(suffix_value >= 0);
  auto bit_width_payload =
      buffer_.value_stores_->ints().Add(suffix_value).AsTokenPayload();

  return LexTokenWithPayload(kind, bit_width_payload, byte_offset);
}

auto Lexer::LexWordAsDollarIntLiteralToken(llvm::StringRef word,
                                           int32_t byte_offset) -> LexResult {
  if (!word.starts_with('$')) {
    return LexResult::NoMatch();
  }

  if (!has_leading_space_) {
    auto prev_token = buffer_.tokens().end()[-1];
    auto kind = buffer_.GetKind(prev_token);
    if (kind.is_word()) {
      CARBON_DIAGNOSTIC(
          CharacterOnlyAllowedAtStart, Error,
          "`$` is only allowed at the start of a positional parameter");
      emitter_.Emit(word.begin(), CharacterOnlyAllowedAtStart);

      auto& prev_token_info = buffer_.token_infos_.Get(prev_token);
      auto prev_token_text_size = buffer_.GetTokenText(prev_token).size();
      prev_token_info = TokenInfo(TokenKind::Error, has_leading_space_,
                                  prev_token_text_size + word.size(),
                                  prev_token_info.byte_offset());
      return LexResult(TokenIndex(buffer_.token_infos_.size() - 1));
    }
  }

  auto diagnose_invalid_char = [&]() {
    CARBON_DIAGNOSTIC(
        InvalidCharacterInDollarIntLiteral, Error,
        "Positional parameters can only contain digits after `$`");
    emitter_.Emit(word.begin() + 1, InvalidCharacterInDollarIntLiteral);
    return LexTokenWithPayload(TokenKind::Error, word.size(), byte_offset);
  };

  if (word.size() < 2) {
    CARBON_DIAGNOSTIC(DollarIntLiteralMissingNumber, Error,
                      "Expected digits after `$`");
    emitter_.Emit(word.begin() + 1, DollarIntLiteralMissingNumber);
    return LexTokenWithPayload(TokenKind::Error, word.size(), byte_offset);
  }
  if (word[1] == '0' && word.size() > 2) {
    CARBON_DIAGNOSTIC(
        DollarIntLiteralLeadingZero, Error,
        "Leading zeroes are not allowed in positional parameters");
    emitter_.Emit(word.begin() + 1, DollarIntLiteralLeadingZero);
    return LexTokenWithPayload(TokenKind::Error, word.size(), byte_offset);
  }
  if ((word[1] < '0' || word[1] > '9')) {
    return diagnose_invalid_char();
  }

  auto suffix = word.substr(1);
  int64_t suffix_value;
  constexpr ssize_t DigitLimit =
      std::numeric_limits<decltype(suffix_value)>::digits10;
  if (suffix.size() > DigitLimit) {
    // See if this is not actually a dollar int literal.
    if (!llvm::all_of(suffix, IsDecimalDigit)) {
      return diagnose_invalid_char();
    }

    // Otherwise, diagnose and produce an error token.
    CARBON_DIAGNOSTIC(TooManyDollarIntDigits, Error,
                      "found a positional parameter using {0} digits, "
                      "which is greater than the limit of {1}",
                      size_t, size_t);
    emitter_.Emit(word.begin() + 1, TooManyDollarIntDigits, suffix.size(),
                  DigitLimit);
    return LexTokenWithPayload(TokenKind::Error, word.size(), byte_offset);
  }

  suffix_value = suffix[0] - '0';
  for (char c : suffix.drop_front()) {
    if (!IsDecimalDigit(c)) {
      return diagnose_invalid_char();
    }
    suffix_value = suffix_value * 10 + (c - '0');
  }
  CARBON_CHECK(suffix_value >= 0);
  return LexTokenWithPayload(
      TokenKind::DollarIntLiteral,
      buffer_.value_stores_->ints().Add(suffix_value).AsTokenPayload(),
      byte_offset);
}

auto Lexer::LexKeywordOrIdentifier(llvm::StringRef source_text,
                                   ssize_t& position) -> LexResult {
  if (static_cast<unsigned char>(source_text[position]) > 0x7F) {
    // TODO: Need to add support for Unicode lexing.
    return LexError(source_text, position);
  }
  CARBON_CHECK(
      IsIdStartByteTable[static_cast<unsigned char>(source_text[position])]);

  // Capture the position before we step past the token.
  int32_t byte_offset = position;

  // Take the valid characters off the front of the source buffer.
  llvm::StringRef identifier_text =
      ScanForIdentifierPrefix(source_text.substr(position));
  CARBON_CHECK(!identifier_text.empty(), "Must have at least one character!");
  position += identifier_text.size();

  // Check if the text is a type literal, and if so form such a literal.
  if (LexResult result =
          LexWordAsTypeLiteralToken(identifier_text, byte_offset)) {
    return result;
  }
  if (LexResult result =
          LexWordAsDollarIntLiteralToken(identifier_text, byte_offset)) {
    return result;
  }

  // Check if the text matches a keyword token, and if so use that.
  TokenKind kind = llvm::StringSwitch<TokenKind>(identifier_text)
#define CARBON_KEYWORD_TOKEN(Name, Spelling) .Case(Spelling, TokenKind::Name)
#include "toolchain/lex/token_kind.def"
                       .Default(TokenKind::Error);
  if (kind != TokenKind::Error) {
    return LexToken(kind, byte_offset);
  }

  // Otherwise we have a generic identifier.
  return LexTokenWithPayload(
      TokenKind::Identifier,
      buffer_.value_stores_->identifiers().Add(identifier_text).index,
      byte_offset);
}

auto Lexer::LexHash(llvm::StringRef source_text, ssize_t& position)
    -> LexResult {
  // For `r#`, we already lexed an `r` identifier token. Detect that case and
  // replace that token with a raw identifier. We do this to keep identifier
  // lexing as fast as possible.

  // Look for the `r` token. Note that this is always in bounds because we
  // create a start of file token.
  auto& prev_token_info =
      buffer_.token_infos_.Get(TokenIndex(buffer_.token_infos_.size() - 1));

  // If the previous token isn't the identifier `r`, or the character after `#`
  // isn't the start of an identifier, this is not a raw identifier.
  if (prev_token_info.kind() != TokenKind::Identifier ||
      source_text[position - 1] != 'r' ||
      position + 1 == static_cast<ssize_t>(source_text.size()) ||
      !IsIdStartByteTable[static_cast<unsigned char>(
          source_text[position + 1])] ||
      prev_token_info.byte_offset() != static_cast<int32_t>(position) - 1) {
    [[clang::musttail]] return LexStringLiteral(source_text, position);
  }
  CARBON_DCHECK(buffer_.value_stores_->identifiers().Get(
                    prev_token_info.ident_id()) == "r");

  // Take the valid characters off the front of the source buffer.
  llvm::StringRef identifier_text =
      ScanForIdentifierPrefix(source_text.substr(position + 1));
  CARBON_CHECK(!identifier_text.empty(), "Must have at least one character!");
  position += 1 + identifier_text.size();

  // Replace the `r` identifier's value with the raw identifier.
  // TODO: This token doesn't carry any indicator that it's raw, so
  // diagnostics are unclear.
  prev_token_info.set_ident_id(
      buffer_.value_stores_->identifiers().Add(identifier_text));
  return LexResult(TokenIndex(buffer_.token_infos_.size() - 1));
}

auto Lexer::LexError(llvm::StringRef source_text, ssize_t& position)
    -> LexResult {
  llvm::StringRef error_text =
      source_text.substr(position).take_while([](char c) {
        if (IsAlnum(c)) {
          return false;
        }
        switch (c) {
          case '_':
          case '\t':
          case '\n':
            return false;
          default:
            break;
        }
        return llvm::StringSwitch<bool>(llvm::StringRef(&c, 1))
#define CARBON_SYMBOL_TOKEN(Name, Spelling) .StartsWith(Spelling, false)
#include "toolchain/lex/token_kind.def"
            .Default(true);
      });
  if (error_text.empty()) {
    // TODO: Reimplement this to use the lexer properly. In the meantime,
    // guarantee that we eat at least one byte.
    error_text = source_text.substr(position, 1);
  }

  auto token =
      LexTokenWithPayload(TokenKind::Error, error_text.size(), position);
  CARBON_DIAGNOSTIC(UnrecognizedCharacters, Error,
                    "encountered unrecognized characters while parsing");
  emitter_.Emit(error_text.begin(), UnrecognizedCharacters);

  position += error_text.size();
  return token;
}

auto Lexer::LexFileStart(llvm::StringRef source_text, ssize_t& position)
    -> void {
  CARBON_CHECK(position == 0);

  // Before lexing any source text, add the start-of-file token so that code
  // can assume a non-empty token buffer for the rest of lexing.
  LexToken(TokenKind::FileStart, 0);

  // The file start also represents whitespace.
  NoteWhitespace();

  // Also skip any horizontal whitespace and record the indentation of the
  // first line.
  CARBON_CHECK(current_line_info().start == 0);
  AdvanceToLine(source_text, position, /*to_line_index=*/LineIndex(0));
}

auto Lexer::LexFileEnd(llvm::StringRef source_text, ssize_t position) -> void {
  CARBON_CHECK(position == static_cast<ssize_t>(source_text.size()));
  // Check if the last line is empty and not the first line (and only). If so,
  // re-pin the last line to be the prior one so that diagnostics and editors
  // can treat newlines as terminators even though we internally handle them
  // as separators in case of a missing newline on the last line. We do this
  // here instead of detecting this when we see the newline to avoid more
  // conditions along that fast path.
  if (position == current_line_info().start && line_index_.index != 0) {
    --line_index_.index;
    --position;
  }

  // The end-of-file token is always considered to be whitespace.
  NoteWhitespace();

  LexToken(TokenKind::FileEnd, position);
}

auto Lexer::Finalize() -> void {
  // If we had any mismatched brackets, issue diagnostics and fix them.
  if (has_mismatched_brackets_ || !open_groups_.empty()) {
    DiagnoseAndFixMismatchedBrackets();
  }

  // Reject source files with so many tokens that we may have exceeded the
  // number of bits in `token_payload_`.
  //
  // Note that we rely on this check also catching the case where there are too
  // many identifiers to fit an `IdentifierId` into a `token_payload_`, and
  // likewise for `IntId` and so on. If we start adding any of those IDs prior
  // to lexing, we may need to also limit the number of those IDs here.
  if (buffer_.token_infos_.size() > TokenIndex::Max) {
    CARBON_DIAGNOSTIC(TooManyTokens, Error,
                      "too many tokens in source file; try splitting into "
                      "multiple source files");
    // Subtract one to leave room for the `FileEnd` token.
    token_emitter_.Emit(TokenIndex(TokenIndex::Max - 1), TooManyTokens);
    // TODO: Convert tokens after the token limit to error tokens to avoid
    // misinterpretation by consumers of the tokenized buffer.
  }
}

// A list of pending insertions to make into a tokenized buffer for error
// recovery. These are buffered so that we can perform them in linear time.
class Lexer::ErrorRecoveryBuffer {
 public:
  // `buffer` must not be null.
  explicit ErrorRecoveryBuffer(TokenizedBuffer* buffer) : buffer_(buffer) {}

  auto empty() const -> bool {
    return insertions_.empty() && !any_error_tokens_;
  }

  // Insert a recovery token of kind `kind` before `insert_before`. Multiple
  // insertions before the same token are applied in reverse of the order they
  // were requested (LIFO). Returns an id that `GetInsertedTokenIndex` maps to
  // the inserted token, once `Apply` has run.
  auto InsertBefore(TokenIndex insert_before, TokenKind kind) -> int {
    return AddInsertion(insert_before, kind, /*is_after=*/false);
  }

  // Insert a recovery token of kind `kind` after `insert_after`. Multiple
  // insertions after the same token are applied in the order they were
  // requested (FIFO). Returns an id as `InsertBefore` does.
  auto InsertAfter(TokenIndex insert_after, TokenKind kind) -> int {
    return AddInsertion(insert_after, kind, /*is_after=*/true);
  }

  // Replace the given token with an error token. We do this immediately,
  // because we don't benefit from buffering it.
  auto ReplaceWithError(TokenIndex token) -> void {
    auto& token_info = buffer_->token_infos_.Get(token);
    int error_length = buffer_->GetTokenText(token).size();
    token_info = token_info.AsError(error_length);
    any_error_tokens_ = true;
  }

  // Merge the recovery tokens into the token list of the tokenized buffer.
  auto Apply() -> void {
    llvm::sort(insertions_);

    ValueStore<TokenIndex, TokenInfo> old_tokens =
        std::exchange(buffer_->token_infos_, {});
    int new_size = old_tokens.size() + insertions_.size();
    buffer_->token_infos_.Reserve(new_size);
    buffer_->recovery_tokens_.resize(new_size);
    inserted_token_index_.assign(insertions_.size(), TokenIndex::None);
    new_token_index_.assign(old_tokens.size(), TokenIndex::None);

    size_t ins_idx = 0;
    for (TokenIndex old_idx(0);
         old_idx.index < static_cast<int>(old_tokens.size()); ++old_idx.index) {
      if (ins_idx == insertions_.size() ||
          insertions_[ins_idx].target() != old_idx) {
        // Nothing is inserted before this token: it keeps its leading
        // whitespace and simply moves across.
        new_token_index_[old_idx.index] =
            buffer_->token_infos_.Add(old_tokens.Get(old_idx));
        continue;
      }

      // At least one insertion goes before this token. `insertions_` is sorted
      // by target, so they are exactly the next run of entries. The first of
      // them takes over the token's leading whitespace, so that the whitespace
      // stays at the start of the run.
      bool orig_leading_space = old_tokens.Get(old_idx).has_leading_space();
      bool is_first = true;
      for (; ins_idx < insertions_.size() &&
             insertions_[ins_idx].target() == old_idx;
           ++ins_idx) {
        TokenInfo info = insertions_[ins_idx].info.WithLeadingSpace(
            is_first ? orig_leading_space : false);
        is_first = false;
        TokenIndex added = buffer_->AddToken(info);
        buffer_->recovery_tokens_.set(added.index);
        inserted_token_index_[insertions_[ins_idx].insertion_order] = added;
      }
      new_token_index_[old_idx.index] = buffer_->token_infos_.Add(
          old_tokens.Get(old_idx).WithLeadingSpace(false));
    }
  }

  // Maps a token index from before `Apply` to the same token's index after it.
  // `Apply` must have run, except that with no insertions to apply this is the
  // identity either way.
  auto GetNewTokenIndex(TokenIndex old_index) const -> TokenIndex {
    if (insertions_.empty()) {
      return old_index;
    }
    CARBON_CHECK(!new_token_index_.empty(),
                 "Token indexes are only renumbered by `Apply`.");
    return new_token_index_[old_index.index];
  }

  // Maps an id returned by `InsertBefore` or `InsertAfter` to the token that
  // insertion added. `Apply` must have run.
  auto GetInsertedTokenIndex(int insertion_id) const -> TokenIndex {
    return inserted_token_index_[insertion_id];
  }

  // Perform bracket matching to fix cross-references between tokens. This must
  // be done after all recovery is performed and all brackets match, because
  // recovery will change token indexes.
  auto FixTokenCrossReferences() -> void {
    llvm::SmallVector<TokenIndex> open_groups;
    for (auto token : buffer_->tokens()) {
      auto kind = buffer_->GetKind(token);
      if (kind.is_opening_symbol()) {
        open_groups.push_back(token);
      } else if (kind.is_closing_symbol()) {
        CARBON_CHECK(!open_groups.empty(), "Failed to balance brackets");
        auto opening_token = open_groups.pop_back_val();

        CARBON_CHECK(kind == buffer_->token_infos_.Get(opening_token)
                                 .kind()
                                 .closing_symbol(),
                     "Failed to balance brackets");
        auto& opening_token_info = buffer_->token_infos_.Get(opening_token);
        auto& closing_token_info = buffer_->token_infos_.Get(token);
        opening_token_info.set_closing_token_index(token);
        closing_token_info.set_opening_token_index(opening_token);
      }
    }
  }

 private:
  struct Insertion {
    TokenIndex anchor;
    bool is_after;
    // Where this insertion came in the sequence of requests, which is both the
    // id handed back to the caller and the tiebreak between insertions that go
    // in the same place.
    int insertion_order;
    TokenInfo info;

    // The token this insertion goes before, indexed in the stream as it was
    // before `Apply`. `Apply` sorts insertions by this.
    auto target() const -> TokenIndex {
      return is_after ? TokenIndex(anchor.index + 1) : anchor;
    }

    // Orders insertions as `Apply` emits them, so that sorting produces the
    // final token order. Insertions are grouped by the token they go before,
    // because `Apply` walks the old tokens in order, and within a group:
    //
    // - An insertion requested as "after the previous token" comes before one
    //   requested as "before this token", so that each stays on the side of
    //   the gap it was anchored to.
    // - A closing bracket comes before an opening one, so that a group ending
    //   in the gap is closed before a new group is opened in it.
    // - Otherwise the request order breaks the tie, in the direction that puts
    //   the earliest request nearest its anchor: first-requested first for
    //   insertions after a token, last-requested first for insertions before
    //   one.
    friend auto operator<(const Insertion& lhs, const Insertion& rhs) -> bool {
      if (lhs.target() != rhs.target()) {
        return lhs.target() < rhs.target();
      }
      if (lhs.is_after != rhs.is_after) {
        return lhs.is_after;
      }
      bool lhs_is_closing = lhs.info.kind().is_closing_symbol();
      bool rhs_is_closing = rhs.info.kind().is_closing_symbol();
      if (lhs_is_closing != rhs_is_closing) {
        return lhs_is_closing;
      }
      return lhs.is_after ? lhs.insertion_order < rhs.insertion_order
                          : lhs.insertion_order > rhs.insertion_order;
    }
  };

  auto AddInsertion(TokenIndex anchor, TokenKind kind, bool is_after) -> int {
    CARBON_CHECK(anchor.index >= 0, "Invalid anchor token index.");
    CARBON_CHECK(anchor.index < static_cast<int>(buffer_->token_infos_.size()),
                 "Cannot insert past the end of file token.");
    if (!is_after) {
      CARBON_CHECK(anchor.index > 0,
                   "Cannot insert before the start of file token.");
    }

    bool insert_leading_space = false;
    int32_t byte_offset = 0;

    if (!is_after) {
      insert_leading_space = buffer_->HasLeadingWhitespace(anchor);
      TokenIndex insert_after_idx(anchor.index - 1);
      const auto& prev_info = buffer_->token_infos_.Get(insert_after_idx);
      byte_offset = prev_info.byte_offset() +
                    buffer_->GetTokenText(insert_after_idx).size();
    } else {
      const auto& anchor_info = buffer_->token_infos_.Get(anchor);
      byte_offset =
          anchor_info.byte_offset() + buffer_->GetTokenText(anchor).size();
      TokenIndex next_tok(anchor.index + 1);
      if (next_tok.index < static_cast<int>(buffer_->token_infos_.size())) {
        insert_leading_space = buffer_->HasLeadingWhitespace(next_tok);
      }
    }

    int insertion_order = static_cast<int>(insertions_.size());
    insertions_.push_back({
        .anchor = anchor,
        .is_after = is_after,
        .insertion_order = insertion_order,
        .info = TokenInfo(kind, insert_leading_space, byte_offset),
    });
    return insertion_order;
  }

  TokenizedBuffer* buffer_;
  llvm::SmallVector<Insertion> insertions_;
  bool any_error_tokens_ = false;

  // Filled in by `Apply`: which token each insertion added, indexed by the id
  // `AddInsertion` handed out.
  llvm::SmallVector<TokenIndex> inserted_token_index_;

  // Filled in by `Apply`: the new index of each token that existed before it,
  // indexed by that token's old index.
  llvm::SmallVector<TokenIndex> new_token_index_;
};

// Returns true if the token kind forms a complete primary expression on its
// own: an identifier, a literal, `self`, a type keyword, and so on.
static auto IsLeafTokenKind(TokenKind kind) -> bool {
  switch (kind) {
    case TokenKind::Identifier:
    case TokenKind::IntLiteral:
    case TokenKind::RealLiteral:
    case TokenKind::StringLiteral:
    case TokenKind::CharLiteral:
    case TokenKind::IntTypeLiteral:
    case TokenKind::UnsignedIntTypeLiteral:
    case TokenKind::FloatTypeLiteral:
    case TokenKind::True:
    case TokenKind::False:
    case TokenKind::SelfValueIdentifier:
    case TokenKind::SelfTypeIdentifier:
    case TokenKind::Underscore:
    case TokenKind::Bool:
    case TokenKind::Type:
    case TokenKind::Auto:
    case TokenKind::Array:
    case TokenKind::Str:
    case TokenKind::Char:
    case TokenKind::Core:
    case TokenKind::Cpp:
      return true;
    default:
      return false;
  }
}

static auto CollectMismatchedBracketTokens(const TokenizedBuffer& buffer)
    -> llvm::SmallVector<MismatchedBracketToken> {
  llvm::SmallVector<MismatchedBracketToken> input_tokens;
  input_tokens.reserve(buffer.size());

  // Where the previous token in the full stream ended (comments are not
  // tokens, so they don't interrupt this).
  int32_t prev_end_byte = -1;
  int32_t prev_line_index = -1;

  for (auto it = buffer.tokens().begin(); it != buffer.tokens().end(); ++it) {
    TokenIndex token = *it;
    auto kind = buffer.GetKind(token);
    int32_t byte_offset = buffer.GetByteOffset(token);
    auto token_line = buffer.GetLine(token);
    bool has_wide_leading_space = prev_end_byte >= 0 &&
                                  token_line.index == prev_line_index &&
                                  byte_offset - prev_end_byte >= 2;
    prev_end_byte =
        byte_offset + static_cast<int32_t>(buffer.GetTokenText(token).size());
    prev_line_index = token_line.index;

    bool is_paren_keyword = false;
    bool is_else_keyword = false;
    BracketTokenKind bracket_kind;
    switch (kind) {
      case TokenKind::OpenParen:
        bracket_kind = BracketTokenKind::OpenParen;
        break;
      case TokenKind::OpenCurlyBrace:
        bracket_kind = BracketTokenKind::OpenCurlyBrace;
        break;
      case TokenKind::OpenSquareBracket:
        bracket_kind = BracketTokenKind::OpenSquareBracket;
        break;
      case TokenKind::CloseParen:
        bracket_kind = BracketTokenKind::CloseParen;
        break;
      case TokenKind::CloseCurlyBrace:
        bracket_kind = BracketTokenKind::CloseCurlyBrace;
        break;
      case TokenKind::CloseSquareBracket:
        bracket_kind = BracketTokenKind::CloseSquareBracket;
        break;
      case TokenKind::Semi:
        bracket_kind = BracketTokenKind::Semi;
        break;
      case TokenKind::Comma:
        bracket_kind = BracketTokenKind::Comma;
        break;
      case TokenKind::Period:
        bracket_kind = BracketTokenKind::Period;
        break;
      case TokenKind::If:
      case TokenKind::While:
      case TokenKind::For:
      case TokenKind::Match:
        bracket_kind = BracketTokenKind::StatementIntroducer;
        is_paren_keyword = true;
        break;
      case TokenKind::Else: {
        bracket_kind = BracketTokenKind::StatementIntroducer;
        // Only a statement `else` (followed by `{` or `if`) normally follows
        // a `}`; a ternary `if..then..else` is followed by an expression.
        auto else_next = std::next(it);
        if (else_next != buffer.tokens().end()) {
          auto next_kind = buffer.GetKind(*else_next);
          is_else_keyword = next_kind == TokenKind::OpenCurlyBrace ||
                            next_kind == TokenKind::If;
        }
        break;
      }
#define CARBON_DECL_INTRODUCER_TOKEN(kind, name) case TokenKind::kind:
#include "toolchain/lex/token_kind.def"
      case TokenKind::Abstract:
      case TokenKind::Case:
      case TokenKind::Continue:
      case TokenKind::Default:
      case TokenKind::Eval:
      case TokenKind::Extend:
      case TokenKind::Final:
      case TokenKind::Friend:
      case TokenKind::Inline:
      case TokenKind::MustEval:
      case TokenKind::Observe:
      case TokenKind::Override:
      case TokenKind::Private:
      case TokenKind::Protected:
      case TokenKind::Return:
      case TokenKind::Returned:
      case TokenKind::Static:
      case TokenKind::Virtual:
        bracket_kind = BracketTokenKind::StatementIntroducer;
        break;
      case TokenKind::Forall:
        bracket_kind = BracketTokenKind::Other;
        is_paren_keyword = true;
        break;
      case TokenKind::Equal:
        bracket_kind = BracketTokenKind::Assignment;
        break;
      case TokenKind::As:
        bracket_kind = BracketTokenKind::As;
        break;
      case TokenKind::MinusGreater:
      case TokenKind::Where:
        bracket_kind = BracketTokenKind::StructuralOp;
        break;
      case TokenKind::Ref:
      case TokenKind::Unused:
      case TokenKind::Template:
      case TokenKind::Const:
        bracket_kind = BracketTokenKind::ModifierKeyword;
        break;
      case TokenKind::EqualEqual:
      case TokenKind::ExclaimEqual:
      case TokenKind::Less:
      case TokenKind::LessEqual:
      case TokenKind::Greater:
      case TokenKind::GreaterEqual:
      case TokenKind::And:
      case TokenKind::Or:
        bracket_kind = BracketTokenKind::ComparisonOp;
        break;
      case TokenKind::FileStart:
        continue;
      case TokenKind::FileEnd:
        bracket_kind = BracketTokenKind::FileEnd;
        break;
      default:
        bracket_kind = IsLeafTokenKind(kind) ? BracketTokenKind::Leaf
                                             : BracketTokenKind::Other;
        break;
    }

    auto line = token_line;
    int32_t line_indent =
        (kind == TokenKind::FileEnd) ? 0 : buffer.GetIndentColumnNumber(line);

    auto next_it = std::next(it);
    bool is_at_end_of_line =
        (next_it == buffer.tokens().end() || buffer.GetLine(*next_it) != line);

    bool is_struct_brace = false;
    if (kind == TokenKind::OpenCurlyBrace) {
      if (next_it != buffer.tokens().end()) {
        auto next_kind = buffer.GetKind(*next_it);
        if (next_kind == TokenKind::Period ||
            next_kind == TokenKind::CloseCurlyBrace) {
          is_struct_brace = true;
        } else if (next_kind == TokenKind::Identifier) {
          auto next2_it = std::next(next_it);
          if (next2_it != buffer.tokens().end() &&
              buffer.GetKind(*next2_it) == TokenKind::Colon) {
            is_struct_brace = true;
          }
        }
      }
    }

    input_tokens.push_back(MismatchedBracketToken{
        .token_index = token,
        .kind = bracket_kind,
        .line = line.index,
        .line_indent = line_indent,
        .is_at_end_of_line = is_at_end_of_line,
        .is_struct_brace = is_struct_brace,
        .is_paren_keyword = is_paren_keyword,
        .is_else_keyword = is_else_keyword,
        .has_leading_space = buffer.HasLeadingWhitespace(token),
        .has_wide_leading_space = has_wide_leading_space,
    });
  }

  return input_tokens;
}

// The source position where `token` starts. Note that this can't go through
// `GetTokenText`, which returns the kind's fixed spelling rather than a pointer
// into the source for most token kinds.
static auto TokenStartPosition(const TokenizedBuffer& buffer, TokenIndex token)
    -> const char* {
  return buffer.source().text().begin() + buffer.GetByteOffset(token);
}

// The source position just past the end of `token`.
static auto TokenEndPosition(const TokenizedBuffer& buffer, TokenIndex token)
    -> const char* {
  return TokenStartPosition(buffer, token) + buffer.GetTokenText(token).size();
}

// Where in the source a bracket that `correction` proposes would be written:
// just past the end of the token it goes after, rather than at the start of the
// token it goes before. For `f(x` on one line and `;` on the next, that points
// the suggestion at the position directly after the `x`, where the `)` belongs,
// instead of down at the `;`. Where the two tokens are separated, the bracket
// is then placed on the side that matches how it would be written.
static auto BracketInsertionPosition(const TokenizedBuffer& buffer,
                                     const BracketCorrection& correction)
    -> const char* {
  TokenIndex insert_after =
      correction.fix_action == BracketFixAction::InsertAfter
          ? correction.fix_token_index
          : TokenIndex(correction.fix_token_index.index - 1);
  // Nothing precedes the first token, so an insertion before it goes at the
  // very start of the file.
  if (!insert_after.has_value()) {
    return TokenStartPosition(buffer, correction.fix_token_index);
  }

  const char* pos = TokenEndPosition(buffer, insert_after);
  const char* end = buffer.source().text().end();

  // Advances over the spaces starting at `from`, up to `limit` of them.
  auto skip_spaces = [end](const char* from, int32_t limit) {
    for (; limit > 0 && from != end && *from == ' '; --limit) {
      ++from;
    }
    return from;
  };
  constexpr int32_t NoLimit = std::numeric_limits<int32_t>::max();

  if (correction.fix_token_kind == TokenKind::CloseCurlyBrace) {
    // A `}` closing a multi-line scope goes on a line of its own, indented to
    // line up with the `{` it closes. Point at that column when the next line
    // is indented far enough to have a position there, and otherwise at the
    // start of it: a diagnostic can't name the virtual space to the right of
    // the end of a line, so this is as close as the source gets.
    if (const char* line_end = skip_spaces(pos, NoLimit);
        line_end != end && *line_end == '\n') {
      auto open_line = buffer.GetLine(correction.diagnostic_token_index);
      pos = skip_spaces(line_end + 1,
                        buffer.GetIndentColumnNumber(open_line) - 1);
    }
  } else if (correction.fix_token_kind.is_opening_symbol()) {
    // An opening bracket binds to what comes after it, so it belongs on the far
    // side of any space separating it from the token it follows.
    pos = skip_spaces(pos, NoLimit);
  }
  return pos;
}

// If brackets didn't pair or nest properly, find a set of places to insert
// brackets to fix the nesting, issue suitable diagnostics, and update the
// token list to describe the fixes.
auto Lexer::DiagnoseAndFixMismatchedBrackets() -> void {
  ErrorRecoveryBuffer fixes(&buffer_);
  auto input_tokens = CollectMismatchedBracketTokens(buffer_);
  auto corrections = FixMismatchedBrackets(input_tokens);

  // For each correction, the insertion it requested, or -1 if it made none.
  // Applying the fixes renumbers the token stream, so this is what lets the
  // corrections be reported against the indexes the caller will see.
  llvm::SmallVector<int> insertion_ids(corrections.size(), -1);

  for (auto [correction_index, correction] : llvm::enumerate(corrections)) {
    CARBON_DIAGNOSTIC(UnmatchedOpening, Error,
                      "opening symbol without a corresponding closing symbol");
    CARBON_DIAGNOSTIC(UnmatchedClosing, Error,
                      "closing symbol without a corresponding opening symbol");
    CARBON_DIAGNOSTIC(PossiblyMissingBracketHere, Note,
                      "possibly missing `{0}` here", Lex::TokenKind);

    // The note names a position between two tokens, which only the source
    // pointer emitter can express.
    auto builder = emitter_.Build(
        TokenStartPosition(buffer_, correction.diagnostic_token_index),
        correction.diagnostic_kind == BracketDiagnosticKind::UnmatchedOpening
            ? UnmatchedOpening
            : UnmatchedClosing);

    if (correction.fix_action == BracketFixAction::ReplaceWithError) {
      fixes.ReplaceWithError(correction.fix_token_index);
    } else if (correction.is_tied) {
      // The cheapest repairs disagree about where this bracket goes, so give up
      // on the bracket rather than suggest one of them.
      fixes.ReplaceWithError(correction.diagnostic_token_index);
    } else {
      builder.Note(BracketInsertionPosition(buffer_, correction),
                   PossiblyMissingBracketHere, correction.fix_token_kind);
      insertion_ids[correction_index] =
          correction.fix_action == BracketFixAction::InsertBefore
              ? fixes.InsertBefore(correction.fix_token_index,
                                   correction.fix_token_kind)
              : fixes.InsertAfter(correction.fix_token_index,
                                  correction.fix_token_kind);
    }

    builder.Emit();
  }

  buffer_.has_errors_ = true;

  if (!fixes.empty()) {
    fixes.Apply();
    fixes.FixTokenCrossReferences();
  }

  if (options_.bracket_corrections) {
    // Report the corrections against the token stream the caller sees, rather
    // than the one recovery started from, which no longer exists. A correction
    // that inserted a bracket names the inserted token itself; one that only
    // proposed a bracket, or replaced one with an error, names the token it
    // still refers to.
    for (auto [correction_index, correction] : llvm::enumerate(corrections)) {
      correction.diagnostic_token_index =
          fixes.GetNewTokenIndex(correction.diagnostic_token_index);
      correction.fix_token_index =
          insertion_ids[correction_index] >= 0
              ? fixes.GetInsertedTokenIndex(insertion_ids[correction_index])
              : fixes.GetNewTokenIndex(correction.fix_token_index);
    }
    *options_.bracket_corrections = std::move(corrections);
  }
}

auto Lex(SharedValueStores& value_stores, SourceBuffer& source,
         LexOptions options) -> TokenizedBuffer {
  auto* consumer =
      options.consumer ? options.consumer : &Diagnostics::ConsoleConsumer();
  auto tokens = Lexer(options, value_stores, source, *consumer).Lex();

  if (options.vlog_stream || options.dump_stream) {
    // Flush diagnostics before printing.
    consumer->Flush();
  }
  CARBON_VLOG_TO(options.vlog_stream, "*** Lex::TokenizedBuffer ***\n{0}",
                 tokens);
  if (options.dump_stream) {
    tokens.Print(*options.dump_stream, options.omit_file_boundary_tokens);
  }
  return tokens;
}

}  // namespace Carbon::Lex