ecs-observable/src/bt/tree-def.ts

import type { World, Entity } from "../index";
import type { EntityDef, EntityDefChild } from "../entity-tree";
import { Task, ChildOf } from "./task";
import { TaskRunner } from "./runner";

// ── Cancel ────────────────────────────────────────────

/**
 * Throw this inside a leaf `run` function to cancel the leaf and its subtree.
 *
 * @example
 * ```ts
 * leaf(() => { throw Cancel; })
 * ```
 */
export const Cancel: unique symbol = Symbol("leaf.cancel");

// ── Tree definition ───────────────────────────────────

/** A leaf function — plain or generator. */
export type LeafFn =
  | ((world: World, dt: number) => void)
  | (() => Generator<number | void, void, number>);

export interface LeafTaskMeta {
  readonly run: LeafFn;
}

export type TaskEntityDef = EntityDef<
  typeof Task.type,
  LeafTaskMeta | undefined
>;

/** Behaviour-tree definition accepted by `buildTree`. */
export type TreeDef = TaskEntityDef;

// ── Entity task factories ──────────────────────────────

function compactChildren(children: readonly EntityDefChild[]): EntityDef[] {
  return children.filter(Boolean) as EntityDef[];
}

function task(
  kind: typeof Task.type.kind,
  children: readonly EntityDefChild[] = [],
  meta?: LeafTaskMeta,
): TaskEntityDef {
  return {
    kind: "entity",
    component: Task,
    value: { kind },
    children: compactChildren(children),
    meta,
  };
}

/** Create a leaf task entity definition. */
export function leaf(
  run: LeafFn,
  children: readonly EntityDefChild[] = [],
): TaskEntityDef {
  return task("leaf", children, { run });
}

/** Create a sequential task entity definition. */
export function sequential(
  children: readonly EntityDefChild[] = [],
): TaskEntityDef {
  return task("sequential", children);
}

/** Create a parallel task entity definition. */
export function parallel(
  children: readonly EntityDefChild[] = [],
): TaskEntityDef {
  return task("parallel", children);
}

/** Create a selector task entity definition. */
export function selector(
  children: readonly EntityDefChild[] = [],
): TaskEntityDef {
  return task("selector", children);
}

/** Create a random task entity definition. */
export function random(
  children: readonly EntityDefChild[] = [],
): TaskEntityDef {
  return task("random", children);
}

/**
 * Create a repeat task entity definition.
 *
 * `repeat(child)` and `repeat([child, metadataEntity])` are both supported.
 * The runner operates only on child entities that have the `Task` component.
 */
export function repeat(child: EntityDef): TaskEntityDef;
export function repeat(children?: readonly EntityDefChild[]): TaskEntityDef;
export function repeat(
  childOrChildren: EntityDef | readonly EntityDefChild[] = [],
): TaskEntityDef {
  return task(
    "repeat",
    Array.isArray(childOrChildren) ? childOrChildren : [childOrChildren],
  );
}

// ── Builder ───────────────────────────────────────────

/**
 * Materialize a behaviour-tree definition into ECS entities and return a
 * fully-wired `TaskRunner`.
 *
 * Definitions are `EntityDef` trees produced by the task factories (`leaf`,
 * `sequential`, `parallel`, `selector`, `random`, `repeat`) and generic
 * single-component entity factories. Non-task child entities are materialized
 * into the ECS tree but ignored by `TaskRunner` execution.
 *
 * Leaf `run` functions:
 * - **Plain function** — runs once per tick. `return` = success. `throw` = fail.
 *   `throw Cancel` = cancel.
 * - **Generator function** — each `yield` suspends until next tick. The value
 *   yielded is the desired delay in ms (or `undefined` for next frame).
 *   Generator completion = success. `throw` = fail. `throw Cancel` = cancel.
 */
export function buildTree(world: World, def: TreeDef): TaskRunner {
  const leafHandlers = new Map<Entity, LeafFn>();
  // Track generator iterators for multi-frame leaves
  const generators = new Map<Entity, Generator<number | void, void, number>>();

  function build(def: EntityDef, parent?: Entity): Entity {
    const entity = world.spawn();
    world.add(entity, def.component, def.value);

    if (parent !== undefined) {
      world.relate(entity, ChildOf, parent);
    }

    if (def.component === Task) {
      const taskData = world.get(entity, Task);
      if (taskData.kind === "leaf") {
        const run = (def.meta as LeafTaskMeta | undefined)?.run;
        if (!run) {
          throw new Error("Leaf task entity is missing a run function");
        }
        leafHandlers.set(entity, run);
      }
    }

    for (const child of def.children) {
      build(child, entity);
    }

    return entity;
  }

  const root = build(def);

  if (!world.has(root, Task)) {
    throw new Error("buildTree root must be a task entity");
  }

  const runner = new TaskRunner(world);
  runner.root = root;

  runner.onLeaf = (_w, entity, dt) => {
    const handler = leafHandlers.get(entity);
    if (!handler) return;

    try {
      // Check if this leaf has an active generator
      let gen = generators.get(entity);

      if (gen) {
        // Resume existing generator
        const result = gen.next(dt);
        if (result.done) {
          generators.delete(entity);
          runner.succeed(entity);
        }
        // If not done, leaf stays Running — nothing to do
      } else {
        // First invocation — call the handler
        const ret = handler(_w, dt);

        // Check if it returned a generator
        if (ret != null && typeof (ret as any).next === "function") {
          const gen = ret as Generator<number | void, void, number>;
          generators.set(entity, gen);
          const result = gen.next(dt);
          if (result.done) {
            generators.delete(entity);
            runner.succeed(entity);
          }
          // Not done → leaf stays Running
        } else {
          // Plain function — returned undefined → success
          runner.succeed(entity);
        }
      }
    } catch (err) {
      // Clean up generator if one was active
      generators.delete(entity);

      if (err === Cancel) {
        runner.cancel(entity);
      } else {
        runner.fail(entity);
      }
    }
  };

  runner.onTerminal = (_w, entity) => {
    // Clean up generator when a leaf reaches terminal by external means
    generators.delete(entity);
  };

  return runner;
}