feuerwerk/src/core/model/hierarchy.ts

import { type Removed, removeWhere } from '../util/list';
import type { SceneObject } from './SceneObject';
import { descendantsBelow, parentChain } from './walk';

/**
 * `objects` stays a flat list in creation order; the tree is implied by `SceneObject.parentId`.
 * Anything holding such a list fits – a scene, or the template inside a container asset.
 */
export interface ObjectTree {
  objects: SceneObject[];
}

export function findObject(tree: ObjectTree, id: string): SceneObject | undefined {
  return tree.objects.find((o) => o.id === id);
}

function childrenOf(tree: ObjectTree, parentId: string): SceneObject[] {
  return tree.objects.filter((o) => o.parentId === parentId);
}

/** Children of `parentId` (`''` = top level) in sibling order – what a reorder rearranges. */
export function childIdsOf(tree: ObjectTree, parentId: string): string[] {
  return childrenOf(tree, parentId).map((o) => o.id);
}

function parentOf(tree: ObjectTree, object: SceneObject): SceneObject | undefined {
  return object.parentId ? findObject(tree, object.parentId) : undefined;
}

/** Depth-first, parents before children. */
export function descendantsOf(tree: ObjectTree, id: string): SceneObject[] {
  return descendantsBelow(
    id,
    (parentId) => childrenOf(tree, parentId),
    (o) => o.id,
  );
}

export function subtreeIds(tree: ObjectTree, rootIds: Iterable<string>): Set<string> {
  const ids = new Set<string>();
  for (const rootId of rootIds) {
    ids.add(rootId);
    for (const child of descendantsOf(tree, rootId)) ids.add(child.id);
  }
  return ids;
}

/** Returns the removed objects with their former indices (see `restoreAt`). */
export function removeSubtrees(tree: ObjectTree, rootIds: readonly string[]): Removed<SceneObject>[] {
  const ids = subtreeIds(tree, rootIds);
  return removeWhere(tree.objects, (o) => ids.has(o.id));
}

/** An object and every parent above it, nearest first. */
export function chainOf(tree: ObjectTree, object: SceneObject): SceneObject[] {
  return [...parentChain(object, (o) => parentOf(tree, o))];
}

/** Nearest parent first. */
export function ancestorIdsOf(tree: ObjectTree, id: string): string[] {
  const object = findObject(tree, id);
  return object
    ? chainOf(tree, object)
        .slice(1)
        .map((o) => o.id)
    : [];
}

export function isSelfOrAncestor(tree: ObjectTree, ancestorId: string, objectId: string): boolean {
  const object = findObject(tree, objectId);
  return object !== undefined && chainOf(tree, object).some((o) => o.id === ancestorId);
}

/** Of the given ids, those with no ancestor among the ids themselves. */
export function subtreeRoots(tree: ObjectTree, ids: readonly string[]): string[] {
  const set = new Set(ids);
  return ids.filter((id) => !ancestorIdsOf(tree, id).some((ancestorId) => set.has(ancestorId)));
}

export interface TreeEntry {
  object: SceneObject;
  depth: number;
}

/**
 * Top-level objects in storage order, each followed by its subtree – the order the hierarchy panel
 * shows and the tie-break order for equal `z`.
 */
export function treeOrder(tree: ObjectTree): TreeEntry[] {
  const result: TreeEntry[] = [];
  const ids = new Set(tree.objects.map((o) => o.id));
  const visit = (parentId: string, depth: number) => {
    for (const object of tree.objects) {
      if (object.parentId !== parentId) continue;
      result.push({ object, depth });
      visit(object.id, depth + 1);
    }
  };
  visit('', 0);
  // Objects whose parent is missing are listed at the top level rather than dropped.
  for (const object of tree.objects) {
    if (object.parentId && !ids.has(object.parentId) && !result.some((e) => e.object === object))
      result.push({ object, depth: 0 });
  }
  return result;
}

export interface VisibleEntry extends TreeEntry {
  hasChildren: boolean;
}

/**
 * Tree order with the subtrees of collapsed objects left out. Anything not in `expanded` counts as
 * collapsed, so a freshly loaded scene shows only its top level.
 */
export function visibleTreeOrder(tree: ObjectTree, expanded: ReadonlySet<string>): VisibleEntry[] {
  const parents = new Set(tree.objects.map((o) => o.parentId).filter(Boolean));
  const result: VisibleEntry[] = [];
  // Depth of the collapsed row whose subtree is being skipped (-1 = showing rows).
  let collapsedAt = -1;
  for (const entry of treeOrder(tree)) {
    if (collapsedAt >= 0) {
      if (entry.depth > collapsedAt) continue;
      collapsedAt = -1;
    }
    const hasChildren = parents.has(entry.object.id);
    result.push({ ...entry, hasChildren });
    if (hasChildren && !expanded.has(entry.object.id)) collapsedAt = entry.depth;
  }
  return result;
}