use std::collections::{HashMap, VecDeque};
use crate::ast::{CharVal, Node, NodeId, NumVal, Witness};
use crate::check::CheckedGrammar;
use crate::error::GenError;
use crate::rng::SplitMix64;
pub const DEFAULT_SPREAD: usize = 3;
pub const DEFAULT_GEN_DEPTH: usize = 16;
pub const DEFAULT_MAX_OUTPUT_LEN: usize = 1 << 20;
pub const DEFAULT_MAX_STEPS: u64 = 1 << 24;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct GenOptions {
pub max_depth: usize,
pub spread: usize,
pub coverage: bool,
pub preserve_case: bool,
pub max_output_len: Option<usize>,
pub max_steps: Option<u64>,
}
impl Default for GenOptions {
fn default() -> Self {
Self {
max_depth: DEFAULT_GEN_DEPTH,
spread: DEFAULT_SPREAD,
coverage: false,
preserve_case: false,
max_output_len: Some(DEFAULT_MAX_OUTPUT_LEN),
max_steps: Some(DEFAULT_MAX_STEPS),
}
}
}
pub struct Generator<'g> {
grammar: &'g CheckedGrammar,
rng: SplitMix64,
options: GenOptions,
coverage: Coverage,
pending: Vec<Unit>,
steps: u64,
emitted: usize,
}
impl<'g> Generator<'g> {
#[must_use]
pub fn new(grammar: &'g CheckedGrammar, seed: u64) -> Self {
Self {
grammar,
rng: SplitMix64::new(seed),
options: GenOptions::default(),
coverage: Coverage::new(grammar),
pending: Vec::new(),
steps: 0,
emitted: 0,
}
}
#[must_use]
pub fn with_options(mut self, options: GenOptions) -> Self {
self.options = options;
self
}
#[must_use]
pub const fn steps(&self) -> u64 {
self.steps
}
pub fn uncovered(&mut self, rule: &str) -> Result<usize, GenError> {
let body = self.start_rule(rule)?;
Ok(self.coverage.uncovered_reachable(self.grammar, body))
}
pub fn generate(&mut self, rule: &str) -> Result<String, GenError> {
let body = self.start_rule(rule)?;
let mut out = String::new();
self.steps = 0;
self.emitted = 0;
self.pending.clear();
self.walk(body, 0, &mut out)?;
for unit in std::mem::take(&mut self.pending) {
self.coverage.cover(unit);
}
Ok(out)
}
fn start_rule(&self, rule: &str) -> Result<NodeId, GenError> {
let found = self
.grammar
.rule(rule)
.ok_or_else(|| GenError::UnknownRule(rule.to_owned()))?;
self.grammar.can_generate(rule)?;
if !self.grammar.min_len(found.body).is_finite() {
return Err(GenError::NoFiniteExpansion {
rule: found.name.as_str().to_owned(),
});
}
Ok(found.body)
}
fn walk(&mut self, node: NodeId, depth: usize, out: &mut String) -> Result<(), GenError> {
self.step()?;
match self.grammar.node(node) {
Node::Alt { branches } => {
let branches = branches.clone();
let chosen = self.choose_branch(node, &branches, depth);
self.walk(branches[chosen], depth, out)
}
Node::Concat { items } => {
for item in items.clone() {
self.walk(item, depth, out)?;
}
Ok(())
}
Node::Repeat { repeat, body, .. } => {
let (repeat, body) = (*repeat, *body);
let count = self.choose_count(repeat.min, repeat.max, body, depth);
for _ in 0..count {
self.walk(body, depth, out)?;
}
Ok(())
}
Node::Optional { body } => {
let body = *body;
if self.choose_branch_of_optional(node, body, depth) == 0 {
self.walk(body, depth, out)?;
}
Ok(())
}
Node::RuleRef { .. } => {
let target = self
.grammar
.target(node)
.expect("`check` resolved every reference");
let body = self.grammar.rule_by_id(target).body;
self.walk(body, depth + 1, out)
}
Node::CharVal(value) => {
let value = value.clone();
self.emit_char_val(&value, out)
}
Node::NumVal { value, .. } => {
let value = value.clone();
self.emit_num_val(&value, out)
}
Node::ProseVal(_) => {
debug_assert!(false, "prose reached despite the start-rule gate");
Ok(())
}
}
}
fn choose_branch(&mut self, node: NodeId, branches: &[NodeId], depth: usize) -> usize {
let usable: Vec<usize> = (0..branches.len())
.filter(|index| self.grammar.min_len(branches[*index]).is_finite())
.collect();
debug_assert!(
!usable.is_empty(),
"a productive alternation has a productive branch"
);
let targets: Vec<NodeId> = usable.iter().map(|branch| branches[*branch]).collect();
let chosen = self.select(node, &usable, &targets, depth);
self.pending.push(Unit {
node,
branch: chosen as u32,
});
chosen
}
fn choose_branch_of_optional(&mut self, node: NodeId, body: NodeId, depth: usize) -> usize {
let (usable, targets) = if self.grammar.min_len(body).is_finite() {
(vec![0, 1], vec![body, node])
} else {
(vec![1], vec![node])
};
let chosen = self.select(node, &usable, &targets, depth);
self.pending.push(Unit {
node,
branch: chosen as u32,
});
chosen
}
fn select(
&mut self,
node: NodeId,
usable: &[usize],
targets: &[NodeId],
depth: usize,
) -> usize {
if self.options.coverage {
let untaken = usable.iter().copied().find(|branch| {
let unit = Unit {
node,
branch: *branch as u32,
};
!self.coverage.is_covered(unit) && !self.pending.contains(&unit)
});
if let Some(branch) = untaken {
return branch;
}
let mut nearest: Option<(u32, usize)> = None;
for (position, branch) in usable.iter().copied().enumerate() {
let target = targets[position];
if target == node {
continue;
}
let distance = self.coverage.distance_to_uncovered(target);
if distance != UNREACHABLE && nearest.is_none_or(|(best, _)| distance < best) {
nearest = Some((distance, branch));
}
}
if let Some((_, branch)) = nearest {
return branch;
}
}
if depth >= self.options.max_depth {
if let Some(Witness::Branch(branch)) = self.grammar.witness(node)
&& usable.contains(&(branch as usize))
{
return branch as usize;
}
}
let pick = self.rng.below(usable.len() as u64) as usize;
usable[pick]
}
fn choose_count(&mut self, min: u64, max: Option<u64>, body: NodeId, depth: usize) -> u64 {
if self.options.coverage
&& max.is_none_or(|max| max >= 1)
&& self.coverage.distance_to_uncovered(body) != UNREACHABLE
{
return min.max(1);
}
if depth >= self.options.max_depth {
return min; }
let ceiling = max
.unwrap_or(u64::MAX)
.min(min.saturating_add(self.options.spread as u64));
if ceiling <= min {
return min;
}
min + self.rng.below(ceiling - min + 1)
}
fn emit_char_val(&mut self, value: &CharVal, out: &mut String) -> Result<(), GenError> {
for ch in value.value.chars() {
let emitted = if value.case_sensitive || self.options.preserve_case {
ch
} else if self.rng.bool() {
ch.to_ascii_uppercase()
} else {
ch.to_ascii_lowercase()
};
self.push(emitted, out)?;
}
Ok(())
}
fn emit_num_val(&mut self, value: &NumVal, out: &mut String) -> Result<(), GenError> {
match value {
NumVal::Scalar(scalar) => {
let ch = scalar_at(*scalar).expect("`check` refused unrepresentable terminals");
self.push(ch, out)
}
NumVal::Range { lo, hi } => {
let count = NumVal::scalars_in(*lo, *hi);
debug_assert!(count > 0, "`check` refused empty ranges");
let index = self.rng.below(count);
let ch = NumVal::nth_scalar(*lo, *hi, index).expect("index is in range");
self.push(ch, out)
}
NumVal::Concat(values) => {
for scalar in values {
let ch = scalar_at(*scalar).expect("`check` refused unrepresentable terminals");
self.push(ch, out)?;
}
Ok(())
}
}
}
fn push(&mut self, ch: char, out: &mut String) -> Result<(), GenError> {
if self
.options
.max_output_len
.is_some_and(|limit| self.emitted >= limit)
{
return Err(GenError::OutputLimit);
}
self.emitted += 1;
out.push(ch);
Ok(())
}
fn step(&mut self) -> Result<(), GenError> {
self.steps += 1;
if self
.options
.max_steps
.is_some_and(|limit| self.steps > limit)
{
return Err(GenError::StepLimit);
}
Ok(())
}
}
fn scalar_at(value: u64) -> Option<char> {
char::from_u32(u32::try_from(value).ok()?)
}
const UNREACHABLE: u32 = u32::MAX;
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct Unit {
node: NodeId,
branch: u32,
}
struct Coverage {
units: Vec<Unit>,
index: HashMap<Unit, usize>,
covered: Vec<bool>,
reachable: HashMap<NodeId, Vec<usize>>,
edges: Vec<Vec<NodeId>>,
distance: Vec<u32>,
stale: bool,
}
impl Coverage {
fn new(grammar: &CheckedGrammar) -> Self {
let node_count = grammar.node_count();
let mut units = Vec::new();
let mut edges = vec![Vec::new(); node_count];
for (index, outgoing) in edges.iter_mut().enumerate() {
let node = NodeId::from_index(index as u32);
match grammar.node(node) {
Node::Alt { branches } => {
for (branch, child) in branches.iter().enumerate() {
if grammar.min_len(*child).is_finite() {
units.push(Unit {
node,
branch: branch as u32,
});
outgoing.push(*child);
}
}
}
Node::Optional { body } => {
if grammar.min_len(*body).is_finite() {
units.push(Unit { node, branch: 0 });
outgoing.push(*body);
}
units.push(Unit { node, branch: 1 });
}
Node::Concat { items } => outgoing.extend(items.iter().copied()),
Node::Repeat { repeat, body, .. } => {
if !repeat.is_never() {
outgoing.push(*body);
}
}
Node::RuleRef { .. } => {
if let Some(target) = grammar.target(node) {
outgoing.push(grammar.rule_by_id(target).body);
}
}
Node::CharVal(_) | Node::NumVal { .. } | Node::ProseVal(_) => {}
}
}
let index = units
.iter()
.enumerate()
.map(|(position, unit)| (*unit, position))
.collect();
let count = units.len();
Self {
units,
index,
covered: vec![false; count],
reachable: HashMap::new(),
edges,
distance: vec![UNREACHABLE; node_count],
stale: true,
}
}
fn uncovered_reachable(&mut self, grammar: &CheckedGrammar, body: NodeId) -> usize {
self.reachable_from(grammar, body);
self.reachable[&body]
.iter()
.filter(|position| !self.covered[**position])
.count()
}
fn reachable_from(&mut self, grammar: &CheckedGrammar, body: NodeId) -> &[usize] {
if !self.reachable.contains_key(&body) {
let mut seen = vec![false; grammar.node_count()];
let mut stack = vec![body];
seen[body.index()] = true;
let mut found = Vec::new();
while let Some(node) = stack.pop() {
for (position, unit) in self.units.iter().enumerate() {
if unit.node == node {
found.push(position);
}
}
for next in &self.edges[node.index()] {
if !seen[next.index()] {
seen[next.index()] = true;
stack.push(*next);
}
}
}
found.sort_unstable();
self.reachable.insert(body, found);
}
&self.reachable[&body]
}
fn cover(&mut self, unit: Unit) {
if let Some(position) = self.index.get(&unit)
&& !self.covered[*position]
{
self.covered[*position] = true;
self.stale = true;
}
}
fn refresh(&mut self) {
if !self.stale {
return;
}
self.stale = false;
let node_count = self.distance.len();
self.distance.fill(UNREACHABLE);
let mut incoming: Vec<Vec<usize>> = vec![Vec::new(); node_count];
for (from, targets) in self.edges.iter().enumerate() {
for target in targets {
incoming[target.index()].push(from);
}
}
let mut queue = VecDeque::new();
for (position, unit) in self.units.iter().enumerate() {
if !self.covered[position] {
let index = unit.node.index();
if self.distance[index] != 0 {
self.distance[index] = 0;
queue.push_back(index);
}
}
}
while let Some(node) = queue.pop_front() {
let next = self.distance[node] + 1;
for previous in &incoming[node] {
if self.distance[*previous] > next {
self.distance[*previous] = next;
queue.push_back(*previous);
}
}
}
}
fn distance_to_uncovered(&mut self, node: NodeId) -> u32 {
self.refresh();
self.distance[node.index()]
}
fn is_covered(&self, unit: Unit) -> bool {
self.index
.get(&unit)
.is_some_and(|position| self.covered[*position])
}
}