///|
/// Limits are checked before visiting a node. Root depth is zero; both limits
/// must be nonnegative. Iterative algorithms also terminate on cyclic AST arrays.
pub(all) enum TreeError {
  InvalidLimits(Int32, Int32)
  DepthLimitExceeded(Int32)
  NodeLimitExceeded(Int32)
  ChildCountMismatch(String, Int32, Int32)
} derive(Eq)

///|
pub(all) struct TreeVisit {
  node : Node
  depth : Int32
}

///|
/// Direct children in AST field order (including ids, comments and annotations).
/// Returns a fresh array, not an alias of an AST-owned array.
pub fn Node::children(self : Node) -> Array[Node] {
  let result : Array[Node] = []
  fn add_optional(n : Node?) {
    if n is Some(n) {
      result.push(n)
    }
  }
  match self.kind {
    Resource(nodes) | Pattern(nodes) => result.append(nodes)
    Message(id, value, attrs, comment) => {
      result.push(id)
      add_optional(value)
      result.append(attrs)
      add_optional(comment)
    }
    Term(id, value, attrs, comment) => {
      result.push(id)
      result.push(value)
      result.append(attrs)
      add_optional(comment)
    }
    Placeable(n) | VariableReference(n) => result.push(n)
    MessageReference(id, attribute) => {
      result.push(id)
      add_optional(attribute)
    }
    TermReference(id, attribute, arguments) => {
      result.push(id)
      add_optional(attribute)
      add_optional(arguments)
    }
    FunctionReference(id, args)
    | Attribute(id, args)
    | NamedArgument(id, args) => {
      result.push(id)
      result.push(args)
    }
    SelectExpression(selector, variants) => {
      result.push(selector)
      result.append(variants)
    }
    CallArguments(positional, named) => {
      result.append(positional)
      result.append(named)
    }
    Variant(key, value, _) => {
      result.push(key)
      result.push(value)
    }
    Junk(_, annotations) => result.append(annotations)
    _ => ()
  }
  result
}

///|
/// Replace all direct children without changing constructor, optional-field
/// presence, array lengths, scalar payloads or span. This is a structural
/// primitive, not a Fluent grammar validator. Even leaves require zero children.
pub fn Node::with_children(
  self : Node,
  children : Array[Node],
) -> Result[Node, TreeError] {
  let expected = self.children().length()
  if expected != children.length() {
    return Err(
      ChildCountMismatch(self.type_name(), expected, children.length()),
    )
  }
  let mut cursor = 0
  fn next() -> Node {
    let n = children[cursor]
    cursor = cursor + 1
    n
  }
  fn optional(previous : Node?) -> Node? {
    match previous {
      Some(_) => Some(next())
      None => None
    }
  }
  fn many(previous : Array[Node]) -> Array[Node] {
    previous.map(_ => next())
  }
  let kind = match self.kind {
    Resource(ns) => Resource(many(ns))
    Message(id, value, attrs, comment) => {
      ignore(id)
      Message(next(), optional(value), many(attrs), optional(comment))
    }
    Term(_, _, attrs, comment) =>
      Term(next(), next(), many(attrs), optional(comment))
    Pattern(ns) => Pattern(many(ns))
    Placeable(_) => Placeable(next())
    MessageReference(_, attr) => MessageReference(next(), optional(attr))
    TermReference(_, attr, args) =>
      TermReference(next(), optional(attr), optional(args))
    VariableReference(_) => VariableReference(next())
    FunctionReference(_, _) => FunctionReference(next(), next())
    SelectExpression(_, variants) => SelectExpression(next(), many(variants))
    CallArguments(positional, named) =>
      CallArguments(many(positional), many(named))
    Attribute(_, _) => Attribute(next(), next())
    Variant(_, _, default) => Variant(next(), next(), default)
    NamedArgument(_, _) => NamedArgument(next(), next())
    Junk(text, annotations) => Junk(text, many(annotations))
    other => other
  }
  Ok({ kind, span: self.span, })
}

///|
/// Preorder enumeration. Shared subtrees are visited once per occurrence, not
/// deduplicated. No recursive call stack is used. Failure returns no partial list.
pub fn Node::walk(
  self : Node,
  max_depth~ : Int32 = 256,
  max_nodes~ : Int32 = 100000,
) -> Result[Array[TreeVisit], TreeError] {
  if max_depth < 0 || max_nodes < 0 {
    return Err(InvalidLimits(max_depth, max_nodes))
  }
  let pending : Array[TreeVisit] = [{ node: self, depth: 0, }]
  let visits : Array[TreeVisit] = []
  while !pending.is_empty() {
    let visit = pending.pop().unwrap()
    if visit.depth > max_depth {
      return Err(DepthLimitExceeded(max_depth))
    }
    if visits.length() >= max_nodes {
      return Err(NodeLimitExceeded(max_nodes))
    }
    visits.push(visit)
    let children = visit.node.children()
    // Bound the pending queue as well as the output, including very wide trees.
    if children.length() > max_nodes - visits.length() - pending.length() {
      return Err(NodeLimitExceeded(max_nodes))
    }
    for i = children.length() - 1; i >= 0; i = i - 1 {
      pending.push({ node: children[i], depth: visit.depth + 1, })
    }
  }
  Ok(visits)
}

///|
/// Bottom-up, left-to-right rewrite with original spans retained, even when the
/// callback supplies a synthetic span. Newly returned subtrees are not revisited.
/// The callback must not mutate input AST arrays. Limits are validated for the
/// entire input before invoking any callback, so traversal errors have no effects.
pub fn Node::transform(
  self : Node,
  rewrite : (Node) -> Node,
  max_depth~ : Int32 = 256,
  max_nodes~ : Int32 = 100000,
) -> Result[Node, TreeError] {
  let visits = match self.walk(max_depth~, max_nodes~) {
    Ok(visits) => visits
    Err(error) => return Err(error)
  }
  // An explicit continuation stack preserves left-to-right callback order.
  let pending : Array[(Node, Bool)] = [(self, false)]
  let output : Array[Node] = []
  ignore(visits)
  while !pending.is_empty() {
    let (node, ready) = pending.pop().unwrap()
    let children = node.children()
    if !ready {
      pending.push((node, true))
      for i = children.length() - 1; i >= 0; i = i - 1 {
        pending.push((children[i], false))
      }
    } else {
      let start = output.length() - children.length()
      let rebuilt_children = children.mapi((i, _) => output[start + i])
      for _ in children {
        ignore(output.pop())
      }
      let rebuilt = match node.with_children(rebuilt_children) {
        Ok(n) => n
        Err(error) => return Err(error)
      }
      let changed = rewrite(rebuilt)
      output.push({ kind: changed.kind, span: node.span, })
    }
  }
  Ok(output.pop().unwrap())
}

///|
pub extend TreeError with Eq::{equal, not_equal}