ujsx/src/host/config.ts

import type { UNode, UElement, URoot, ComponentFn, UComponent } from "../core/schema.js";
import { isURoot, isUPrimitive } from "../core/schema.js";
import { Context } from "../core/context.js";
import type { Fiber } from "./fiber.js";
import { disposeFiber } from "./fiber.js";
import { reconcileProps, reconcileChildren, commitMutations } from "./reconcile.js";
import type { CommitContext } from "./reconcile.js";

export interface HostConfig<TTag extends string, Instance, RootCtx> {
  name: string;
  createRootContext(container: unknown, options?: Record<string, unknown>, context?: Context): RootCtx;
  finalizeRoot?(ctx: RootCtx): void;
  createInstance(tag: TTag, props: Record<string, unknown>, ctx: RootCtx, parent?: Instance): Instance;
  createTextInstance(text: string, ctx: RootCtx, parent?: Instance): Instance;
  appendChild(parent: Instance, child: Instance, ctx: RootCtx): void;
  insertBefore?(parent: Instance, child: Instance, before: Instance, ctx: RootCtx): void;
  removeChild?(parent: Instance, child: Instance, ctx: RootCtx): void;
  prepareUpdate?(
    instance: Instance,
    tag: TTag,
    prevProps: Record<string, unknown>,
    nextProps: Record<string, unknown>,
    ctx: RootCtx,
  ): unknown | null;
  commitUpdate?(
    instance: Instance,
    payload: unknown,
    tag: TTag,
    prevProps: Record<string, unknown>,
    nextProps: Record<string, unknown>,
    ctx: RootCtx,
  ): void;
  emit?(type: string, id: string, payload: unknown): void;
  finalizeInstance?(instance: Instance, ctx: RootCtx): void;
}

export interface Root<TTag extends string, Instance, RootCtx> {
  host: HostConfig<TTag, Instance, RootCtx>;
  ctx: RootCtx;
  container: unknown;
  context: Context;
  rootFiber: Fiber<Instance> | null;
  render(node: UNode): void;
  unmount(): void;
}

