///|
/// Key for pending pattern nodes during compilation.
struct TodoKey {
  node : Int
  reg : Int
} derive(Eq, Hash)

///|
/// Flattened pattern nodes.
priv enum FlatNode {
  Var(String)
  VarIf(String, (Id) -> Bool)
  Sym(String)
  Num(Float)
  Node(String, Array[Int])
}

///|
priv struct FlatPattern {
  nodes : Array[FlatNode]
  root : Int
}

///|
pub enum ENodeOrReg {
  Node(ENode)
  Reg(Int)
}

///|
/// Instruction program for the backtracking matcher.
pub enum Instruction {
  Bind(NodeOp, Int, Int, Int) // op, arity, src_reg, out_reg
  Compare(Int, Int)
  Lookup(Array[ENodeOrReg], Int) // term, target reg
  Scan(Int) // out reg
}

///|
pub struct Program {
  instructions : Array[Instruction]
  subst : Map[String, Int]
  var_if : Map[String, (Id) -> Bool]
}

///|
priv struct Machine {
  reg : Array[Id]
  mut lookup : Array[Id]
}

///|
fn Machine::new() -> Machine {
  Machine::{ reg: Array::new(), lookup: Array::new() }
}

///|
fn Machine::trim(self : Machine, len : Int) -> Unit {
  while self.reg.length() > len {
    let _ = self.reg.pop()
  }
}

///|
fn flatten_pattern(pat : Pattern) -> FlatPattern {
  let nodes : Array[FlatNode] = Array::new()
  fn go(
    pat : Pattern,
    nodes : Array[FlatNode],
    counter : Int,
  ) -> (Int, Int) {
    match pat {
      Pattern::Wild => {
        let name = "_wild\{counter}"
        nodes.push(FlatNode::Var(name))
        (nodes.length() - 1, counter + 1)
      }
      Pattern::Var(name) => {
        nodes.push(FlatNode::Var(name))
        (nodes.length() - 1, counter)
      }
      Pattern::VarIf(name, pred) => {
        nodes.push(FlatNode::VarIf(name, pred))
        (nodes.length() - 1, counter)
      }
      Pattern::Sym(sym) => {
        nodes.push(FlatNode::Sym(sym))
        (nodes.length() - 1, counter)
      }
      Pattern::Num(n) => {
        nodes.push(FlatNode::Num(n))
        (nodes.length() - 1, counter)
      }
      Pattern::Node(op, children) => {
        let mut next = counter
        let child_ids : Array[Int] = Array::new()
        for child in children {
          let (cid, updated) = go(child, nodes, next)
          next = updated
          child_ids.push(cid)
        }
        nodes.push(FlatNode::Node(op, child_ids))
        (nodes.length() - 1, next)
      }
    }
  }
  let (root, _) = go(pat, nodes, 0)
  FlatPattern::{ nodes, root }
}

///|
fn compute_info(
  flat : FlatPattern,
) -> (Array[Map[String, Bool]], Array[Int]) {
  let free_vars : Array[Map[String, Bool]] = Array::new()
  let sizes : Array[Int] = Array::new()
  for _ in 0.. {
        let m = Map::new()
        m.set(name, true)
        free_vars[idx] = m
        sizes[idx] = 0
      }
      FlatNode::VarIf(name, _) => {
        let m = Map::new()
        m.set(name, true)
        free_vars[idx] = m
        sizes[idx] = 0
      }
      FlatNode::Sym(_) | FlatNode::Num(_) => {
        free_vars[idx] = Map::new()
        sizes[idx] = 1
      }
      FlatNode::Node(_, children) => {
        let m = Map::new()
        let mut size = 1
        for child in children {
          for pair in free_vars[child].iter() {
            let (k, v) = pair
            if v {
              m.set(k, true)
            }
          }
          size = size + sizes[child]
        }
        free_vars[idx] = m
        sizes[idx] = size
      }
    }
  }
  (free_vars, sizes)
}

///|
fn build_lookup_term(
  flat : FlatPattern,
  root : Int,
  v2r : Map[String, Int],
) -> Array[ENodeOrReg] {
  let term : Array[ENodeOrReg] = Array::new()
  let pos : Map[Int, Int] = Map::new()
  fn visit(
    flat : FlatPattern,
    id : Int,
    v2r : Map[String, Int],
    term : Array[ENodeOrReg],
    pos : Map[Int, Int],
  ) -> Unit {
    match pos.get(id) {
      Some(_) => return
      None => ()
    }
    let node = flat.nodes[id]
    match node {
      FlatNode::Node(_, children) =>
        for child in children {
          visit(flat, child, v2r, term, pos)
        }
      _ => ()
    }
    let entry = match node {
      FlatNode::Var(name) => ENodeOrReg::Reg(v2r.get(name).unwrap())
      FlatNode::VarIf(name, _) => ENodeOrReg::Reg(v2r.get(name).unwrap())
      FlatNode::Sym(sym) =>
        ENodeOrReg::Node(ENode::{ op: NodeOp::Symbol(sym), children: [] })
      FlatNode::Num(n) =>
        ENodeOrReg::Node(ENode::{ op: NodeOp::Number(n), children: [] })
      FlatNode::Node(op, children) => {
        let term_children = children.map(child => pos.get(child).unwrap())
        ENodeOrReg::Node(ENode::{
          op: NodeOp::Name(op),
          children: term_children,
        })
      }
    }
    term.push(entry)
    pos.set(id, term.length() - 1)
  }
  visit(flat, root, v2r, term, pos)
  term
}

