Optimize token kinds and whitespace flags (#8044)

Currently, these lookup the token in its `ValueStore`, and then access
the kind and flag from that. This is near-optimal in the lexer while
_writing_ the tokens, but ends up being considerably slower in parsing
when the overhead of digging the token out of the value store dominates
the actual use case.

This PR does an struct-of-arrays reorganization of the kinds and
whitespace flags, but only after lexing. This lets lexing write a linear
series of info structs, and only afterward re-organize into a
query-optimized representation that is a struct of arrays.

However, to make this a net performance gain, we have to create the
arrays _very_ efficiently. We can avoid hand-written assembly here by
very carefully and precisely writing the inner loop to build these
arrays so that LLVM will vectorize it into high-quality SIMD code for us
across both x86 and Arm. We also need the arrays to have a particular
shape (for example, using a `SmallVector<bool>` rather than a
`BitVector`) in order to preserve this vectorized creation. It burns a
tiny amount of memory, but is key to the whole thing being a net win of
5% to 10% across parsing.

Benchmark results (./scripts/bench_runner.py --runs=12,
-march=x86-64-v3):

- BM_CompileApiFileDenseDecls<Lang::Carbon, Phase::Parse>/16384: CPU
Time: -9.045%
- BM_CompileApiFileDenseDecls<Lang::Carbon, Phase::Parse>/65536: CPU
Time: -6.526%
- BM_CompileImplFile<Lang::Carbon, Phase::Parse>/16384: CPU Time:
-9.275%
- BM_CompileImplFile<Lang::Carbon, Phase::Parse>/65536: CPU Time:
-8.382%

Assisted-by: Antigravity with both Gemini and Claude

toolchain/base/value_store.h

@@ -15,6 +15,7 @@
 
 #include "common/check.h"
 #include "common/ostream.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallVector.h"
@@ -351,6 +352,13 @@ class ValueStore
     return EnumerateRange(*this);
   }
 