export function createRoot<TTag extends string, Instance, RootCtx>(
  host: HostConfig<TTag, Instance, RootCtx>,
  container: unknown,
  options?: Record<string, unknown>,
  context?: Context,
): Root<TTag, Instance, RootCtx> {
  const ctx = host.createRootContext(container, options, context);
  const rootContext = context ?? new Context();

  function mountNode(node: UNode, parentFiber: Fiber<Instance> | null): Fiber<Instance> | undefined {
    if (node == null || node === false) return undefined;

    if (isUPrimitive(node)) {
      const text = node === null ? "" : String(node);
      const parentInst = parentFiber?.instance;
      const t = host.createTextInstance(text, ctx, parentInst);
      if (parentInst) host.appendChild(parentInst, t, ctx);
      host.emit?.("instance.create", `text_${Date.now()}`, { kind: "text", value: text });
      const fiber: Fiber<Instance> = {
        instance: t,
        tag: "#text",
        props: { text },
        key: undefined,
        children: [],
        parent: parentFiber,
        effect: null,
        signalDisposers: [],
        prevProps: null,
        cachedNode: node,
        disposed: false,
        hash: null,
      };
      if (parentFiber) parentFiber.children.push(fiber);
      return fiber;
    }

    if (isURoot(node)) {
      for (const child of node.children) {
        mountNode(child, parentFiber);
      }
      return parentFiber ?? undefined;
    }

    const el = node as UElement;

    if (typeof el.type === "function") {
      const component = el.type as ComponentFn;
      const out = component({ ...el.props, children: el.children });
      host.emit?.("component.invoke", `comp_${Date.now()}`, { type: (component as unknown as UComponent).displayName ?? "anonymous" });
      return mountNode(out, parentFiber);
    }

    const tag = el.type as TTag;
    const parentInst = parentFiber?.instance;
    const inst = host.createInstance(tag, el.props as Record<string, unknown>, ctx, parentInst);
    host.emit?.("instance.create", `${tag}_${Date.now()}`, { kind: "element", tag, props: el.props });

    const fiber: Fiber<Instance> = {
      instance: inst,
      tag,
      props: el.props as Record<string, unknown>,
      key: el.key,
      children: [],
      parent: parentFiber,
      effect: null,
      signalDisposers: [],
      prevProps: null,
      cachedNode: node,
      disposed: false,
      hash: null,
    };

    for (const child of el.children) {
      mountNode(child, fiber);
    }

    if (parentInst) host.appendChild(parentInst, inst, ctx);
    if (parentFiber) parentFiber.children.push(fiber);
    return fiber;
  }

  function createFiberForInsert(node: UNode, parentFiber: Fiber<Instance>): Fiber<Instance> {
    if (isUPrimitive(node)) {
      const text = node === null ? "" : String(node);
      const t = host.createTextInstance(text, ctx, parentFiber.instance);
      host.emit?.("instance.create", `text_${Date.now()}`, { kind: "text", value: text });
      return {
        instance: t,
        tag: "#text",
        props: { text },
        key: undefined,
        children: [],
        parent: parentFiber,
        effect: null,
        signalDisposers: [],
        prevProps: null,
        cachedNode: node,
        disposed: false,
        hash: null,
      };
    }

    const el = node as UElement;

    if (typeof el.type === "function") {
      const component = el.type as ComponentFn;
      const out = component({ ...el.props, children: el.children });
      host.emit?.("component.invoke", `comp_${Date.now()}`, { type: (component as unknown as UComponent).displayName ?? "anonymous" });
      return createFiberForInsert(out, parentFiber);
    }

    const tag = el.type as TTag;
    const inst = host.createInstance(tag, el.props as Record<string, unknown>, ctx, parentFiber.instance);
    host.emit?.("instance.create", `${tag}_${Date.now()}`, { kind: "element", tag, props: el.props });

    const fiber: Fiber<Instance> = {
      instance: inst,
      tag,
      props: el.props as Record<string, unknown>,
      key: el.key,
      children: [],
      parent: parentFiber,
      effect: null,
      signalDisposers: [],
      prevProps: null,
      cachedNode: node,
      disposed: false,
      hash: null,
    };

    for (const child of el.children) {
      const childFiber = createFiberForInsert(child, fiber);
      if (childFiber) {
        host.appendChild(inst, childFiber.instance, ctx);
        fiber.children.push(childFiber);
      }
    }

    return fiber;
  }

  function resolveUNode(node: UNode): UNode[] {
    if (node == null || node === false) return [];
    if (isUPrimitive(node)) return [node];
    if (isURoot(node)) {
      const result: UNode[] = [];
      for (const child of (node as URoot).children) {
        result.push(...resolveUNode(child));
      }
      return result;
    }
    const el = node as UElement;
    if (typeof el.type === "function") {
      const component = el.type as ComponentFn;
      const out = component({ ...el.props, children: el.children });
      return resolveUNode(out);
    }
    return [node];
  }

  function reconcileNode(fiber: Fiber<Instance>, nextNode: UNode): void {
    if (isUPrimitive(nextNode)) {
      if (fiber.tag === "#text") {
        reconcileProps(fiber, nextNode, host as HostConfig<string, Instance, unknown>, ctx as unknown);
      }
      return;
    }

    if (isURoot(nextNode)) {
      const rootChildren = (nextNode as URoot).children;
      const classification = reconcileChildren(
        fiber.children,
        rootChildren,
      );
      for (const m of classification.matched) {
        reconcileProps(m.oldFiber, m.newChild, host as HostConfig<string, Instance, unknown>, ctx as unknown);
      }
      const commitCtx: CommitContext<Instance> = {
        host: host as HostConfig<string, Instance, unknown>,
        ctx: ctx as unknown,
        createFiber: (node, parent) => createFiberForInsert(node, parent),
      };
      commitMutations(fiber, classification, commitCtx);
      return;
    }

    const el = nextNode as UElement;

    if (typeof el.type === "function") {
      const component = el.type as ComponentFn;
      const out = component({ ...el.props, children: el.children });
      reconcileNode(fiber, out);
      return;
    }

    if (fiber.tag !== el.type) return;

    reconcileProps(fiber, el, host as HostConfig<string, Instance, unknown>, ctx as unknown);

    const classification = reconcileChildren(fiber.children, el.children);
    for (const m of classification.matched) {
      reconcileProps(m.oldFiber, m.newChild, host as HostConfig<string, Instance, unknown>, ctx as unknown);
    }
    const commitCtx: CommitContext<Instance> = {
      host: host as HostConfig<string, Instance, unknown>,
      ctx: ctx as unknown,
      createFiber: (node, parent) => createFiberForInsert(node, parent),
    };
    commitMutations(fiber, classification, commitCtx);
  }

  return {
    host,
    ctx,
    container,
    context: rootContext,
    rootFiber: null,
    render(node: UNode) {
      if (this.rootFiber) {
        const payloadChildren = isURoot(node) ? (node as URoot).children : [node];
        const resolvedChildren: UNode[] = [];
        for (const child of payloadChildren) {
          resolvedChildren.push(...resolveUNode(child));
        }
        const classification = reconcileChildren(
          this.rootFiber.children,
          resolvedChildren,
        );
        for (const m of classification.matched) {
          reconcileNode(m.oldFiber, m.newChild);
        }
        const commitCtx: CommitContext<Instance> = {
          host: host as HostConfig<string, Instance, unknown>,
          ctx: ctx as unknown,
          createFiber: (node, parent) => createFiberForInsert(node, parent),
        };
        commitMutations(this.rootFiber, classification, commitCtx);
        host.emit?.("root.render", `root_${Date.now()}`, { childCount: payloadChildren.length });
        return;
      }

      const root: Fiber<Instance> = {
        instance: undefined as unknown as Instance,
        tag: "#root",
        props: {},
        key: undefined,
        children: [],
        parent: null,
        effect: null,
        signalDisposers: [],
        prevProps: null,
        cachedNode: null,
        disposed: false,
        hash: null,
      };
      const payloadChildren = isURoot(node) ? (node as URoot).children : [node];
      for (const child of payloadChildren) {
        mountNode(child, root);
      }
      this.rootFiber = root;
      host.finalizeRoot?.(ctx);
      host.emit?.("root.render", `root_${Date.now()}`, { childCount: payloadChildren.length });
    },
    unmount() {
      const rootFiber = this.rootFiber;
      if (!rootFiber) return;

      disposeFiber(rootFiber, host as import("./fiber.js").HostLike<Instance, RootCtx>, ctx);

      for (const child of rootFiber.children) {
        host.removeChild?.(rootFiber.instance as never, child.instance as never, ctx as never);
      }

      host.finalizeRoot?.(ctx);
      host.emit?.("root.unmount", `root_${Date.now()}`, {});
      this.rootFiber = null;
    },
  };
}