///|
priv struct Compiler {
  v2r : Map[String, Int]
  mut free_vars : Array[Map[String, Bool]]
  mut subtree_size : Array[Int]
  todo_nodes : Map[TodoKey, Int]
  instructions : Array[Instruction]
  mut next_reg : Int
  var_if : Map[String, (Id) -> Bool]
}

///|
fn Compiler::new() -> Compiler {
  Compiler::{
    v2r: Map::new(),
    free_vars: Array::new(),
    subtree_size: Array::new(),
    todo_nodes: Map::new(),
    instructions: Array::new(),
    next_reg: 0,
    var_if: Map::new(),
  }
}

///|
fn Compiler::add_todo(self : Compiler, flat : FlatPattern, id : Int, reg : Int) -> Unit {
  let node = flat.nodes[id]
  match node {
    FlatNode::Var(name) => match self.v2r.get(name) {
      Some(existing) => self.instructions.push(Instruction::Compare(reg, existing))
      None => self.v2r.set(name, reg)
    }
    FlatNode::VarIf(name, pred) => {
      if !self.var_if.contains(name) {
        self.var_if.set(name, pred)
      }
      match self.v2r.get(name) {
        Some(existing) =>
          self.instructions.push(Instruction::Compare(reg, existing))
        None => self.v2r.set(name, reg)
      }
    }
    FlatNode::Sym(_) | FlatNode::Num(_) | FlatNode::Node(_, _) => {
      self.todo_nodes.set(TodoKey::{ node: id, reg }, id)
    }
  }
}

///|
fn Compiler::load_pattern(self : Compiler, flat : FlatPattern) -> Unit {
  let (free_vars, sizes) = compute_info(flat)
  self.free_vars = free_vars
  self.subtree_size = sizes
}

///|
fn Compiler::is_ground_now(self : Compiler, id : Int) -> Bool {
  let free = self.free_vars[id]
  for pair in free.iter() {
    let (name, present) = pair
    if present && !self.v2r.contains(name) {
      return false
    }
  }
  true
}

///|
fn Compiler::next(self : Compiler, flat : FlatPattern) -> (TodoKey, FlatNode)? {
  let mut best_key : (Bool, Int, Int)? = None
  let mut best : (TodoKey, FlatNode)? = None
  for pair in self.todo_nodes.iter() {
    let (k, _) = pair
    let idx = k.node
    let free = self.free_vars[idx]
    let mut n_bound = 0
    for fv in free.iter() {
      let (name, present) = fv
      if present && self.v2r.contains(name) {
        n_bound = n_bound + 1
      }
    }
    let n_free = free.length() - n_bound
    let size = self.subtree_size[idx]
    let key = (n_free == 0, n_free, -size)
    match best_key {
      None => {
        best_key = Some(key)
        best = Some((k, flat.nodes[idx]))
      }
      Some(existing) =>
        if key > existing {
          best_key = Some(key)
          best = Some((k, flat.nodes[idx]))
        }
    }
  }
  match best {
    Some((k, node)) => {
      self.todo_nodes.remove(k)
      Some((k, node))
    }
    None => None
  }
}

