From 5e2fb7d8bb0e0b1f06927b369ee1f6c8e41e732c Mon Sep 17 00:00:00 2001 From: Max Brunsfeld Date: Sat, 23 May 2026 18:50:26 -0700 Subject: [PATCH] Store compressed parse state symbols and values in one array --- crates/generate/src/parser.h.inc | 11 ++- crates/generate/src/render.rs | 145 +++++++++++++++++-------------- lib/src/language.h | 30 ++++--- lib/src/parser.h | 11 ++- lib/src/wasm_store.c | 45 +++++----- 5 files changed, 138 insertions(+), 104 deletions(-) diff --git a/crates/generate/src/parser.h.inc b/crates/generate/src/parser.h.inc index ee67cfc44..2b8eb9937 100644 --- a/crates/generate/src/parser.h.inc +++ b/crates/generate/src/parser.h.inc @@ -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) { diff --git a/crates/generate/src/render.rs b/crates/generate/src/render.rs index a8657e8bc..8613b0244 100644 --- a/crates/generate/src/render.rs +++ b/crates/generate/src/render.rs @@ -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> = 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 = Vec::new(); - let mut values: Vec = 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 = 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::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); - 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); + 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); + } 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); diff --git a/lib/src/language.h b/lib/src/language.h index 4396d69aa..b696e5ade 100644 --- a/lib/src/language.h +++ b/lib/src/language.h @@ -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--; } diff --git a/lib/src/parser.h b/lib/src/parser.h index ee67cfc44..2b8eb9937 100644 --- a/lib/src/parser.h +++ b/lib/src/parser.h @@ -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) { diff --git a/lib/src/wasm_store.c b/lib/src/wasm_store.c index fe2248e68..904e9bdec 100644 --- a/lib/src/wasm_store.c +++ b/lib/src/wasm_store.c @@ -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);