use std::sync::OnceLock;
use rudb_common::{Error, Result};
use crate::generated::keywords::KEYWORDS;
use crate::generated::rules::{CHILDREN, NODES, RULES, SYMBOLS};
use crate::matcher::SUGGESTIONS;
use crate::rules::{Node, Op, Suggestion};
const UNREACHABLE: u32 = u32::MAX;
const START: &str = "Statement";
const BUDGET: u32 = 60;
const REPEATS: u32 = 3;
const STACK: u32 = 512;
#[derive(Debug, Clone)]
pub struct Catalog {
pub tables: Vec<Table>,
pub functions: Vec<String>,
pub table_functions: Vec<String>,
pub types: Vec<String>,
pub schemas: Vec<String>,
pub catalogs: Vec<String>,
pub pragmas: Vec<String>,
pub settings: Vec<String>,
pub files: Vec<String>,
pub variables: Vec<String>,
}
#[derive(Debug, Clone)]
pub struct Table {
pub name: String,
pub columns: Vec<String>,
}
impl Default for Catalog {
fn default() -> Self {
Self {
tables: vec![
Table {
name: "t".into(),
columns: ["a", "b", "c"].iter().map(|name| (*name).into()).collect(),
},
Table {
name: "u".into(),
columns: ["x", "y"].iter().map(|name| (*name).into()).collect(),
},
],
functions: names(&["abs", "length", "upper", "count", "coalesce"]),
table_functions: names(&["range", "generate_series"]),
types: names(&["INTEGER", "VARCHAR", "DOUBLE", "BOOLEAN", "DATE"]),
schemas: names(&["main"]),
catalogs: names(&["memory"]),
pragmas: names(&["database_list", "show_tables"]),
settings: names(&["threads", "memory_limit"]),
files: names(&["'data.parquet'", "'out.csv'"]),
variables: names(&["v"]),
}
}
}
fn names(from: &[&str]) -> Vec<String> {
from.iter().map(|name| (*name).to_string()).collect()
}
const NUMBERS: [&str; 7] = ["0", "1", "2", "42", "1.5", "1e3", "9223372036854775807"];
const STRINGS: [&str; 4] = ["'a'", "''", "'it''s'", "'é'"];
const OPERATORS: [&str; 5] = ["||", "<<", ">>", "@>", "&&"];
#[derive(Debug, Clone)]
pub struct Generator {
catalog: Catalog,
budget: u32,
repeats: u32,
}
impl Default for Generator {
fn default() -> Self {
Self { catalog: Catalog::default(), budget: BUDGET, repeats: REPEATS }
}
}
impl Generator {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_catalog(catalog: Catalog) -> Self {
Self { catalog, ..Self::default() }
}
#[must_use]
pub fn budget(mut self, tokens: u32) -> Self {
self.budget = tokens.max(1);
self
}
#[must_use]
pub fn repeats(mut self, times: u32) -> Self {
self.repeats = times.max(1);
self
}
#[must_use]
pub fn statement(&self, seed: u64) -> String {
self.from_rule(START, seed).expect("Statement is a rule")
}
pub fn from_rule(&self, rule: &str, seed: u64) -> Result<String> {
let index = RULES
.binary_search_by(|candidate| candidate.name.cmp(rule))
.map_err(|_| Error::parser(format!("no rule named {rule}")))?;
let root = RULES[index].root;
if costs()[root as usize] == UNREACHABLE {
return Err(Error::parser(format!("nothing finite can be written from {rule}")));
}
let mut run = Run {
catalog: &self.catalog,
random: Random::new(seed),
pieces: Vec::new(),
spent: 0,
budget: self.budget,
repeats: self.repeats,
path: vec![0; RULES.len()],
stack: 0,
tight: false,
table: 0,
};
run.table = run.random.below(self.catalog.tables.len().max(1));
run.node(root);
Ok(run.pieces.join(" "))
}
}
struct Run<'a> {
catalog: &'a Catalog,
random: Random,
pieces: Vec<String>,
spent: u32,
budget: u32,
repeats: u32,
path: Vec<u32>,
stack: u32,
tight: bool,
table: usize,
}
impl Run<'_> {
fn cheap(&self) -> bool {
self.tight || self.spent >= self.budget
}
fn emit(&mut self, text: impl Into<String>) {
self.pieces.push(text.into());
self.spent += 1;
}
fn node(&mut self, index: u32) {
let node = NODES[index as usize];
match node.op {
Op::Rule => self.rule(node.a, node.b),
Op::Sequence => {
for child in node.children() {
self.node(*child);
}
}
Op::Choice => {
let child = self.alternative(node.children(), self.cheap());
self.node(child);
}
Op::Optional => {
if !self.cheap() && self.random.chance(3) {
self.node(node.a);
}
}
Op::Repeat => {
let times = if self.cheap() { 1 } else { self.random.count(1, 3) };
for _ in 0..times {
self.node(node.a);
}
}
Op::Keyword => {
let word = KEYWORDS[node.a as usize].0.to_uppercase();
self.emit(word);
}
Op::KeywordClass => self.keyword_in(node.a),
Op::Symbol => self.emit(SYMBOLS[node.a as usize]),
Op::Identifier => self.name(SUGGESTIONS[node.a as usize]),
Op::Number => {
let number = self.random.pick(&NUMBERS);
self.emit(number);
}
Op::String => {
let text = self.random.pick(&STRINGS);
self.emit(text);
}
Op::Operator => {
let operator = self.random.pick(&OPERATORS);
self.emit(operator);
}
Op::EndOfInput => {}
}
}
fn rule(&mut self, rule: u32, root: u32) {
let index = rule as usize;
let was = self.tight;
self.tight = was || self.path[index] >= self.repeats || self.stack >= STACK;
self.path[index] += 1;
self.stack += 1;
self.node(root);
self.stack -= 1;
self.path[index] -= 1;
self.tight = was;
}
fn alternative(&mut self, children: &[u32], cheap: bool) -> u32 {
let costs = costs();
if cheap {
let mut best = children[0];
for child in children {
if costs[*child as usize] < costs[best as usize] {
best = *child;
}
}
return best;
}
let weights: Vec<u64> =
children.iter().map(|child| weight(costs[*child as usize])).collect();
let total: u64 = weights.iter().sum();
if total == 0 {
return children[0];
}
let mut pick = self.random.next() % total;
for (child, weight) in children.iter().zip(&weights) {
if pick < *weight {
return *child;
}
pick -= *weight;
}
children[children.len() - 1]
}
fn keyword_in(&mut self, mask: u32) {
let words = keywords_in(mask as u8);
if words.is_empty() {
return;
}
let index = self.random.below(words.len());
let word = KEYWORDS[words[index] as usize].0.to_uppercase();
self.emit(word);
}
fn name(&mut self, suggestion: Suggestion) {
let catalog = self.catalog;
let table = catalog.tables.get(self.table);
let pool = match suggestion {
Suggestion::TableName => {
let name = table.map_or("t", |table| table.name.as_str()).to_string();
self.emit(name);
return;
}
Suggestion::ColumnName => {
let columns = table.map(|table| table.columns.as_slice()).unwrap_or_default();
if columns.is_empty() {
self.emit("a");
} else {
let index = self.random.below(columns.len());
let name = columns[index].clone();
self.emit(name);
}
return;
}
Suggestion::Variable => &catalog.variables,
Suggestion::ScalarFunctionName => &catalog.functions,
Suggestion::TableFunctionName => &catalog.table_functions,
Suggestion::TypeName => &catalog.types,
Suggestion::SchemaName => &catalog.schemas,
Suggestion::CatalogName => &catalog.catalogs,
Suggestion::PragmaName => &catalog.pragmas,
Suggestion::SettingName => &catalog.settings,
Suggestion::FileName => &catalog.files,
};
if pool.is_empty() {
self.emit("a");
return;
}
let index = self.random.below(pool.len());
let name = pool[index].clone();
self.emit(name);
}
}
fn weight(cost: u32) -> u64 {
if cost == UNREACHABLE {
return 0;
}
(16 / (u64::from(cost) + 1)).max(1)
}
fn costs() -> &'static [u32] {
static COSTS: OnceLock<Box<[u32]>> = OnceLock::new();
COSTS.get_or_init(build_costs)
}
fn build_costs() -> Box<[u32]> {
let mut costs = vec![UNREACHABLE; NODES.len()];
loop {
let mut moved = false;
for (index, node) in NODES.iter().enumerate() {
let value = cost_of(*node, &costs);
if value < costs[index] {
costs[index] = value;
moved = true;
}
}
if !moved {
return costs.into_boxed_slice();
}
}
}
fn cost_of(node: Node, costs: &[u32]) -> u32 {
match node.op {
Op::EndOfInput | Op::Optional => 0,
Op::Keyword
| Op::KeywordClass
| Op::Symbol
| Op::Identifier
| Op::Number
| Op::String
| Op::Operator => 1,
Op::Rule => costs[node.b as usize],
Op::Repeat => costs[node.a as usize],
Op::Sequence => CHILDREN[node.a as usize..(node.a + node.b) as usize]
.iter()
.fold(0, |total, child| total.saturating_add(costs[*child as usize])),
Op::Choice => CHILDREN[node.a as usize..(node.a + node.b) as usize]
.iter()
.map(|child| costs[*child as usize])
.min()
.unwrap_or(UNREACHABLE),
}
}
fn keywords_in(mask: u8) -> &'static [u16] {
static BY_CLASS: OnceLock<[Vec<u16>; 8]> = OnceLock::new();
let by_class = BY_CLASS.get_or_init(|| {
let mut lists: [Vec<u16>; 8] = Default::default();
for (index, (_, classes)) in KEYWORDS.iter().enumerate() {
for (bit, list) in lists.iter_mut().enumerate() {
if classes & (1 << bit) != 0 {
list.push(index as u16);
}
}
}
lists
});
match (0..8).find(|bit| mask & (1 << bit) != 0) {
Some(bit) => &by_class[bit],
None => &[],
}
}
struct Random(u64);
impl Random {
fn new(seed: u64) -> Self {
Self(seed.wrapping_mul(0x2545_f491_4f6c_dd1d) | 1)
}
fn next(&mut self) -> u64 {
self.0 ^= self.0 << 13;
self.0 ^= self.0 >> 7;
self.0 ^= self.0 << 17;
self.0
}
fn below(&mut self, bound: usize) -> usize {
(self.next() % bound.max(1) as u64) as usize
}
fn count(&mut self, low: usize, high: usize) -> usize {
low + self.below(high - low + 1)
}
fn chance(&mut self, one_in: u64) -> bool {
self.next() % one_in == 0
}
fn pick<T: Copy>(&mut self, from: &[T]) -> T {
from[self.below(from.len())]
}
}
#[cfg(test)]
mod tests {
use super::{Catalog, Generator, Random, Run, Table, UNREACHABLE, costs};
use crate::generated::rules::RULES;
use crate::matcher::parse_from;
use crate::tokenize::tokenize;
const SEEDS: u64 = 3000;
fn seeds() -> std::ops::RangeInclusive<u64> {
let first = setting("RUDB_GRAMMAR_SEED", 1);
let count = setting("RUDB_GRAMMAR_SEEDS", SEEDS).max(1);
first..=first.saturating_add(count - 1)
}
fn setting(name: &str, fallback: u64) -> u64 {
match std::env::var(name) {
Ok(text) => text
.trim()
.parse()
.unwrap_or_else(|_| panic!("{name} is {text}, which is not a number")),
Err(_) => fallback,
}
}
fn statements(generator: &Generator, rule: &str) -> Vec<String> {
seeds().map(|seed| generator.from_rule(rule, seed).expect("a rule")).collect()
}
fn cheapest(rule: &str) -> String {
let catalog = Catalog::default();
let root = RULES[RULES.binary_search_by(|r| r.name.cmp(rule)).expect("a rule")].root;
let mut run = Run {
catalog: &catalog,
random: Random::new(1),
pieces: Vec::new(),
spent: 0,
budget: 1,
repeats: 1,
path: vec![0; RULES.len()],
stack: 0,
tight: true,
table: 0,
};
run.node(root);
run.pieces.join(" ")
}
#[test]
fn the_same_seed_writes_the_same_statement() {
let generator = Generator::new();
for seed in [1, 2, 99, 10_000] {
assert_eq!(generator.statement(seed), generator.statement(seed));
}
assert_ne!(generator.statement(1), generator.statement(2));
}
#[test]
fn every_rule_in_the_table_can_be_written_down() {
let costs = costs();
let unreachable: Vec<&str> = RULES
.iter()
.filter(|rule| costs[rule.root as usize] == UNREACHABLE)
.map(|rule| rule.name)
.collect();
assert!(unreachable.is_empty(), "no finite text exists for {unreachable:?}");
}
#[test]
fn the_cheapest_text_a_rule_has_is_as_long_as_the_cost_table_says() {
for rule in &RULES {
let text = cheapest(rule.name);
let written = if text.is_empty() { 0 } else { text.split(' ').count() as u32 };
assert_eq!(written, costs()[rule.root as usize], "{} wrote {text:?}", rule.name);
}
}
#[test]
fn what_it_writes_always_tokenizes() {
for text in statements(&Generator::new(), "Statement") {
tokenize(&text).unwrap_or_else(|error| panic!("{text}\n{error}"));
}
}
#[test]
fn the_generator_writes_statements_that_parse() {
let written = statements(&Generator::new(), "Statement");
let parsed =
written.iter().filter(|text| parse_from(text, "Statement", true).is_ok()).count();
let total = written.len();
println!("{parsed} of {total} statements parse");
assert!(parsed * 100 / total >= 85, "{parsed} of {total} parse");
}
#[test]
fn the_filter_and_the_walk_without_it_agree_on_what_it_writes() {
for text in statements(&Generator::new(), "Statement") {
let filtered = parse_from(&text, "Statement", true);
let whole = parse_from(&text, "Statement", false);
assert_eq!(filtered.is_ok(), whole.is_ok(), "the filter changed the answer on {text}");
}
}
#[test]
fn the_names_it_writes_are_the_catalog_it_was_handed() {
let catalog = Catalog {
tables: vec![Table { name: "zork".into(), columns: vec!["quux".into()] }],
functions: vec!["frob".into()],
..Catalog::default()
};
let generator = Generator::with_catalog(catalog.clone());
let written = statements(&generator, "Statement");
let total = written.len();
let mut seen_a_column = false;
for text in written {
for piece in text.split(' ') {
if !piece.starts_with(|first: char| first.is_ascii_lowercase()) {
continue;
}
seen_a_column |= piece == "quux";
assert!(known(&catalog, piece), "{piece} is not a name the catalog has, in {text}");
}
}
assert!(seen_a_column, "no statement in {total} mentioned a column");
}
fn known(catalog: &Catalog, piece: &str) -> bool {
catalog
.tables
.iter()
.any(|table| table.name == piece || table.columns.iter().any(|column| column == piece))
|| [
&catalog.functions,
&catalog.table_functions,
&catalog.types,
&catalog.schemas,
&catalog.catalogs,
&catalog.pragmas,
&catalog.settings,
&catalog.files,
&catalog.variables,
]
.iter()
.any(|pool| pool.iter().any(|name| name == piece))
}
#[test]
fn a_query_is_what_comes_out_of_the_rule_that_writes_queries() {
let generator = Generator::new();
let queries = statements(&generator, "SelectStatement");
let parsed =
queries.iter().filter(|text| parse_from(text, "SelectStatement", true).is_ok()).count();
let total = queries.len();
println!("{parsed} of {total} queries parse");
assert!(parsed * 100 / total >= 80, "{parsed} of {total} parse");
assert!(queries.iter().any(|text| text.contains("SELECT")));
}
#[test]
fn a_budget_of_one_token_still_writes_a_whole_statement() {
assert_eq!(cheapest("Statement"), "SELECT");
assert!(parse_from("SELECT", "Statement", true).is_ok());
}
#[test]
fn the_shortest_explain_there_is_does_not_parse_and_that_is_the_grammar() {
assert_eq!(cheapest("ExplainStatement"), "EXPLAIN ANALYZE");
assert!(parse_from("EXPLAIN ANALYZE", "ExplainStatement", true).is_err());
assert!(parse_from("EXPLAIN ANALYZE SELECT 1", "Statement", true).is_ok());
}
#[test]
fn an_unknown_rule_says_so() {
let error = Generator::new().from_rule("NoSuchRule", 1).unwrap_err();
assert!(error.to_string().contains("no rule named NoSuchRule"), "{error}");
}
}