import { ASTRootBase } from './nodes';

export function* iterChildNodes(node: ASTRootBase): IterableIterator<ASTRootBase> {
  const type = Object.getPrototypeOf(node).constructor;
  for (let attributeName of type._attributes) {
    // @ts-ignore
    const attribute = node[attributeName];

    if (attribute instanceof Array) {
      for (let arrayItem of attribute) {
        if (arrayItem instanceof ASTRootBase) {
          yield arrayItem;
        }
      }
    } else if (attribute instanceof ASTRootBase) {
      yield attribute;
    }
  }
}

export function* walk(node: ASTRootBase): IterableIterator<ASTRootBase> {
  const todo = [node];
  while (todo.length > 0) {
    const current = todo.shift()!;
    todo.push(...iterChildNodes(current));
    yield current;
  }
}

export class NodeVisitor {
  visit(node: ASTRootBase): any {
    const name = node.constructor.name;
    // @ts-ignore
    const handler = this[`visit${name}`];
    if (handler === undefined) {
      return this.genericVisit(node);
    } else {
      return handler(node);
    }
  }

  genericVisit(node: ASTRootBase): void {
    for (let attribute of iterChildNodes(node)) {
      this.visit(attribute);
    }
  }
}