diff --git a/toolchain/base/value_store.h b/toolchain/base/value_store.h index bea2e2509675..445b439fd6a1 100644 --- a/toolchain/base/value_store.h +++ b/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(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); + 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. diff --git a/toolchain/lex/lex.cpp b/toolchain/lex/lex.cpp index b4d17bc9952b..4760a86b5cf9 100644 --- a/toolchain/lex/lex.cpp +++ b/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. diff --git a/toolchain/lex/token_info.h b/toolchain/lex/token_info.h index b38c04c7f940..363013c82094 100644 --- a/toolchain/lex/token_info.h +++ b/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; diff --git a/toolchain/lex/tokenized_buffer.cpp b/toolchain/lex/tokenized_buffer.cpp index f42b9f88e874..4b7423a50ebf 100644 --- a/toolchain/lex/tokenized_buffer.cpp +++ b/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_); } diff --git a/toolchain/lex/tokenized_buffer.h b/toolchain/lex/tokenized_buffer.h index 9e278e6a06e0..a73641b10e18 100644 --- a/toolchain/lex/tokenized_buffer.h +++ b/toolchain/lex/tokenized_buffer.h @@ -7,6 +7,7 @@ #include +#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 { 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 { // recovery tokens inserted inline by lexing, such as matching brackets, but // excludes recovery tokens added after lexing finished. auto tokens() const -> llvm::iterator_range { - 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 { 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 { ValueStore 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 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` rather than + // `llvm::BitVector` to facilitate vectorized population in + // `BuildTokenArrays`. + llvm::SmallVector token_has_leading_space_; + ValueStore line_infos_; // Comments in the file. @@ -402,19 +452,28 @@ class TokenizedBuffer : public Printable { }; 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(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(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 { diff --git a/toolchain/lex/tokenized_buffer_test.cpp b/toolchain/lex/tokenized_buffer_test.cpp index 4808d3545ea6..4bb81f0d28c6 100644 --- a/toolchain/lex/tokenized_buffer_test.cpp +++ b/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)