use std::sync::Arc;
use rustc_hash::FxHashMap;
use crate::api::context::Context;
use crate::api::expr::{Ex, Expr, Numeric, SimplifyOpts};
use crate::base::arena::Arena;
use crate::base::errors::SymplexError;
use crate::base::node::{ExprId, ExprNode};
use crate::base::walk;
use crate::transforms::expand::ExpandOpts;
use crate::transforms::pattern::{
self, MATCH_BUDGET, MatchResult, Pattern, RawStep, Substitution, WildId,
};
#[derive(Clone, Debug)]
pub struct Step {
pub rule_name: String,
pub before: Ex,
pub after: Ex,
}
impl Step {
fn from_raw(template: &Ex, raw: RawStep) -> Step {
Step {
rule_name: raw.rule_name,
before: template.wrap(raw.before),
after: template.wrap(raw.after),
}
}
}
#[derive(Clone, Debug, Default)]
pub struct Bindings {
map: FxHashMap<String, Ex>,
}
impl Bindings {
#[must_use]
pub fn get(&self, name: &str) -> Option<&Ex> {
self.map.get(name)
}
#[must_use]
pub fn len(&self) -> usize {
self.map.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
pub fn iter(&self) -> impl Iterator<Item = (&str, &Ex)> {
self.map.iter().map(|(k, v)| (k.as_str(), v))
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum RewriteStrategy {
#[default]
BottomUp,
TopDown,
Innermost,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RewriteOpts {
pub max_iterations: usize,
pub strategy: RewriteStrategy,
}
impl Default for RewriteOpts {
fn default() -> Self {
RewriteOpts {
max_iterations: 50,
strategy: RewriteStrategy::BottomUp,
}
}
}
impl RewriteOpts {
#[must_use]
pub fn single_pass() -> Self {
RewriteOpts {
max_iterations: 1,
strategy: RewriteStrategy::BottomUp,
}
}
#[must_use]
pub fn max_iterations(mut self, n: usize) -> Self {
self.max_iterations = n;
self
}
#[must_use]
pub fn strategy(mut self, strategy: RewriteStrategy) -> Self {
self.strategy = strategy;
self
}
}
pub const MAX_REWRITE_OPS: usize = 100_000;
const MAX_NODE_REWRITES: usize = 32;
const MAX_INNERMOST_NESTING: usize = 8;
type GuardFn = Arc<dyn Fn(&Bindings) -> bool + Send + Sync>;
type RhsFn = Arc<dyn Fn(&Bindings) -> Option<Ex> + Send + Sync>;
#[derive(Clone)]
enum Rhs {
Template(ExprId),
Closure(RhsFn),
}
#[derive(Clone)]
enum Guard {
None,
Arena(fn(&Arena, &Substitution) -> bool),
Closure(GuardFn),
}
#[derive(Clone)]
pub struct Rule {
name: String,
lhs: Ex,
pattern: Pattern,
wild_names: FxHashMap<WildId, String>,
rhs: Rhs,
guard: Guard,
}
impl std::fmt::Debug for Rule {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Rule")
.field("name", &self.name)
.field("lhs", &self.lhs)
.finish_non_exhaustive()
}
}
impl Rule {
#[must_use]
pub fn new(name: impl Into<String>, lhs: &Ex, rhs: &Ex) -> Rule {
let rhs_id = lhs.checked_id(rhs);
let (pattern, wild_names) = Self::compile_lhs(lhs);
Rule {
name: name.into(),
lhs: lhs.clone(),
pattern,
wild_names,
rhs: Rhs::Template(rhs_id),
guard: Guard::None,
}
}
pub fn try_new(name: impl Into<String>, lhs: &Ex, rhs: &Ex) -> Result<Rule, SymplexError> {
let rule = Rule::new(name, lhs, rhs);
if rule.pattern.wilds.contains_key(&rule.pattern.root) {
return Err(SymplexError::InvalidArgument {
operation: "Rule::try_new",
reason: "left-hand side is a bare wildcard and would match every expression"
.to_string(),
});
}
let lhs_names: Vec<&String> = rule.wild_names.values().collect();
let unbound = {
let inner = lhs.inner.read();
let rhs_id = rule.lhs.checked_id(rhs);
walk::post_order_ids(&inner.arena, rhs_id)
.into_iter()
.filter(|&id| pattern::is_wild_symbol(&inner.arena, id))
.filter_map(|id| match inner.arena.node(id) {
ExprNode::Symbol(sid) => Some(inner.arena.symbol_name(*sid).to_string()),
_ => None,
})
.find(|n| !lhs_names.contains(&n))
};
if let Some(n) = unbound {
return Err(SymplexError::InvalidArgument {
operation: "Rule::try_new",
reason: format!(
"wildcard `{n}` appears in the right-hand side but not in the left-hand side"
),
});
}
Ok(rule)
}
#[must_use]
pub fn new_with_guard(
name: impl Into<String>,
lhs: &Ex,
rhs: &Ex,
guard: impl Fn(&Bindings) -> bool + Send + Sync + 'static,
) -> Rule {
let mut rule = Rule::new(name, lhs, rhs);
rule.guard = Guard::Closure(Arc::new(guard));
rule
}
#[must_use]
pub fn new_fn(
name: impl Into<String>,
lhs: &Ex,
f: impl Fn(&Bindings) -> Option<Ex> + Send + Sync + 'static,
) -> Rule {
let (pattern, wild_names) = Self::compile_lhs(lhs);
Rule {
name: name.into(),
lhs: lhs.clone(),
pattern,
wild_names,
rhs: Rhs::Closure(Arc::new(f)),
guard: Guard::None,
}
}
#[must_use]
pub fn from_macro_rule(ctx: &Context, rule: pattern::Rule) -> Rule {
let pattern::Rule {
name,
pattern,
template,
condition,
} = rule;
let lhs = Ex::from_raw_parts(ctx.id, Arc::clone(&ctx.inner), pattern.root);
let wild_names = {
let inner = ctx.inner.read();
pattern
.wilds
.iter()
.map(|(&id, &wid)| {
let name = match inner.arena.node(id) {
ExprNode::Symbol(sid) => inner.arena.symbol_name(*sid).to_string(),
_ => format!("_w{}", wid.0),
};
(wid, name)
})
.collect()
};
Rule {
name: name.to_string(),
lhs,
pattern,
wild_names,
rhs: Rhs::Template(template),
guard: match condition {
Some(f) => Guard::Arena(f),
None => Guard::None,
},
}
}
fn compile_lhs(lhs: &Ex) -> (Pattern, FxHashMap<WildId, String>) {
let inner = lhs.inner.read();
pattern::pattern_from_expr(&inner.arena, lhs.raw_id())
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub fn lhs(&self) -> &Ex {
&self.lhs
}
#[must_use]
pub fn wildcards(&self) -> Vec<String> {
let mut v: Vec<String> = self.wild_names.values().cloned().collect();
v.sort();
v
}
#[must_use]
pub fn matches(&self, expr: &Ex) -> Option<Bindings> {
let id = self.lhs.checked_id(expr);
let matches = {
let mut inner = self.lhs.inner.write();
pattern::match_all(&mut inner.arena, &self.pattern, id, false, 8, MATCH_BUDGET)
};
matches.into_iter().find_map(|m| {
self.check_guard(&m)
.map(|b| b.unwrap_or_else(|| self.bindings_of(&m.bindings)))
})
}
#[must_use]
pub fn apply(&self, expr: &Ex) -> Option<Ex> {
let id = self.lhs.checked_id(expr);
self.apply_id(id, true).map(|r| self.lhs.wrap(r))
}
fn bindings_of(&self, subs: &Substitution) -> Bindings {
let mut map = FxHashMap::default();
for (&wid, &id) in subs {
let name = self
.wild_names
.get(&wid)
.cloned()
.unwrap_or_else(|| format!("_w{}", wid.0));
map.insert(name, self.lhs.wrap(id));
}
Bindings { map }
}
fn check_guard(&self, m: &MatchResult) -> Option<Option<Bindings>> {
match &self.guard {
Guard::None => Some(None),
Guard::Arena(f) => {
let inner = self.lhs.inner.read();
if f(&inner.arena, &m.bindings) {
Some(None)
} else {
None
}
}
Guard::Closure(g) => {
let b = self.bindings_of(&m.bindings);
if g(&b) { Some(Some(b)) } else { None }
}
}
}
pub(crate) fn apply_id(&self, id: ExprId, allow_partial: bool) -> Option<ExprId> {
let max_results = match (&self.guard, &self.rhs) {
(Guard::None, Rhs::Template(_)) => 1,
_ => 8,
};
let matches = {
let mut inner = self.lhs.inner.write();
pattern::match_all(
&mut inner.arena,
&self.pattern,
id,
allow_partial,
max_results,
MATCH_BUDGET,
)
};
for m in matches {
let Some(guard_bindings) = self.check_guard(&m) else {
continue;
};
let replacement = match &self.rhs {
Rhs::Template(t) => {
let mut inner = self.lhs.inner.write();
pattern::instantiate(&mut inner.arena, *t, &self.pattern.wilds, &m.bindings)
}
Rhs::Closure(f) => {
let b = guard_bindings.unwrap_or_else(|| self.bindings_of(&m.bindings));
match f(&b) {
Some(r) => self.lhs.checked_id(&r),
None => continue,
}
}
};
let mut inner = self.lhs.inner.write();
let result = pattern::reattach_leftover(&mut inner.arena, id, replacement, &m.leftover);
return Some(result);
}
None
}
}
#[derive(Clone, Debug, Default)]
pub struct RuleSet {
rules: Vec<Rule>,
}
impl RuleSet {
#[must_use]
pub fn new() -> Self {
RuleSet::default()
}
#[must_use]
pub fn from_rules(rules: Vec<Rule>) -> Self {
RuleSet { rules }
}
#[must_use]
pub fn from_macro_rules(ctx: &Context, rules: Vec<pattern::Rule>) -> Self {
RuleSet {
rules: rules
.into_iter()
.map(|r| Rule::from_macro_rule(ctx, r))
.collect(),
}
}
#[must_use]
pub fn standard(ctx: &Context) -> Self {
let raw = {
let mut inner = ctx.inner.write();
pattern::basic_rules(&mut inner.arena)
};
RuleSet::from_macro_rules(ctx, raw)
}
pub fn push(&mut self, rule: Rule) {
self.rules.push(rule);
}
#[must_use]
pub fn with(mut self, rule: Rule) -> Self {
self.rules.push(rule);
self
}
pub fn extend(&mut self, other: &RuleSet) {
self.rules.extend(other.rules.iter().cloned());
}
#[must_use]
pub fn len(&self) -> usize {
self.rules.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.rules.is_empty()
}
pub fn iter(&self) -> impl Iterator<Item = &Rule> {
self.rules.iter()
}
#[must_use]
pub fn rules(&self) -> &[Rule] {
&self.rules
}
fn apply_first(&self, id: ExprId) -> Option<(ExprId, &str)> {
for rule in &self.rules {
if let Some(r) = rule.apply_id(id, true)
&& r != id
{
return Some((r, rule.name.as_str()));
}
}
None
}
}
impl FromIterator<Rule> for RuleSet {
fn from_iter<I: IntoIterator<Item = Rule>>(iter: I) -> Self {
RuleSet {
rules: iter.into_iter().collect(),
}
}
}
impl From<Vec<Rule>> for RuleSet {
fn from(rules: Vec<Rule>) -> Self {
RuleSet { rules }
}
}
impl std::ops::Index<usize> for RuleSet {
type Output = Rule;
fn index(&self, i: usize) -> &Rule {
&self.rules[i]
}
}
struct Driver<'a> {
template: &'a Ex,
rules: &'a RuleSet,
steps: Vec<RawStep>,
trace: bool,
}
impl Driver<'_> {
fn with_arena<R>(&self, f: impl FnOnce(&Arena) -> R) -> R {
let inner = self.template.inner.read();
f(&inner.arena)
}
fn with_arena_mut<R>(&self, f: impl FnOnce(&mut Arena) -> R) -> R {
let mut inner = self.template.inner.write();
f(&mut inner.arena)
}
fn record(&mut self, name: &str, before: ExprId, after: ExprId) {
if self.trace {
self.steps.push(RawStep {
rule_name: name.to_string(),
before,
after,
});
}
}
fn rewrite_at(&mut self, id: ExprId) -> Option<ExprId> {
let (new, name) = self.rules.apply_first(id)?;
let name = name.to_string();
self.record(&name, id, new);
Some(new)
}
fn rewrite_at_fixpoint(&mut self, id: ExprId) -> ExprId {
let mut current = id;
for _ in 0..MAX_NODE_REWRITES {
match self.rewrite_at(current) {
Some(new) if new != current => current = new,
_ => break,
}
}
current
}
fn pass_bottom_up(&mut self, root: ExprId) -> ExprId {
let post_order = self.with_arena(|a| walk::post_order_ids(a, root));
let mut cache: FxHashMap<ExprId, ExprId> = FxHashMap::default();
for &id in &post_order {
let rebuilt = self.with_arena_mut(|a| {
if a.node(id).is_atom() {
id
} else {
walk::rebuild_with_cache(a, id, &cache)
}
});
let rewritten = self.rewrite_at(rebuilt).unwrap_or(rebuilt);
cache.insert(id, rewritten);
}
cache.get(&root).copied().unwrap_or(root)
}
fn pass_innermost(&mut self, root: ExprId, nesting: usize) -> ExprId {
let post_order = self.with_arena(|a| walk::post_order_ids(a, root));
let mut cache: FxHashMap<ExprId, ExprId> = FxHashMap::default();
for &id in &post_order {
let rebuilt = self.with_arena_mut(|a| {
if a.node(id).is_atom() {
id
} else {
walk::rebuild_with_cache(a, id, &cache)
}
});
let mut current = rebuilt;
for _ in 0..MAX_NODE_REWRITES {
let Some(new) = self.rewrite_at(current) else {
break;
};
if new == current {
break;
}
current = if nesting < MAX_INNERMOST_NESTING {
self.pass_innermost(new, nesting + 1)
} else {
new
};
}
cache.insert(id, current);
}
cache.get(&root).copied().unwrap_or(root)
}
fn pass_top_down(&mut self, root: ExprId) -> ExprId {
let mut root_rw: FxHashMap<ExprId, ExprId> = FxHashMap::default();
let mut done: FxHashMap<ExprId, ExprId> = FxHashMap::default();
let r0 = self.rewrite_at_fixpoint(root);
root_rw.insert(root, r0);
let mut stack: Vec<(ExprId, bool)> = vec![(r0, false)];
while let Some(&(rid, expanded)) = stack.last() {
if done.contains_key(&rid) {
stack.pop();
continue;
}
if !expanded {
if let Some(top) = stack.last_mut() {
top.1 = true;
}
let children = self.with_arena(|a| a.children(rid));
for &c in children.iter().rev() {
let rc = match root_rw.get(&c) {
Some(&rc) => rc,
None => {
let rc = self.rewrite_at_fixpoint(c);
root_rw.insert(c, rc);
rc
}
};
if !done.contains_key(&rc) {
stack.push((rc, false));
}
}
} else {
stack.pop();
let children = self.with_arena(|a| a.children(rid));
let mut cache: FxHashMap<ExprId, ExprId> = FxHashMap::default();
for &c in &children {
let rc = root_rw.get(&c).copied().unwrap_or(c);
let fc = done.get(&rc).copied().unwrap_or(rc);
cache.insert(c, fc);
}
let rebuilt = self.with_arena_mut(|a| {
if a.node(rid).is_atom() {
rid
} else {
walk::rebuild_with_cache(a, rid, &cache)
}
});
done.insert(rid, rebuilt);
}
}
done.get(&r0).copied().unwrap_or(r0)
}
fn run(&mut self, root: ExprId, opts: &RewriteOpts) -> ExprId {
let mut current = root;
for _ in 0..opts.max_iterations {
let steps_before = self.steps.len();
let next = match opts.strategy {
RewriteStrategy::BottomUp => self.pass_bottom_up(current),
RewriteStrategy::TopDown => self.pass_top_down(current),
RewriteStrategy::Innermost => self.pass_innermost(current, 0),
};
if next == current {
break;
}
let size = self.with_arena(|a| pattern::tree_size_capped(a, next, MAX_REWRITE_OPS + 1));
if size > MAX_REWRITE_OPS {
tracing::debug!(size, "rewrite: tree-size guard triggered, stopping");
self.steps.truncate(steps_before);
break;
}
current = next;
}
current
}
}
impl Expr<Numeric> {
#[must_use = "returns the rewritten form; does not modify in place"]
pub fn rewrite(&self, rules: &RuleSet) -> Ex {
self.rewrite_with(rules, &RewriteOpts::default())
}
#[must_use = "returns the rewritten form; does not modify in place"]
pub fn rewrite_once(&self, rules: &RuleSet) -> Ex {
self.rewrite_with(rules, &RewriteOpts::single_pass())
}
#[must_use = "returns the rewritten form and trace; does not modify in place"]
pub fn rewrite_traced(&self, rules: &RuleSet) -> (Ex, Vec<Step>) {
self.rewrite_with_traced(rules, &RewriteOpts::default())
}
#[must_use = "returns the rewritten form; does not modify in place"]
pub fn rewrite_with(&self, rules: &RuleSet, opts: &RewriteOpts) -> Ex {
let mut driver = Driver {
template: self,
rules,
steps: Vec::new(),
trace: false,
};
let id = driver.run(self.raw_id(), opts);
self.wrap(id)
}
#[must_use = "returns the rewritten form and trace; does not modify in place"]
pub fn rewrite_with_traced(&self, rules: &RuleSet, opts: &RewriteOpts) -> (Ex, Vec<Step>) {
let mut driver = Driver {
template: self,
rules,
steps: Vec::new(),
trace: true,
};
let id = driver.run(self.raw_id(), opts);
let steps = driver
.steps
.into_iter()
.map(|s| Step::from_raw(self, s))
.collect();
(self.wrap(id), steps)
}
#[must_use = "returns the simplified form; does not modify in place"]
pub fn simplify_with_rules(&self, extra: &RuleSet) -> Ex {
let max = SimplifyOpts::default().max_iterations.max(1);
let mut current = self.clone();
for _ in 0..max {
let simplified = current.simplify();
let rewritten = simplified.rewrite(extra);
if rewritten == current {
return rewritten;
}
current = rewritten;
}
current
}
#[must_use = "returns the simplified form and trace; does not modify in place"]
pub fn simplify_traced(&self, opts: &SimplifyOpts) -> (Ex, Vec<Step>) {
let opts = opts.clone().trace();
let result = {
let mut inner = self.inner.write();
crate::simplify::simplify_engine::unified_simplify(
&mut inner.arena,
self.raw_id(),
&opts,
)
};
let steps = result
.steps
.into_iter()
.map(|s| Step::from_raw(self, s))
.collect();
(self.wrap(result.expr), steps)
}
}
impl Expr<Numeric> {
fn transform(&self, f: impl FnOnce(&mut Arena, ExprId) -> ExprId) -> Ex {
let id = {
let mut inner = self.inner.write();
f(&mut inner.arena, self.raw_id())
};
self.wrap(id)
}
#[must_use = "returns the expanded form; does not modify in place"]
pub fn expand_with(&self, opts: &ExpandOpts) -> Ex {
let opts = *opts;
self.transform(move |a, id| crate::transforms::expand::expand_with(a, id, &opts))
}
#[must_use = "returns the expanded form; does not modify in place"]
pub fn expand_power_base(&self, force: bool) -> Ex {
let opts = ExpandOpts::none().power_base(true).force(force);
self.expand_with(&opts)
}
#[must_use = "returns the expanded form; does not modify in place"]
pub fn expand_power_exp(&self, force: bool) -> Ex {
let opts = ExpandOpts::none().power_exp(true).force(force);
self.expand_with(&opts)
}
#[must_use = "returns the expanded form; does not modify in place"]
pub fn expand_multinomial(&self) -> Ex {
self.expand_with(&ExpandOpts::none().multinomial(true))
}
#[must_use = "returns the expanded form; does not modify in place"]
pub fn expand_log_with(&self, force: bool) -> Ex {
self.transform(move |a, id| crate::simplify::log_expand::expand_log_with(a, id, force))
}
#[must_use = "returns the combined form; does not modify in place"]
pub fn log_combine_with(&self, force: bool) -> Ex {
self.transform(move |a, id| crate::simplify::log_combine::log_combine_with(a, id, force))
}
#[must_use = "returns the denested form; does not modify in place"]
pub fn sqrtdenest(&self) -> Ex {
self.transform(crate::simplify::radsimp::sqrtdenest)
}
#[must_use = "returns the normalised form; does not modify in place"]
pub fn signsimp(&self) -> Ex {
self.transform(crate::simplify::factor_terms::signsimp)
}
#[must_use = "returns the denested form; does not modify in place"]
pub fn powdenest(&self, force: bool) -> Ex {
self.transform(move |a, id| crate::simplify::powsimp::powdenest_with(a, id, force))
}
#[must_use = "returns the collected form; does not modify in place"]
pub fn rcollect(&self, vars: &[&Ex]) -> Ex {
let ids: Vec<ExprId> = vars.iter().map(|v| self.checked_id(v)).collect();
self.transform(move |a, id| crate::simplify::factor_terms::rcollect(a, id, &ids))
}
#[must_use = "returns the collected form; does not modify in place"]
pub fn collect_const(&self) -> Ex {
self.transform(crate::simplify::factor_terms::collect_const)
}
#[must_use = "returns the recognised form; does not modify in place"]
pub fn nsimplify_with_constants(&self, constants: &[&Ex], tolerance: f64) -> Ex {
let ids: Vec<ExprId> = constants.iter().map(|c| self.checked_id(c)).collect();
self.transform(move |a, id| {
crate::simplify::nsimplify::nsimplify_with_constants(a, id, &ids, tolerance)
})
}
#[must_use = "returns the recognised form; does not modify in place"]
pub fn nsimplify(&self, tolerance: f64) -> Ex {
self.transform(move |a, id| {
let consts = crate::simplify::nsimplify::default_constants(a);
crate::simplify::nsimplify::nsimplify_with_constants(a, id, &consts, tolerance)
})
}
#[must_use]
pub fn separate_vars_additive(&self, vars: &[&Ex]) -> Vec<(Vec<Ex>, Ex)> {
let ids: Vec<ExprId> = vars.iter().map(|v| self.checked_id(v)).collect();
let raw = {
let mut inner = self.inner.write();
crate::domains::separatevars::separatevars_additive(
&mut inner.arena,
self.raw_id(),
&ids,
)
};
raw.into_iter()
.map(|(deps, sum)| {
(
deps.into_iter().map(|d| self.wrap(d)).collect(),
self.wrap(sum),
)
})
.collect()
}
#[must_use]
pub fn separate_vars_dict(&self, vars: &[&Ex]) -> Option<Vec<(Ex, Ex)>> {
let ids: Vec<ExprId> = vars.iter().map(|v| self.checked_id(v)).collect();
let factors = {
let mut inner = self.inner.write();
crate::domains::separatevars::separatevars_dict(&mut inner.arena, self.raw_id(), &ids)?
};
Some(
vars.iter()
.zip(factors)
.map(|(v, f)| ((*v).clone(), self.wrap(f)))
.collect(),
)
}
#[must_use = "returns the substituted form; does not modify in place"]
pub fn subs_algebraic(&self, old: &Ex, new: &Ex) -> Ex {
let old_id = self.checked_id(old);
let new_id = self.checked_id(new);
self.transform(move |a, id| pattern::subs_algebraic(a, id, old_id, new_id))
}
}