///|
fn Compiler::compile(
  self : Compiler,
  flat : FlatPattern,
  binder? : String? = None,
) -> Unit {
  self.load_pattern(flat)
  let root = flat.root
  let mut next_out = self.next_reg

  if binder is Some(name) {
    match self.v2r.get(name) {
      Some(existing) => self.add_todo(flat, root, existing)
      None => {
        next_out = next_out + 1
        if !self.instructions.is_empty() {
          self.instructions.push(Instruction::Scan(self.next_reg))
        }
        self.add_todo(flat, root, self.next_reg)
        self.v2r.set(name, self.next_reg)
      }
    }
  } else {
    next_out = next_out + 1
    if !self.instructions.is_empty() {
      self.instructions.push(Instruction::Scan(self.next_reg))
    }
    self.add_todo(flat, root, self.next_reg)
  }

  loop () {
    _ =>
      match self.next(flat) {
        None => break ()
        Some((key, node)) => {
          let reg = key.reg
          let (op, children) = match node {
            FlatNode::Sym(sym) => (NodeOp::Symbol(sym), Array::new())
            FlatNode::Num(n) => (NodeOp::Number(n), Array::new())
            FlatNode::Node(op, kids) => (NodeOp::Name(op), kids)
            _ => continue ()
          }
          if self.is_ground_now(key.node) && !children.is_empty() {
            let term = build_lookup_term(flat, key.node, self.v2r)
            self.instructions.push(Instruction::Lookup(term, reg))
            continue ()
          }
          let out = next_out
          next_out = next_out + children.length()
          self.instructions.push(Instruction::Bind(op, children.length(), reg, out))
          for idx in 0.. Program {
  let subst = Map::new()
  for pair in self.v2r.iter() {
    let (name, reg) = pair
    subst.set(name, reg)
  }
  Program::{ instructions: self.instructions, subst, var_if: self.var_if }
}

///|
pub fn Program::compile_from_pattern(pat : Pattern) -> Program {
  let flat = flatten_pattern(pat)
  let compiler = Compiler::new()
  compiler.compile(flat, binder=None)
  compiler.extract()
}

///|
pub fn Program::compile_from_multi(patterns : Array[(String, Pattern)]) -> Program {
  let compiler = Compiler::new()
  for pair in patterns {
    let (binder, pat) = pair
    let flat = flatten_pattern(pat)
    compiler.compile(flat, binder=Some(binder))
  }
  compiler.extract()
}

///|
fn exec(
  machine : Machine,
  program : Program,
  egraph : EGraph,
  idx : Int,
  results : Array[Map[String, Id]],
  limit : Int?,
) -> Bool {
  match limit {
    Some(lim) if results.length() >= lim => return false
    _ => ()
  }
  if idx >= program.instructions.length() {
    if !egraph.allow_cycles() {
      if machine.reg.length() > 1 {
        let first = egraph.find(machine.reg[0])
        for i in 1..
          if !(pred)(id) {
            return true
          }
        None => return true
      }
    }
    results.push(subst)
    return match limit {
      Some(lim) => results.length() < lim
      None => true
    }
  }
  let inst = program.instructions[idx]
  match inst {
    Instruction::Compare(a, b) => {
      if egraph.find(machine.reg[a]) != egraph.find(machine.reg[b]) {
        return true
      }
      exec(machine, program, egraph, idx + 1, results, limit)
    }
    Instruction::Bind(op, arity, src, out) => {
      let class_id = egraph.find(machine.reg[src])
      match egraph.class_for(class_id) {
        None => true
        Some(class) => {
          let saved = machine.reg.length()
          for node_idx in class.nodes {
            let node = egraph.nodes[node_idx]
            if node.op != op {
              continue
            }
            if node.children.length() != arity {
              continue
            }
            machine.trim(out)
            for child in node.children {
              machine.reg.push(child)
            }
            if !exec(
              machine, program, egraph, idx + 1, results, limit,
            ) {
              return false
            }
            machine.trim(saved)
          }
          true
        }
      }
    }
    Instruction::Scan(out) => {
      let saved = machine.reg.length()
      for class_id in egraph.class_ids() {
        machine.trim(out)
        machine.reg.push(class_id)
        if !exec(
          machine, program, egraph, idx + 1, results, limit,
        ) {
          return false
        }
        machine.trim(saved)
      }
      true
    }
    Instruction::Lookup(term, target_reg) => {
      machine.lookup = Array::new()
      let mut ok = true
      for entry in term {
        match entry {
          ENodeOrReg::Reg(r) => machine.lookup.push(egraph.find(machine.reg[r]))
          ENodeOrReg::Node(template) => {
            let children_ids = template.children.map(idx => machine.lookup[idx])
            let enode = ENode::{ op: template.op, children: children_ids }
            match egraph.lookup(enode) {
              Some(id) => machine.lookup.push(id)
              None => {
                ok = false
                break
              }
            }
          }
        }
      }
      if !ok {
        return true
      }
      let target = egraph.find(machine.reg[target_reg])
      match machine.lookup.last() {
        Some(id) if id == target =>
          exec(machine, program, egraph, idx + 1, results, limit)
        _ => true
      }
    }
  }
}

///|
pub fn Program::run_with_limit(
  self : Program,
  egraph : EGraph,
  eclass : Id,
  limit? : Int? = None,
) -> Array[Map[String, Id]] {
  let results : Array[Map[String, Id]] = Array::new()
  match limit {
    Some(lim) if lim <= 0 => return results
    _ => ()
  }
  let machine = Machine::new()
  machine.reg.push(egraph.find(eclass))
  ignore(exec(machine, self, egraph, 0, results, limit))
  results
}