Store compressed parse state symbols and values in one array

This commit is contained in:
Max Brunsfeld 2026-05-23 18:50:26 -07:00
parent 591847a7d0
commit 5e2fb7d8bb
5 changed files with 138 additions and 104 deletions

View file

@ -156,9 +156,14 @@ struct TSLanguage {
// [0, large_state_count), the CSR tier covers
// [large_state_count, large_state_count + csr_state_count), and the
// Small tier covers the rest.
const uint32_t *parse_table_row_offsets;
const uint16_t *parse_table_columns;
const uint16_t *parse_table_values;
//
// CSR rows live in a single flat array `compressed_parse_table` of
// interleaved (symbol, value) pairs sorted by symbol.
// `compressed_parse_table_map[i]` is the start of row `i` as a uint16_t
// index into `compressed_parse_table`; each row occupies `2 * NNZ_i`
// uint16_t slots.
const uint32_t *compressed_parse_table_map;
const uint16_t *compressed_parse_table;
};
static inline bool set_contains(const TSCharacterRange *ranges, uint32_t len, int32_t lookahead) {

View file

@ -433,7 +433,11 @@ impl Generator {
}
// Recount tier boundaries from the (now-sorted) repr vector.
self.large_state_count = self.state_repr.iter().filter(|r| **r == Repr::Dense).count();
self.large_state_count = self
.state_repr
.iter()
.filter(|r| **r == Repr::Dense)
.count();
self.csr_state_count = self.state_repr.iter().filter(|r| **r == Repr::Csr).count();
}
@ -566,11 +570,7 @@ impl Generator {
}
let group_count = term_groups.len() + nonterm_groups.len();
out.push([
symbol_count * 2,
nnz * 4 + 4,
2 * nnz + 4 * group_count + 6,
]);
out.push([symbol_count * 2, nnz * 4 + 4, 2 * nnz + 4 * group_count + 6]);
}
out
}
@ -619,7 +619,10 @@ impl Generator {
}
add_line!(self, "#define PER_STATE_OPT_BYTES {sum_opt}");
add_line!(self, "#define PER_STATE_OPT_BYTES_PRE_DEDUP {sum_opt_pre_dedup}");
add_line!(
self,
"#define PER_STATE_OPT_BYTES_PRE_DEDUP {sum_opt_pre_dedup}"
);
add_line!(self, "#define PER_STATE_DENSE_BYTES {sum_dense}");
add_line!(self, "#define PER_STATE_CSR_BYTES {sum_csr}");
add_line!(self, "#define PER_STATE_SMALL_BYTES {sum_small}");
@ -642,8 +645,7 @@ impl Generator {
// Picker bias: prefer Csr to Small unless Small is at least
// SMALL_VS_CSR_BIAS fraction smaller (Csr binary search is much
// faster than Small linear scan at parse time).
let small_beats_csr =
(c[2] as f64) <= SMALL_VS_CSR_BIAS * (c[1] as f64);
let small_beats_csr = (c[2] as f64) <= SMALL_VS_CSR_BIAS * (c[1] as f64);
let pick = if c[0] <= c[1] && (c[0] <= c[2] || !small_beats_csr) {
Repr::Dense
} else if !small_beats_csr || c[1] <= c[2] {
@ -660,7 +662,10 @@ impl Generator {
// since we still verify the promotion is locally cheaper.
let mut by_hash: FxHashMap<u64, Vec<usize>> = FxHashMap::default();
for i in 0..n {
by_hash.entry(self.canonical_small_hash(i)).or_default().push(i);
by_hash
.entry(self.canonical_small_hash(i))
.or_default()
.push(i);
}
for (_, group) in by_hash.into_iter().filter(|(_, g)| g.len() > 1) {
@ -1604,10 +1609,7 @@ impl Generator {
&mut next_parse_action_list_index,
);
self.add_three_way_parse_table(
&mut parse_table_entries,
&mut next_parse_action_list_index,
);
self.add_three_way_parse_table(&mut parse_table_entries, &mut next_parse_action_list_index);
if next_parse_action_list_index >= usize::from(u16::MAX) {
Err(RenderError::ParseTable(next_parse_action_list_index))?;
@ -1629,9 +1631,11 @@ impl Generator {
/// * `[0, LARGE_STATE_COUNT)` - Dense
/// `ts_parse_table[LARGE_STATE_COUNT][SYMBOL_COUNT]`, O(1) lookup.
/// * `[LARGE_STATE_COUNT, LARGE_STATE_COUNT + CSR_STATE_COUNT)` - CSR
/// `ts_parse_table_row_offsets[CSR_STATE_COUNT + 1]`,
/// `ts_parse_table_columns[NNZ]`, `ts_parse_table_values[NNZ]`,
/// O(log n) binary search on columns within the row.
/// `ts_compressed_parse_table_map[CSR_STATE_COUNT + 1]` and
/// `ts_compressed_parse_table[]` (interleaved symbol/value pairs,
/// sorted by symbol within each row). Map entries are uint16_t
/// indices into `ts_compressed_parse_table`. O(log n) binary search
/// by entry within each row.
/// * remainder - Small
/// grouped sparse `ts_small_parse_table[]` with
/// `ts_small_parse_table_map[]` for offsets, O(group_count + nnz)
@ -1645,7 +1649,11 @@ impl Generator {
let dense_count = self.large_state_count;
let csr_count = self.csr_state_count;
self.add_dense_section(dense_count, parse_table_entries, next_parse_action_list_index);
self.add_dense_section(
dense_count,
parse_table_entries,
next_parse_action_list_index,
);
self.add_csr_section(
dense_count,
csr_count,
@ -1670,7 +1678,10 @@ impl Generator {
if dense_count == 0 {
// The struct field still references ts_parse_table; keep a 1x1
// placeholder so the static initializer compiles.
add_line!(self, "static const uint16_t ts_parse_table[1][1] = {{{{0}}}};");
add_line!(
self,
"static const uint16_t ts_parse_table[1][1] = {{{{0}}}};"
);
add_line!(self, "");
return;
}
@ -1729,10 +1740,16 @@ impl Generator {
return;
}
let mut row_offsets = Vec::with_capacity(csr_count + 1);
let mut columns: Vec<u16> = Vec::new();
let mut values: Vec<u16> = Vec::new();
let mut entries_buf: Vec<(u16, u16)> = Vec::new();
// Each row contributes `2 * NNZ` uint16_t slots to `entries` (interleaved
// symbol/value pairs). `row_offsets[i]` is the uint16_t index where row
// `i` begins. The trailing entry is the total uint16_t length.
let mut row_offsets: Vec<u32> = Vec::with_capacity(csr_count + 1);
// Per-row buffer of (column, value, symbol-identifier) so we can sort by
// column and still emit a readable `sym_xxx, value` pair for each entry.
let mut entries_buf: Vec<(u16, u16, String)> = Vec::new();
// Per-row rendered text lines (e.g. `sym_foo, 42,`).
let mut rendered_rows: Vec<Vec<String>> = Vec::with_capacity(csr_count);
let mut total_u16: u32 = 0;
for (state_idx, state) in self
.parse_table
@ -1755,7 +1772,8 @@ impl Generator {
GotoAction::Goto(s) => *s,
GotoAction::ShiftExtra => state_idx,
};
entries_buf.push((col as u16, state_id as u16));
let id = self.symbol_ids[symbol].clone();
entries_buf.push((col as u16, state_id as u16, id));
}
for (symbol, entry) in &state.terminal_entries {
let col = match self.symbol_order.get(symbol) {
@ -1770,21 +1788,24 @@ impl Generator {
parse_table_entries,
next_parse_action_list_index,
);
entries_buf.push((col as u16, entry_id as u16));
let id = self.symbol_ids[symbol].clone();
entries_buf.push((col as u16, entry_id as u16, id));
}
entries_buf.sort_unstable_by_key(|&(col, _)| col);
row_offsets.push(columns.len() as u32);
for &(col, val) in &entries_buf {
columns.push(col);
values.push(val);
entries_buf.sort_unstable_by_key(|&(col, _, _)| col);
row_offsets.push(total_u16);
total_u16 += (entries_buf.len() as u32) * 2;
let mut row_lines = Vec::with_capacity(entries_buf.len());
for (_, val, id) in &entries_buf {
row_lines.push(format!("{id}, {val},"));
}
rendered_rows.push(row_lines);
}
row_offsets.push(columns.len() as u32);
let total_nnz = columns.len();
row_offsets.push(total_u16);
add_line!(
self,
"static const uint32_t ts_parse_table_row_offsets[CSR_STATE_COUNT + 1] = {{",
"static const uint32_t ts_compressed_parse_table_map[CSR_STATE_COUNT + 1] = {{",
);
indent!(self);
for (i, &off) in row_offsets.iter().enumerate() {
@ -1793,23 +1814,27 @@ impl Generator {
dedent!(self);
add_line!(self, "}};");
add_line!(self, "");
add_line!(self, "#define PARSE_TABLE_NNZ {total_nnz}");
add_line!(self, "");
// Single flat array of (sym, val) pairs. Designated initializers point
// to the uint16_t-index where each CSR row begins, so the source
// layout mirrors `ts_compressed_parse_table_map` directly. Array size
// is inferred from the initializer.
add_line!(
self,
"static const uint16_t ts_parse_table_columns[PARSE_TABLE_NNZ] = {{",
"static const uint16_t ts_compressed_parse_table[] = {{",
);
indent!(self);
self.add_chunked_u16_array(&columns);
for (i, row) in rendered_rows.iter().enumerate() {
if row.is_empty() {
continue;
}
let off = row_offsets[i];
add_line!(self, "[{off}] =");
indent!(self);
for line in row {
add_line!(self, "{line}");
}
dedent!(self);
add_line!(self, "}};");
add_line!(self, "");
add_line!(
self,
"static const uint16_t ts_parse_table_values[PARSE_TABLE_NNZ] = {{",
);
indent!(self);
self.add_chunked_u16_array(&values);
}
dedent!(self);
add_line!(self, "}};");
add_line!(self, "");
@ -1900,7 +1925,11 @@ impl Generator {
seen_data.insert(key, next_table_index);
small_state_indices.push(next_table_index);
add_line!(self, "[{next_table_index}] = {},", values_with_symbols.len());
add_line!(
self,
"[{next_table_index}] = {},",
values_with_symbols.len()
);
indent!(self);
next_table_index += 1;
@ -1924,7 +1953,10 @@ impl Generator {
add_line!(self, "}};");
add_line!(self, "");
add_line!(self, "static const uint32_t ts_small_parse_table_map[] = {{");
add_line!(
self,
"static const uint32_t ts_small_parse_table_map[] = {{"
);
indent!(self);
for i in small_offset..n {
add_line!(
@ -1938,18 +1970,6 @@ impl Generator {
add_line!(self, "");
}
fn add_chunked_u16_array(&mut self, data: &[u16]) {
for chunk in data.chunks(16) {
add_whitespace!(self);
for val in chunk {
add!(self, "{val}, ");
}
add_line!(self, "");
}
}
fn add_parse_action_list(&mut self, parse_table_entries: Vec<(usize, ParseTableEntry)>) {
add_line!(
self,
@ -2156,10 +2176,9 @@ impl Generator {
if self.csr_state_count > 0 {
add_line!(
self,
".parse_table_row_offsets = ts_parse_table_row_offsets,"
".compressed_parse_table_map = ts_compressed_parse_table_map,"
);
add_line!(self, ".parse_table_columns = ts_parse_table_columns,");
add_line!(self, ".parse_table_values = ts_parse_table_values,");
add_line!(self, ".compressed_parse_table = ts_compressed_parse_table,");
}
dedent!(self);

View file

@ -72,6 +72,11 @@ static inline bool ts_language_has_reduce_action(
// (O(log n) binary search on a sparse row), and Small (linear scan over
// grouped symbol/action sections). For ABI < 16, only the Dense and Small
// tiers exist (csr_state_count == 0).
//
// CSR layout: `compressed_parse_table` is a single flat array of interleaved
// (symbol, value) pairs sorted by symbol within each row. `compressed_parse_table_map[i]`
// is the index in `compressed_parse_table` (in uint16_t units) where CSR row `i`
// begins. Each row occupies `2 * NNZ_i` uint16_t slots.
static inline uint16_t ts_language_lookup(
const TSLanguage *self,
TSStateId state,
@ -82,14 +87,16 @@ static inline uint16_t ts_language_lookup(
} else if (self->abi_version >= LANGUAGE_VERSION_WITH_COMPRESSED_TABLES
&& state < self->large_state_count + self->csr_state_count) {
uint32_t csr_state = state - self->large_state_count;
uint32_t start = self->parse_table_row_offsets[csr_state];
uint32_t end = self->parse_table_row_offsets[csr_state + 1];
// Binary search for symbol in columns[start..end]
// Convert uint16_t-array offsets to entry indices (each entry = sym+val pair).
uint32_t start = self->compressed_parse_table_map[csr_state] >> 1;
uint32_t end = self->compressed_parse_table_map[csr_state + 1] >> 1;
// Binary search by entry. Entry i lives at compressed_parse_table[2*i] (sym)
// and compressed_parse_table[2*i+1] (val).
while (start < end) {
uint32_t mid = start + (end - start) / 2;
uint16_t col = self->parse_table_columns[mid];
uint16_t col = self->compressed_parse_table[2 * mid];
if (col == symbol) {
return self->parse_table_values[mid];
return self->compressed_parse_table[2 * mid + 1];
} else if (col < symbol) {
start = mid + 1;
} else {
@ -151,8 +158,9 @@ static inline LookaheadIterator ts_language_lookaheads(
// data is set to NULL to indicate CSR mode.
data = NULL;
uint32_t csr_state = state - self->large_state_count;
uint32_t start = self->parse_table_row_offsets[csr_state];
uint32_t end = self->parse_table_row_offsets[csr_state + 1];
// Map entries are in uint16_t units; convert to entry count.
uint32_t start = self->compressed_parse_table_map[csr_state] >> 1;
uint32_t end = self->compressed_parse_table_map[csr_state + 1] >> 1;
group_count = (uint16_t)(end - start); // NNZ per row <= symbol_count <= UINT16_MAX
} else {
data = &self->parse_table[state * self->symbol_count] - 1;
@ -196,9 +204,11 @@ static inline bool ts_lookahead_iterator__next(LookaheadIterator *self) {
else if (self->data == NULL) {
if (self->group_count == 0) return false;
uint32_t csr_state = self->state - self->language->large_state_count;
uint32_t pos = self->language->parse_table_row_offsets[csr_state + 1] - self->group_count;
self->symbol = self->language->parse_table_columns[pos];
self->table_value = self->language->parse_table_values[pos];
// Map entries are in uint16_t units; convert to entry index, then back to uint16_t index.
uint32_t end_entry = self->language->compressed_parse_table_map[csr_state + 1] >> 1;
uint32_t pos = end_entry - self->group_count;
self->symbol = self->language->compressed_parse_table[2 * pos];
self->table_value = self->language->compressed_parse_table[2 * pos + 1];
self->group_count--;
}

View file

@ -156,9 +156,14 @@ struct TSLanguage {
// [0, large_state_count), the CSR tier covers
// [large_state_count, large_state_count + csr_state_count), and the
// Small tier covers the rest.
const uint32_t *parse_table_row_offsets;
const uint16_t *parse_table_columns;
const uint16_t *parse_table_values;
//
// CSR rows live in a single flat array `compressed_parse_table` of
// interleaved (symbol, value) pairs sorted by symbol.
// `compressed_parse_table_map[i]` is the start of row `i` as a uint16_t
// index into `compressed_parse_table`; each row occupies `2 * NNZ_i`
// uint16_t slots.
const uint32_t *compressed_parse_table_map;
const uint16_t *compressed_parse_table;
};
static inline bool set_contains(const TSCharacterRange *ranges, uint32_t len, int32_t lookahead) {

View file

@ -162,10 +162,11 @@ typedef struct {
int32_t supertype_map_slices;
int32_t supertype_map_entries;
TSLanguageMetadata metadata;
// CSR-compressed parse table (ABI version >= 16)
int32_t parse_table_row_offsets;
int32_t parse_table_columns;
int32_t parse_table_values;
// CSR-compressed parse table (ABI version >= 16). `compressed_parse_table`
// is a single flat array of interleaved (symbol, value) pairs;
// `compressed_parse_table_map` holds uint16_t-indices of row starts.
int32_t compressed_parse_table_map;
int32_t compressed_parse_table;
} LanguageInWasmMemory;
// LexerInWasmMemory - The memory layout of a `TSLexer` when compiled to wasm32.
@ -1337,9 +1338,8 @@ const TSLanguage *ts_wasm_store_load_language(
wasm_language.parse_table,
wasm_language.small_parse_table,
wasm_language.small_parse_table_map,
wasm_language.parse_table_row_offsets,
wasm_language.parse_table_columns,
wasm_language.parse_table_values,
wasm_language.compressed_parse_table_map,
wasm_language.compressed_parse_table,
wasm_language.parse_actions,
wasm_language.symbol_names,
wasm_language.field_names,
@ -1603,32 +1603,28 @@ const TSLanguage *ts_wasm_store_load_language(
language->abi_version >= LANGUAGE_VERSION_WITH_COMPRESSED_TABLES &&
language->csr_state_count > 0
) {
language->parse_table_row_offsets = copy(
language->compressed_parse_table_map = copy(
&wasm_memory,
wasm_language.parse_table_row_offsets,
wasm_language.compressed_parse_table_map,
(language->csr_state_count + 1) * sizeof(uint32_t),
&valid_wasm_memory
);
if (!valid_wasm_memory) goto invalid_language_memory;
uint32_t total_nnz;
// The final map entry is the total length of `compressed_parse_table`
// in uint16_t units (= 2 * total NNZ entries).
uint32_t total_entries_u16;
if (!wasm_memory__read(
&wasm_memory,
wasm_language.parse_table_row_offsets + language->csr_state_count * sizeof(uint32_t),
&total_nnz,
sizeof(total_nnz)
wasm_language.compressed_parse_table_map + language->csr_state_count * sizeof(uint32_t),
&total_entries_u16,
sizeof(total_entries_u16)
)) {
goto invalid_language_memory;
}
language->parse_table_columns = copy(
language->compressed_parse_table = copy(
&wasm_memory,
wasm_language.parse_table_columns,
total_nnz * sizeof(uint16_t),
&valid_wasm_memory
);
language->parse_table_values = copy(
&wasm_memory,
wasm_language.parse_table_values,
total_nnz * sizeof(uint16_t),
wasm_language.compressed_parse_table,
total_entries_u16 * sizeof(uint16_t),
&valid_wasm_memory
);
if (!valid_wasm_memory) goto invalid_language_memory;
@ -2043,9 +2039,8 @@ void ts_wasm_language_release(const TSLanguage *self) {
ts_free((void *)self->reserved_words);
ts_free((void *)self->parse_actions);
ts_free((void *)self->parse_table);
ts_free((void *)self->parse_table_row_offsets);
ts_free((void *)self->parse_table_columns);
ts_free((void *)self->parse_table_values);
ts_free((void *)self->compressed_parse_table_map);
ts_free((void *)self->compressed_parse_table);
ts_free((void *)self->primary_state_ids);
ts_free((void *)self->public_symbol_map);
ts_free((void *)self->small_parse_table);