///|
pub fn default_cost(op : NodeOp, child_costs : Array[Int]) -> Int {
  ignore(op)
  let base = child_costs.fold(init=0, (acc, cost) => acc + cost)
  1 + base
}

///|
pub fn extract_best(
  egraph : EGraph,
  root : Id,
  cost_fn? : (NodeOp, Array[Int]) -> Int = default_cost,
) -> (Int, Expr) raise {
  egraph.rebuild()
  let best_cost : Map[Id, Int] = Map::new()
  let best_node : Map[Id, ENode] = Map::new()

  // Greedily propagate costs until fixed point, mirroring egg's Extractor.
  loop () {
    _ => {
      let mut improved = false
      for pair in egraph.classes.iter() {
        let (id, class) = pair
        let class_id = egraph.find_read(id)
        if class_id != id {
          continue
        }
        let mut candidate_cost : Int? = None
        let mut candidate_node : ENode? = None
        for node_index in class.nodes {
          let node = egraph.nodes[node_index]
          let child_costs : Array[Int] = Array::new()
          let mut ready = true
          for child in node.children {
            let child_root = egraph.find_read(child)
            match best_cost.get(child_root) {
              Some(cost) => child_costs.push(cost)
              None => {
                ready = false
                break
              }
            }
          }
          if !ready {
            if !node.children.is_empty() {
              continue
            }
          }
          let cost = cost_fn(node.op, child_costs)
          match candidate_cost {
            None => {
              candidate_cost = Some(cost)
              candidate_node = Some(node)
            }
            Some(existing) =>
              if cost < existing {
                candidate_cost = Some(cost)
                candidate_node = Some(node)
              }
          }
        }
        match (candidate_cost, candidate_node) {
          (Some(cost), Some(node)) =>
            match best_cost.get(class_id) {
              Some(existing) =>
                if cost < existing {
                  best_cost.set(class_id, cost)
                  best_node.set(class_id, node)
                  improved = true
                }
              None => {
                best_cost.set(class_id, cost)
                best_node.set(class_id, node)
                improved = true
              }
            }
          _ => ()
        }
      }
      if !improved {
        break ()
      }
      continue ()
    }
  }
  let cache : Map[Id, Expr] = Map::new()
  fn build_expr(
    cache : Map[Id, Expr],
    best_node : Map[Id, ENode],
    egraph : EGraph,
    id : Id,
  ) -> (Map[Id, Expr], Expr) raise {
    let root = egraph.find_read(id)
    match cache.get(root) {
      Some(expr) => (cache, expr)
      None => {
        let node = best_node.get(root).unwrap()
        let mut next_cache = cache
        let built_children : Array[Expr] = Array::new()
        for child in node.children {
          let (updated, expr) = build_expr(next_cache, best_node, egraph, child)
          next_cache = updated
          built_children.push(expr)
        }
        let expr = match node.op {
          NodeOp::Name(name) => Expr::Node(name, built_children)
          NodeOp::Symbol(sym) => Expr::Leaf(sym)
          NodeOp::Number(n) => Expr::Leaf(n.to_string())
        }
        next_cache.set(root, expr)
        (next_cache, next_cache.get(root).unwrap())
      }
    }
  }

  let root_id = egraph.find_read(root)
  let (_, best_expr) = build_expr(cache, best_node, egraph, root_id)
  let cost = best_cost.get(root_id).unwrap()
  (cost, best_expr)
}