use std::cmp::Ordering;
use std::collections::HashMap;
use std::fmt::Debug;
use crate::{Analysis, EClass, EGraph, Id, Language, RecExpr};
pub struct Extractor<'a, CF: CostFunction<L>, L: Language, N: Analysis<L>> {
cost_function: CF,
costs: HashMap<Id, (CF::Cost, L)>,
egraph: &'a EGraph<L, N>,
}
pub trait CostFunction<L: Language> {
type Cost: PartialOrd + Debug + Clone;
fn cost<C>(&mut self, enode: &L, costs: C) -> Self::Cost
where
C: FnMut(Id) -> Self::Cost;
fn cost_rec(&mut self, expr: &RecExpr<L>) -> Self::Cost {
let mut costs: HashMap<Id, Self::Cost> = HashMap::default();
for (i, node) in expr.as_ref().iter().enumerate() {
let cost = self.cost(node, |i| costs[&i].clone());
costs.insert(Id::from(i), cost);
}
let last_id = Id::from(expr.as_ref().len() - 1);
costs[&last_id].clone()
}
}
pub struct AstSize;
impl<L: Language> CostFunction<L> for AstSize {
type Cost = usize;
fn cost<C>(&mut self, enode: &L, mut costs: C) -> Self::Cost
where
C: FnMut(Id) -> Self::Cost,
{
enode.fold(1, |sum, id| sum + costs(id))
}
}
pub struct AstDepth;
impl<L: Language> CostFunction<L> for AstDepth {
type Cost = usize;
fn cost<C>(&mut self, enode: &L, mut costs: C) -> Self::Cost
where
C: FnMut(Id) -> Self::Cost,
{
1 + enode.fold(0, |max, id| max.max(costs(id)))
}
}
fn cmp<T: PartialOrd>(a: &Option<T>, b: &Option<T>) -> Ordering {
match (a, b) {
(None, None) => Ordering::Equal,
(None, Some(_)) => Ordering::Greater,
(Some(_), None) => Ordering::Less,
(Some(a), Some(b)) => a.partial_cmp(&b).unwrap(),
}
}
impl<'a, CF, L, N> Extractor<'a, CF, L, N>
where
CF: CostFunction<L>,
L: Language,
N: Analysis<L>,
{
pub fn new(egraph: &'a EGraph<L, N>, cost_function: CF) -> Self {
let costs = HashMap::default();
let mut extractor = Extractor {
costs,
egraph,
cost_function,
};
extractor.find_costs();
extractor
}
pub fn find_best(&mut self, eclass: Id) -> (CF::Cost, RecExpr<L>) {
let mut expr = RecExpr::default();
let (_, cost) = self.find_best_rec(&mut expr, eclass);
(cost, expr)
}
fn find_best_rec(&mut self, expr: &mut RecExpr<L>, eclass: Id) -> (Id, CF::Cost) {
let id = self.egraph.find(eclass);
let (best_cost, best_node) = match self.costs.get(&id) {
Some(result) => result.clone(),
None => panic!("Failed to extract from eclass {}", id),
};
let node = best_node.map_children(|child| self.find_best_rec(expr, child).0);
(expr.add(node), best_cost)
}
fn node_total_cost(&mut self, node: &L) -> Option<CF::Cost> {
let eg = &self.egraph;
let has_cost = |&id| self.costs.contains_key(&eg.find(id));
if node.children().iter().all(has_cost) {
let costs = &self.costs;
let cost_f = |id| costs[&eg.find(id)].0.clone();
Some(self.cost_function.cost(&node, cost_f))
} else {
None
}
}
fn find_costs(&mut self) {
let mut did_something = true;
while did_something {
did_something = false;
for class in self.egraph.classes() {
let pass = self.make_pass(class);
match (self.costs.get(&class.id), pass) {
(None, Some(new)) => {
self.costs.insert(class.id, new);
did_something = true;
}
(Some(old), Some(new)) if new.0 < old.0 => {
self.costs.insert(class.id, new);
did_something = true;
}
_ => (),
}
}
}
for class in self.egraph.classes() {
if !self.costs.contains_key(&class.id) {
log::warn!(
"Failed to compute cost for eclass {}: {:?}",
class.id,
class.nodes
)
}
}
}
fn make_pass(&mut self, eclass: &EClass<L, N::Data>) -> Option<(CF::Cost, L)> {
let (cost, node) = eclass
.iter()
.map(|n| (self.node_total_cost(n), n))
.min_by(|a, b| cmp(&a.0, &b.0))
.unwrap_or_else(|| panic!("Can't extract, eclass is empty: {:#?}", eclass));
cost.map(|c| (c, node.clone()))
}
}