///|
pub(all) struct Match {
  root : Id
  subst : Map[String, Id]
} derive(Show)

///|
pub struct Rewrite {
  name : String
  searcher : Searcher
  applier : Applier
  condition : (EGraph, Map[String, Id]) -> Bool
}

///|
pub struct RewriteStats {
  name : String
  matched : Int
  applied : Int
}

///|
pub fn Rewrite::from_parts(
  name : String,
  searcher : Searcher,
  applier : Applier,
  condition? : (EGraph, Map[String, Id]) -> Bool = (_, _) => true,
) -> Rewrite {
  Rewrite::{ name, searcher, applier, condition }
}

///|
pub fn Rewrite::new(name : String, lhs : Pattern, rhs : Pattern) -> Rewrite {
  Rewrite::from_parts(name, Searcher::pattern(lhs), Applier::pattern(rhs))
}

///|
pub fn Rewrite::with_condition(
  name : String,
  lhs : Pattern,
  rhs : Pattern,
  condition : (EGraph, Map[String, Id]) -> Bool,
) -> Rewrite {
  Rewrite::from_parts(
    name,
    Searcher::pattern(lhs),
    Applier::pattern(rhs),
    condition~,
  )
}

///|
pub fn Rewrite::search(self : Rewrite, egraph : EGraph) -> Array[Match] {
  egraph.rebuild()
  self.searcher.search(egraph)
}

///|
fn Rewrite::apply_match(
  self : Rewrite,
  egraph : EGraph,
  m : Match,
  on_applied : (Match, Id) -> Unit,
) -> Bool raise {
  if !(self.condition)(egraph, m.subst) {
    return false
  }
  let rhs_id = self.applier.apply(egraph, m.subst)
  let root_before = egraph.find(m.root)
  let rhs_root = egraph.find(rhs_id)
  if root_before != rhs_root {
    ignore(egraph.union(root_before, rhs_root))
  }
  on_applied(m, rhs_id)
  true
}

///|
pub fn Rewrite::apply_all_stats(
  self : Rewrite,
  egraph : EGraph,
  on_applied? : (Match, Id) -> Unit = (_, _) => (),
) -> RewriteStats raise {
  let matches = self.search(egraph)
  let mut applied = 0
  for m in matches {
    if self.apply_match(egraph, m, on_applied) {
      applied += 1
    }
  }
  egraph.rebuild()
  RewriteStats::{ name: self.name, matched: matches.length(), applied }
}

///|
pub fn Rewrite::apply_all(self : Rewrite, egraph : EGraph) -> Int raise {
  self.apply_all_stats(egraph).applied
}

///|
/// Apply rewrite to matches that satisfy predicate (e.g., reachable from root).
pub fn Rewrite::apply_filtered(
  self : Rewrite,
  egraph : EGraph,
  pred : (Match) -> Bool,
  on_applied? : (Match, Id) -> Unit = (_, _) => (),
  root_first? : Id? = None,
  limit? : Int? = None,
) -> RewriteStats raise {
  let matches = self.search(egraph)
  let mut applied = 0
  let mut matched = 0
  let root_id = match root_first {
    Some(r) => Some(egraph.find(r))
    None => None
  }
  let filtered = matches.filter(m => pred(m))
  let mut ordered : Array[Match] = Array::new()
  match root_id {
    Some(rid) => {
      ordered.append(filtered.filter(m => egraph.find(m.root) == rid)[:])
      ordered.append(filtered.filter(m => egraph.find(m.root) != rid)[:])
    }
    None => ordered = filtered
  }
  for m in ordered {
    match limit {
      Some(lim) if applied >= lim => break
      _ => ()
    }
    matched = matched + 1
    if self.apply_match(egraph, m, on_applied) {
      applied = applied + 1
    }
  }
  egraph.rebuild()
  RewriteStats::{ name: self.name, matched, applied }
}