+  // Returns the chunks of values in order. Each chunk is a contiguous range of
+  // values, so bulk operations over all values can iterate over contiguous
+  // memory rather than looking up each value by its ID.
+  auto chunks() const [[clang::lifetimebound]] -> auto {
+    return llvm::ArrayRef<Chunk>(chunks_);
+  }
+
   auto GetIdTag() const -> IdTagType { return tag_; }
   auto GetRawIndex(IdT id) const -> int32_t {
     CARBON_DCHECK(id.index >= 0, "{0}", id.index);
@@ -483,7 +491,10 @@ class ValueStore
         -> void
       requires(std::is_copy_constructible_v<ValueT>);
 
+    auto data() const -> const ValueType* { return buf_; }
     auto size() const -> int32_t { return num_; }
+    auto begin() const -> const ValueType* { return buf_; }
+    auto end() const -> const ValueType* { return buf_ + num_; }
 
    private:
     // Verify using an `int32_t` for `num_` is sound.

toolchain/lex/lex.cpp

@@ -1389,7 +1389,9 @@ auto Lexer::CanFormRealLiteral() -> bool {
   if (has_leading_space_) {
     return true;
   }
-  auto kind = buffer_.GetKind(buffer_.tokens().end()[-1]);
+  // The flat kind array isn't built until lexing finishes, so look at the
+  // token info directly.
+  auto kind = buffer_.token_infos_.Get(buffer_.tokens().end()[-1]).kind();
   return kind != TokenKind::Period && kind != TokenKind::MinusGreater;
 }
 
@@ -1683,7 +1685,8 @@ auto Lexer::LexWordAsDollarIntLiteralToken(llvm::StringRef word,
 
   if (!has_leading_space_) {
     auto prev_token = buffer_.tokens().end()[-1];
-    auto kind = buffer_.GetKind(prev_token);
+    // The flat kind array isn't built yet, so look at the token info directly.
+    auto kind = buffer_.token_infos_.Get(prev_token).kind();
     if (kind.is_word()) {
       CARBON_DIAGNOSTIC(
           CharacterOnlyAllowedAtStart, Error,
@@ -1913,6 +1916,11 @@ auto Lexer::LexFileEnd(llvm::StringRef source_text, ssize_t position) -> void {
 }
 
 auto Lexer::Finalize() -> void {
+  // The token sequence is complete, except for any bracket recovery below, so
+  // build the flat token arrays. Bracket recovery reads them, and rebuilds them
+  // if it changes the sequence.
+  buffer_.BuildTokenArrays();
+
   // If we had any mismatched brackets, issue diagnostics and fix them.
   if (has_mismatched_brackets_ || !open_groups_.empty()) {
     DiagnoseAndFixMismatchedBrackets();
@@ -1963,7 +1971,8 @@ class Lexer::ErrorRecoveryBuffer {
   }
 
   // Replace the given token with an error token. We do this immediately,
-  // because we don't benefit from buffering it.
+  // because we don't benefit from buffering it. The flat token arrays are
+  // brought up to date when `Apply` rebuilds them.
   auto ReplaceWithError(TokenIndex token) -> void {
     auto& token_info = buffer_->token_infos_.Get(token);
     int error_length = buffer_->GetTokenText(token).size();
@@ -2014,6 +2023,9 @@ class Lexer::ErrorRecoveryBuffer {
       new_token_index_[old_idx.index] = buffer_->token_infos_.Add(
           old_tokens.Get(old_idx).WithLeadingSpace(false));
     }
+
+    // The token sequence changed, so rebuild the flat token arrays from it.
+    buffer_->BuildTokenArrays();
   }
 
   // Maps a token index from before `Apply` to the same token's index after it.

toolchain/lex/token_info.h

@@ -52,7 +52,7 @@ struct CharLiteralValue {
 class TokenInfo {
  public:
   // The kind for this token.
-  auto kind() const -> TokenKind { return kind_; }
+  auto kind() const -> TokenKind { return TokenKind::Make(kind_); }
 
   // Whether this token is preceded by whitespace. We only store the preceding
   // state, and look at the next token to check for trailing whitespace.
@@ -191,8 +191,8 @@ class TokenInfo {
   // token index. Stores to this field may overflow, but we produce an error
   // in `Lexer::Finalize` if the file has more than `TokenIndex::Max` tokens,
   // so this value never overflows if lexing succeeds.
-  TokenKind kind_;
-  static_assert(sizeof(kind_) == 1, "TokenKind must pack to 8 bits");
+  static_assert(sizeof(TokenKind) == 1, "TokenKind must pack to 8 bits");
+  TokenKind::RawEnumType kind_ : 8;
   bool has_leading_space_ : 1;
   unsigned token_payload_ : PayloadBits;
 

toolchain/lex/tokenized_buffer.cpp

@@ -229,11 +229,56 @@ auto TokenizedBuffer::AddPostLexingRecoveryToken(TokenInfo info) -> TokenIndex {
   }
 
   auto token = token_infos_.Add(info);
+  // Overwrite the first `FileEnd` padding entry and append a new one at the end
+  // to keep `KindDataPadding` trailing entries intact.
+  token_kinds_[token.index] = info.kind();
+  token_kinds_.push_back(TokenKind::FileEnd);
+  token_has_leading_space_.push_back(info.has_leading_space());
   recovery_tokens_.push_back(true);
   ++post_lexing_recovery_tokens_;
   return token;
 }
 
+auto TokenizedBuffer::BuildTokenArrays() -> void {
+  int32_t num_tokens = token_infos_.size();
+  token_kinds_.resize_for_overwrite(num_tokens + KindDataPadding);
+  token_has_leading_space_.resize_for_overwrite(num_tokens);
+
+  TokenKind* kinds = token_kinds_.data();
+  bool* has_leading_space = token_has_leading_space_.data();
+  for (const auto& chunk : token_infos_.chunks()) {
+    // Building these flat arrays is only a performance benefit when we can do
+    // so using especially efficient vectorized processing of the underlying
+    // `TokenInfo` structs that are created during lexing. We rely on Clang to
+    // autovectorize the inner loop below, but have to take some care to ensure
+    // that works.
+    //
+    // Because `TokenInfo`'s bitfield loads as 32-bit words, LLVM's default x86
+    // cost model caps the vector width at `256 / 32 = 8` lanes and emits
+    // narrower shuffles; requesting a width of 16 packs each iteration's
+    // extracted bytes into full 128-bit stores via `vpackusdw`/`vpackuswb`.
+    //
+    // We guard the pragma on `NDEBUG` (rather than `__OPTIMIZE__`, which is
+    // also set under `fastbuild`'s `-Og` where the loop vectorizer is disabled)
+    // and skip it under UBSan (whose per-iteration checks block vectorization),
+    // as Clang otherwise diagnoses `-Wpass-failed` under `-Werror`.
+    //
+    // Last but not least, it is important to iterate over `TokenInfo` by value
+    // (8 bytes) so `kind()` and `has_leading_space()` extract from a single SSA
+    // value. When using a reference, Clang lowers them as two narrower memory
+    // loads that AArch64 vectorizes as duplicate strided `ld2` passes over the
+    // chunk instead of single-pass `ldp` loads.
+#if defined(NDEBUG) && !__has_feature(undefined_behavior_sanitizer)
+#pragma clang loop vectorize_width(16)
+#endif
+    for (TokenInfo info : chunk) {
+      *kinds++ = info.kind();
+      *has_leading_space++ = info.has_leading_space();
+    }
+  }
+  std::fill_n(kinds, KindDataPadding, TokenKind::FileEnd);
+}
+
 auto TokenizedBuffer::GetIndentColumnNumber(LineIndex line) const -> int {
   return line_infos_.Get(line).indent + 1;
 }
@@ -482,6 +527,9 @@ auto TokenizedBuffer::CollectMemUsage(MemUsage& mem_usage,
                                       llvm::StringRef label) const -> void {
   mem_usage.Collect(MemUsage::ConcatLabel(label, "allocator_"), allocator_);
   mem_usage.Collect(MemUsage::ConcatLabel(label, "token_infos_"), token_infos_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "token_kinds_"), token_kinds_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "token_has_leading_space_"),
+                    token_has_leading_space_);
   mem_usage.Collect(MemUsage::ConcatLabel(label, "line_infos_"), line_infos_);
   mem_usage.Collect(MemUsage::ConcatLabel(label, "comments_"), comments_);
 }

toolchain/lex/tokenized_buffer.h

@@ -7,6 +7,7 @@
 
 #include <cstdint>
 
+#include "common/check.h"
 #include "common/ostream.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/SmallVector.h"
@@ -127,7 +128,22 @@ class TokenizedBuffer : public Printable<TokenizedBuffer> {
     LineIndex start_line;
   };
 
+  // The number of `TokenKind::FileEnd` entries guaranteed to follow the last
+  // token in the array returned by `kind_data`.
+  static constexpr int32_t KindDataPadding = 64;
+
   auto GetKind(TokenIndex token) const -> TokenKind;
+
+  // Returns a flat, contiguous array of the kinds of all tokens, indexed by
+  // `TokenIndex`. The first `size()` entries are the kinds of the tokens, and
+  // they are followed by at least `KindDataPadding` entries of
+  // `TokenKind::FileEnd` so that callers can scan past the end of the tokens
+  // without bounds checks, for example with SIMD loads.
+  //
+  // The returned pointer is invalidated by
+  // `AddPostLexingRecoveryTokenAsIdentifier`.
+  auto kind_data() const -> const TokenKind* { return token_kinds_.data(); }
+
   auto GetLine(TokenIndex token) const -> LineIndex;
 
   // Returns the 1-based line number.
@@ -247,10 +263,8 @@ class TokenizedBuffer : public Printable<TokenizedBuffer> {
   // recovery tokens inserted inline by lexing, such as matching brackets, but
   // excludes recovery tokens added after lexing finished.
   auto tokens() const -> llvm::iterator_range<TokenIterator> {
-    return llvm::make_range(
-        TokenIterator(TokenIndex(0)),
-        TokenIterator(
-            TokenIndex(token_infos_.size() - post_lexing_recovery_tokens_)));
+    return llvm::make_range(TokenIterator(TokenIndex(0)),
+                            TokenIterator(TokenIndex(num_lexed_tokens())));
   }
 
   auto size() const -> int { return token_infos_.size(); }
@@ -341,9 +355,22 @@ class TokenizedBuffer : public Printable<TokenizedBuffer> {
 
   auto FindLineIndex(int32_t byte_offset) const -> LineIndex;
 
+  // Returns the number of tokens produced by lexing, excluding recovery tokens
+  // added after lexing finished. This is the end of the `tokens()` range.
+  auto num_lexed_tokens() const -> int32_t {
+    return token_infos_.size() - post_lexing_recovery_tokens_;
+  }
+
   // Adds the token and adjusts the expected tree size.
   auto AddToken(TokenInfo info) -> TokenIndex;
 
+  // Builds `token_kinds_` and `token_has_leading_space_` from `token_infos_`,
+  // including the `FileEnd` padding. The lexer calls this once its token
+  // sequence is complete, and again if bracket recovery rewrites the sequence.
+  // Until then, `GetKind`, `HasLeadingWhitespace`, and `HasTrailingWhitespace`
+  // must not be used; the lexer reads `token_infos_` directly instead.
+  auto BuildTokenArrays() -> void;
+
   // Adds a post-lexing recovery token and adjusts the buffer state.
   auto AddPostLexingRecoveryToken(TokenInfo info) -> TokenIndex;
 
@@ -366,6 +393,29 @@ class TokenizedBuffer : public Printable<TokenizedBuffer> {
 
   ValueStore<TokenIndex, TokenInfo> token_infos_;
 
+  // The kind of each token in `token_infos_`, as a flat array indexed by the
+  // token index, followed by `KindDataPadding` entries of `TokenKind::FileEnd`.
+  // The kind is also stored in each `TokenInfo`, but this flat array makes the
+  // kind, which is by far the hottest token query after lexing, a single
+  // contiguous load rather than a chunked `ValueStore` lookup. It also allows
+  // scanning kinds in bulk; see `kind_data`.
+  //
+  // This is built from `token_infos_` in one pass once lexing is complete, by
+  // `BuildTokenArrays`, rather than maintained while lexing: the lexer's token
+  // emission is hot enough that even a few extra instructions there cost more
+  // than the pass.
+  llvm::SmallVector<TokenKind, 0> token_kinds_;
+
+  // Whether each token in `token_infos_` has leading whitespace, as a flat
+  // array indexed by the token index. This duplicates
+  // `TokenInfo::has_leading_space` for the same reason as `token_kinds_` and is
+  // built at the same time.
+  //
+  // This intentionally uses a byte-per-token `SmallVector<bool, 0>` rather than
+  // `llvm::BitVector` to facilitate vectorized population in
+  // `BuildTokenArrays`.
+  llvm::SmallVector<bool, 0> token_has_leading_space_;
+
   ValueStore<LineIndex, LineInfo> line_infos_;
 
   // Comments in the file.
@@ -402,19 +452,28 @@ class TokenizedBuffer : public Printable<TokenizedBuffer> {
 };
 
 inline auto TokenizedBuffer::GetKind(TokenIndex token) const -> TokenKind {
-  return token_infos_.Get(token).kind();
+  CARBON_DCHECK(token.index >= 0 && token.index < size(), "{0}", token);
+  CARBON_DCHECK(static_cast<size_t>(token.index) < token_kinds_.size(),
+                "Token arrays haven't been built yet");
+  return token_kinds_[token.index];
 }
 
 inline auto TokenizedBuffer::HasLeadingWhitespace(TokenIndex token) const
     -> bool {
-  return token_infos_.Get(token).has_leading_space();
+  CARBON_DCHECK(token.index >= 0 && token.index < size(), "{0}", token);
+  CARBON_DCHECK(
+      static_cast<size_t>(token.index) < token_has_leading_space_.size(),
+      "Token arrays haven't been built yet");
+  return token_has_leading_space_[token.index];
 }
 
 inline auto TokenizedBuffer::HasTrailingWhitespace(TokenIndex token) const
     -> bool {
-  TokenIterator it(token);
-  ++it;
-  return it != tokens().end() && token_infos_.Get(*it).has_leading_space();
+  CARBON_DCHECK(token.index >= 0 && token.index < size(), "{0}", token);
+  // Trailing whitespace is the leading whitespace of the next lexed token, if
+  // there is one.
+  int32_t next = token.index + 1;
+  return next < num_lexed_tokens() && HasLeadingWhitespace(TokenIndex(next));
 }
 
 inline auto TokenizedBuffer::AddToken(TokenInfo info) -> TokenIndex {

toolchain/lex/tokenized_buffer_test.cpp

@@ -714,6 +714,99 @@ TEST_F(LexerTest, Whitespace) {
   ASSERT_EQ(pos + 1, std::size(space));
 }
 
+TEST_F(LexerTest, KindDataAndWhitespaceForLongInput) {
+  // The token kinds and whitespace flags are built in bulk once lexing
+  // finishes, so check them on an input long enough to exercise the bulk path
+  // as well as its tail, with a mix of kinds and whitespace.
+  constexpr llvm::StringLiteral Line =
+      "fn x(a: i32) -> bool { return a+1 == 2; }\n";
+  constexpr int Repeats = 37;
+  std::string source;
+  for (int i = 0; i < Repeats; ++i) {
+    source += Line;
+  }
+  auto& buffer = compile_helper_.GetTokenizedBuffer(source);
+  EXPECT_FALSE(buffer.has_errors());
+
+  struct ExpectedLineToken {
+    TokenKind kind;
+    bool has_leading_space;
+  };
+  const ExpectedLineToken expected_line[] = {
+      {TokenKind::Fn, true},              //
+      {TokenKind::Identifier, true},      //
+      {TokenKind::OpenParen, false},      //
+      {TokenKind::Identifier, false},     //
+      {TokenKind::Colon, false},          //
+      {TokenKind::IntTypeLiteral, true},  //
+      {TokenKind::CloseParen, false},     //
+      {TokenKind::MinusGreater, true},    //
+      {TokenKind::Bool, true},            //
+      {TokenKind::OpenCurlyBrace, true},  //
+      {TokenKind::Return, true},          //
+      {TokenKind::Identifier, true},      //
+      {TokenKind::Plus, false},           //
+      {TokenKind::IntLiteral, false},     //
+      {TokenKind::EqualEqual, true},      //
+      {TokenKind::IntLiteral, true},      //
+      {TokenKind::Semi, false},           //
+      {TokenKind::CloseCurlyBrace, true},
+  };
+  constexpr int TokensPerLine = std::size(expected_line);
+  ASSERT_EQ(buffer.size(), 2 + Repeats * TokensPerLine);
+
+  const TokenKind* kind_data = buffer.kind_data();
+  for (TokenIndex token : buffer.tokens()) {
+    SCOPED_TRACE(
+        llvm::formatv("Token #{0}: '{1}'", token, buffer.GetTokenText(token)));
+    ExpectedLineToken expected;
+    if (token.index == 0) {
+      expected = {TokenKind::FileStart, false};
+    } else if (token.index == buffer.size() - 1) {
+      expected = {TokenKind::FileEnd, true};
+    } else {
+      // The start of the file counts as whitespace, so the first `fn` matches
+      // the pattern too.
+      expected = expected_line[(token.index - 1) % TokensPerLine];
+    }
+    EXPECT_EQ(buffer.GetKind(token), expected.kind);
+    EXPECT_EQ(kind_data[token.index], expected.kind);
+    EXPECT_EQ(buffer.HasLeadingWhitespace(token), expected.has_leading_space);
+    if (token.index + 1 < buffer.size()) {
+      EXPECT_EQ(buffer.HasTrailingWhitespace(token),
+                buffer.HasLeadingWhitespace(TokenIndex(token.index + 1)));
+    } else {
+      EXPECT_FALSE(buffer.HasTrailingWhitespace(token));
+    }
+  }
+  for (int i = 0; i < TokenizedBuffer::KindDataPadding; ++i) {
+    EXPECT_EQ(kind_data[buffer.size() + i], TokenKind::FileEnd) << i;
+  }
+
+  // Adding recovery tokens after lexing extends the arrays and keeps the
+  // padding intact.
+  int lexed_size = buffer.size();
+  for (int i = 0; i < 3; ++i) {
+    TokenIndex fn_token(1 + i * TokensPerLine);
+    TokenIndex recovered =
+        buffer.AddPostLexingRecoveryTokenAsIdentifier(fn_token);
+    EXPECT_EQ(recovered.index, lexed_size + i);
+    EXPECT_EQ(buffer.GetKind(recovered), TokenKind::Identifier);
+    EXPECT_EQ(buffer.GetTokenText(recovered), "fn");
+    EXPECT_EQ(buffer.HasLeadingWhitespace(recovered),
+              buffer.HasLeadingWhitespace(fn_token));
+  }
+  EXPECT_EQ(buffer.size(), lexed_size + 3);
+  EXPECT_EQ(buffer.tokens().end() - buffer.tokens().begin(), lexed_size);
+  kind_data = buffer.kind_data();
+  for (TokenIndex token : buffer.tokens()) {
+    EXPECT_EQ(kind_data[token.index], buffer.GetKind(token));
+  }
+  for (int i = 0; i < TokenizedBuffer::KindDataPadding; ++i) {
+    EXPECT_EQ(kind_data[buffer.size() + i], TokenKind::FileEnd) << i;
+  }
+}
+
 TEST_F(LexerTest, Keywords) {
   TokenKind keywords[] = {
 #define CARBON_TOKEN(TokenName)