@assistant-ui/tap 0.4.5 → 0.5.0

package/src/core/{callResourceFn.ts → helpers/callResourceFn.ts}

@@ -1,4 +1,4 @@
-import { Resource } from "./types";
+import { Resource } from "../types";
 
 /**
  * Renders a resource with the given props.

package/src/core/{commit.ts → helpers/commit.ts}

@@ -1,9 +1,9 @@
-import { ResourceFiber, RenderResult } from "./types";
+import { ResourceFiber, RenderResult } from "../types";
 
-export function commitRender(renderResult: RenderResult): void {
+export function commitAllEffects(renderResult: RenderResult): void {
   const errors: unknown[] = [];
 
-  for (const task of renderResult.commitTasks) {
+  for (const task of renderResult.effectTasks) {
     try {
       task();
     } catch (error) {

package/src/core/helpers/env.ts

@@ -0,0 +1,3 @@
+export const isDevelopment =
+  typeof process !== "undefined" &&
+  (process.env.NODE_ENV === "development" || process.env.NODE_ENV === "test");

package/src/core/{execution-context.ts → helpers/execution-context.ts}

@@ -1,5 +1,5 @@
 import { isDevelopment } from "./env";
-import { ResourceFiber } from "./types";
+import { ResourceFiber } from "../types";
 
 let currentResourceFiber: ResourceFiber<any, any> | null = null;
 

package/src/core/helpers/root.ts

@@ -0,0 +1,67 @@
+import { Cell, ResourceFiber, ResourceFiberRoot } from "../types";
+
+export const createResourceFiberRoot = (
+  dispatchUpdate: (cb: () => boolean) => void,
+): ResourceFiberRoot => {
+  return {
+    version: 0,
+    committedVersion: 0,
+    dispatchUpdate,
+    changelog: [],
+    dirtyCells: [],
+  };
+};
+
+export const commitRoot = (root: ResourceFiberRoot): void => {
+  for (const cell of root.dirtyCells) {
+    cell.dirty = false;
+    cell.queue.clear();
+    cell.current = cell.workInProgress;
+  }
+  root.committedVersion = root.version;
+  root.changelog.length = 0;
+  root.dirtyCells.length = 0;
+};
+
+export const setRootVersion = (
+  root: ResourceFiberRoot,
+  version: number,
+): void => {
+  const rollback = root.version > version;
+  root.version = version;
+  if (rollback) {
+    for (const cell of root.dirtyCells) {
+      cell.dirty = false;
+      cell.queue.clear();
+      cell.workInProgress = cell.current;
+    }
+    root.dirtyCells.length = 0;
+
+    if (version === root.committedVersion) {
+      root.changelog.length = 0;
+    } else {
+      // commit happened without a useEffect update (offscreen API)
+
+      if (root.committedVersion > version)
+        throw new Error("Version is less than committed version");
+
+      while (root.committedVersion + root.changelog.length > version) {
+        root.changelog.pop();
+      }
+
+      root.changelog.forEach((apply) => apply());
+      commitRoot(root);
+    }
+  }
+};
+
+export const markCellDirty = (
+  fiber: ResourceFiber<any, any>,
+  cell: Cell & { type: "reducer" },
+): void => {
+  if (!cell.dirty) {
+    cell.dirty = true;
+    fiber.markDirty?.();
+    fiber.root.dirtyCells.push(cell);
+  }
+};

package/src/core/resource.ts

@@ -1,5 +1,5 @@
 import { ResourceElement } from "./types";
-import { fnSymbol } from "./callResourceFn";
+import { fnSymbol } from "./helpers/callResourceFn";
 
 export function resource<R>(fn: () => R): () => ResourceElement<R, undefined>;
 export function resource<R, P>(

package/src/core/scheduler.ts

@@ -36,7 +36,7 @@ export class UpdateScheduler {
 const scheduleFlush = () => {
   if (flushState.isScheduled) return;
   flushState.isScheduled = true;
-  queueMicrotask(flushScheduled);
+  scheduleMacrotask();
 };
 
 const flushScheduled = () => {
@@ -80,6 +80,18 @@ const flushScheduled = () => {
   }
 };
 
+// Use MessageChannel to schedule flushes as macrotasks (like React's scheduler).
+// This allows more state updates to batch into a single re-render.
+const scheduleMacrotask = (() => {
+  if (typeof MessageChannel !== "undefined") {
+    const channel = new MessageChannel();
+    channel.port1.onmessage = flushScheduled;
+    return () => channel.port2.postMessage(null);
+  }
+  // Fallback for environments without MessageChannel
+  return () => setTimeout(flushScheduled, 0);
+})();
+
 export const flushResourcesSync = <T>(callback: () => T): T => {
   const prev = flushState;
   flushState = {

package/src/core/types.ts

@@ -1,6 +1,5 @@
 import type { tapEffect } from "../hooks/tap-effect";
-import type { tapState } from "../hooks/tap-state";
-import type { fnSymbol } from "./callResourceFn";
+import type { fnSymbol } from "./helpers/callResourceFn";
 
 export type ResourceElement<R, P = any> = {
   readonly type: Resource<R, P> & { [fnSymbol]: (props: P) => R };
@@ -18,11 +17,22 @@ export type ExtractResourceReturnType<T> =
       ? R
       : never;
 
+export interface ReducerQueueEntry {
+  readonly action: any;
+  hasEagerState: boolean;
+  eagerState: any;
+}
+
 export type Cell =
   | {
-      readonly type: "state";
-      value: any;
-      set: (updater: tapState.StateUpdater<any>) => void;
+      readonly type: "reducer";
+      readonly dispatch: (action: any) => void;
+
+      readonly queue: Set<ReducerQueueEntry>;
+      dirty: boolean;
+      workInProgress: any;
+      current: any;
+      reducer: (state: any, action: any) => any;
     }
   | {
       readonly type: "effect";
@@ -39,12 +49,22 @@ export interface EffectTask {
 export interface RenderResult {
   readonly output: any;
   readonly props: any;
-  readonly commitTasks: (() => void)[];
+  readonly effectTasks: (() => void)[];
 }
 
-export interface ResourceFiber<R, P> {
+export interface ResourceFiberRoot {
+  version: number;
+  committedVersion: number;
+  readonly changelog: (() => void)[];
+
   readonly dispatchUpdate: (callback: () => boolean) => void;
+  readonly dirtyCells: (Cell & { type: "reducer" })[];
+}
+
+export interface ResourceFiber<R, P> {
+  readonly root: ResourceFiberRoot;
   readonly type: Resource<R, P>;
+  readonly markDirty: (() => void) | undefined;
   readonly devStrictMode: "root" | "child" | null;
 
   cells: Cell[];

package/src/hooks/tap-effect-event.ts

@@ -1,8 +1,8 @@
 import { tapRef } from "./tap-ref";
 import { tapEffect } from "./tap-effect";
-import { isDevelopment } from "../core/env";
+import { isDevelopment } from "../core/helpers/env";
 import { tapCallback } from "./tap-callback";
-import { getCurrentResourceFiber } from "../core/execution-context";
+import { getCurrentResourceFiber } from "../core/helpers/execution-context";
 
 /**
  * Creates a stable function reference that always calls the most recent version of the callback.
@@ -24,6 +24,7 @@ export function tapEffectEvent<T extends (...args: any[]) => any>(
 ): T {
   const callbackRef = tapRef(callback);
 
+  // TODO this effect needs to run before all userland effects
   tapEffect(() => {
     callbackRef.current = callback;
   });

package/src/hooks/tap-memo.ts

@@ -1,25 +1,30 @@
-import { isDevelopment } from "../core/env";
-import { getCurrentResourceFiber } from "../core/execution-context";
-import { tapRef } from "./tap-ref";
+import { isDevelopment } from "../core/helpers/env";
+import { getCurrentResourceFiber } from "../core/helpers/execution-context";
+import { tapReducerWithDerivedState } from "./tap-reducer";
 import { depsShallowEqual } from "./utils/depsShallowEqual";
 
-export const tapMemo = <T>(fn: () => T, deps: readonly unknown[]) => {
-  const dataRef = tapRef<{ value: T; deps: readonly unknown[] }>();
-  if (!dataRef.current) {
-    if (isDevelopment) {
-      const fiber = getCurrentResourceFiber();
-      if (fiber.devStrictMode) {
-        void fn();
-      }
-    }
+const memoReducer = () => {
+  throw new Error("Memo reducer should not be called");
+};
+
+type MemoState<T> = { value: T; deps: readonly unknown[] };
 
-    dataRef.current = { value: fn(), deps };
-  }
+export const tapMemo = <T>(fn: () => T, deps: readonly unknown[]): T => {
+  const fiber = getCurrentResourceFiber();
+  const [state] = tapReducerWithDerivedState(
+    memoReducer,
+    (state: MemoState<T> | null): MemoState<T> => {
+      if (state && depsShallowEqual(state.deps, deps)) return state;
 
-  if (!depsShallowEqual(dataRef.current.deps, deps)) {
-    dataRef.current.value = fn();
-    dataRef.current.deps = deps;
-  }
+      const value = fn();
+
+      if (isDevelopment && fiber.devStrictMode) {
+        void fn();
+      }
 
-  return dataRef.current.value;
+      return { value, deps };
+    },
+    null,
+  );
+  return state.value;
 };

package/src/hooks/tap-reducer.ts

@@ -0,0 +1,148 @@
+import { isDevelopment } from "../core/helpers/env";
+import { getCurrentResourceFiber } from "../core/helpers/execution-context";
+import { ReducerQueueEntry, ResourceFiber } from "../core/types";
+import { markCellDirty } from "../core/helpers/root";
+import { tapHook } from "./utils/tapHook";
+
+type Dispatch<A> = (action: A) => void;
+
+const dispatchOnFiber = (
+  fiber: ResourceFiber<any, any>,
+  callback: () => (() => void) | null,
+): void => {
+  if (fiber.renderContext) {
+    throw new Error("Resource updated during render");
+  }
+  if (fiber.isNeverMounted) {
+    throw new Error("Resource updated before mount");
+  }
+
+  fiber.root.dispatchUpdate(() => {
+    const result = callback();
+    if (result) {
+      result();
+      fiber.root.changelog.push(result);
+      return true;
+    }
+    return false;
+  });
+};
+
+function tapReducerImpl<S, A, I, R extends S>(
+  reducer: (state: S, action: A) => S,
+  getDerivedState: ((state: S) => R) | undefined,
+  initialArg: S | I,
+  initFn: ((arg: I) => S) | undefined,
+): [R, Dispatch<A>] {
+  const cell = tapHook("reducer", () => {
+    const fiber = getCurrentResourceFiber();
+
+    // First render: compute initial state
+    const initialState = initFn ? initFn(initialArg as I) : initialArg;
+
+    if (isDevelopment && fiber.devStrictMode && initFn) {
+      void initFn(initialArg as I);
+    }
+
+    return {
+      type: "reducer",
+      queue: new Set(),
+      dirty: false,
+      workInProgress: initialState,
+      current: initialState,
+      reducer,
+      dispatch: (action: A) => {
+        const entry: ReducerQueueEntry = {
+          action,
+          hasEagerState: false,
+          eagerState: undefined,
+        };
+
+        dispatchOnFiber(fiber, () => {
+          if (fiber.root.dirtyCells.length === 0 && !entry.hasEagerState) {
+            entry.eagerState = reducer(cell.workInProgress, action);
+            entry.hasEagerState = true;
+
+            if (Object.is(cell.current, entry.eagerState)) return null;
+          }
+
+          return () => {
+            markCellDirty(fiber, cell);
+            cell.queue.add(entry);
+          };
+        });
+      },
+    };
+  });
+
+  const fiber = getCurrentResourceFiber();
+  const sameReducer = reducer === cell.reducer;
+  cell.reducer = reducer;
+
+  for (const item of cell.queue) {
+    if (!item.hasEagerState || !sameReducer) {
+      item.eagerState = reducer(cell.workInProgress, item.action);
+      item.hasEagerState = true;
+    }
+
+    if (isDevelopment && fiber.devStrictMode) {
+      void reducer(cell.workInProgress, item.action);
+    }
+
+    cell.workInProgress = item.eagerState;
+  }
+  cell.queue.clear();
+
+  if (getDerivedState) {
+    const derived = getDerivedState(cell.workInProgress);
+
+    if (!Object.is(derived, cell.workInProgress)) {
+      markCellDirty(fiber, cell);
+      cell.workInProgress = derived;
+    }
+  }
+
+  return [cell.workInProgress, cell.dispatch];
+}
+
+export function tapReducer<S, A>(
+  reducer: (state: S, action: A) => S,
+  initialState: S,
+): [S, Dispatch<A>];
+export function tapReducer<S, A, I>(
+  reducer: (state: S, action: A) => S,
+  initialArg: I,
+  init: (arg: I) => S,
+): [S, Dispatch<A>];
+export function tapReducer<S, A, I>(
+  reducer: (state: S, action: A) => S,
+  initialArg: S | I,
+  init?: (arg: I) => S,
+): [S, Dispatch<A>] {
+  return tapReducerImpl(
+    reducer,
+    undefined,
+    initialArg as S,
+    init as ((arg: S) => S) | undefined,
+  );
+}
+
+export function tapReducerWithDerivedState<S, A, R extends S>(
+  reducer: (state: S, action: A) => S,
+  getDerivedState: (state: S) => R,
+  initialState: S,
+): [R, Dispatch<A>];
+export function tapReducerWithDerivedState<S, A, I, R extends S>(
+  reducer: (state: S, action: A) => S,
+  getDerivedState: (state: S) => R,
+  initialArg: I,
+  init: (arg: I) => S,
+): [R, Dispatch<A>];
+export function tapReducerWithDerivedState<S, A, I, R extends S>(
+  reducer: (state: S, action: A) => S,
+  getDerivedState: (state: S) => R,
+  initialArg: I,
+  init?: (arg: I) => S,
+): [R, Dispatch<A>] {
+  return tapReducerImpl(reducer, getDerivedState, initialArg, init);
+}

package/src/hooks/tap-resource.ts

@@ -7,8 +7,8 @@ import {
   commitResourceFiber,
 } from "../core/ResourceFiber";
 import { tapMemo } from "./tap-memo";
-import { tapState } from "./tap-state";
-import { tapConst } from "./tap-const";
+import { tapRef } from "./tap-ref";
+import { getCurrentResourceFiber } from "../core/helpers/execution-context";
 
 export function tapResource<E extends ResourceElement<any, any>>(
   element: E,
@@ -21,21 +21,21 @@ export function tapResource<E extends ResourceElement<any, any>>(
   element: E,
   propsDeps?: readonly unknown[],
 ): ExtractResourceReturnType<E> {
-  const [version, setVersion] = tapState(0);
-  const rerender = tapConst(() => () => setVersion((v) => v + 1), []);
-
+  const parentFiber = getCurrentResourceFiber();
+  const versionRef = tapRef(0);
   const fiber = tapMemo(() => {
     void element.key;
-    return createResourceFiber(element.type, (callback) => {
-      if (callback()) rerender();
+    return createResourceFiber(element.type, parentFiber.root, () => {
+      versionRef.current++;
+      parentFiber.markDirty?.();
     });
-  }, [element.type, element.key]);
+  }, [element.type, element.key, parentFiber]);
 
   const result = propsDeps
     ? // biome-ignore lint/correctness/useExhaustiveDependencies: user provided deps instead of prop identity
       tapMemo(
         () => renderResourceFiber(fiber, element.props),
-        [fiber, ...propsDeps, version],
+        [fiber, ...propsDeps, versionRef.current],
       )
     : renderResourceFiber(fiber, element.props);
 

package/src/hooks/tap-resources.ts

@@ -6,7 +6,6 @@ import {
 } from "../core/types";
 import { tapEffect } from "./tap-effect";
 import { tapMemo } from "./tap-memo";
-import { tapState } from "./tap-state";
 import { tapCallback } from "./tap-callback";
 import {
   createResourceFiber,
@@ -15,6 +14,8 @@ import {
   commitResourceFiber,
 } from "../core/ResourceFiber";
 import { tapConst } from "./tap-const";
+import { tapRef } from "./tap-ref";
+import { getCurrentResourceFiber } from "../core/helpers/execution-context";
 
 type FiberState = {
   fiber: ResourceFiber<unknown, unknown>;
@@ -28,9 +29,17 @@ export function tapResources<E extends ResourceElement<any, any>>(
   getElements: () => readonly E[],
   getElementsDeps?: readonly unknown[],
 ): ExtractResourceReturnType<E>[] {
-  const [version, setVersion] = tapState(0);
-  const rerender = tapConst(() => () => setVersion((v) => v + 1), []);
-
+  const versionRef = tapRef(0);
+  const version = versionRef.current;
+
+  const parentFiber = tapConst(getCurrentResourceFiber, []);
+  const markDirty = tapConst(
+    () => () => {
+      versionRef.current++;
+      parentFiber.markDirty?.();
+    },
+    [],
+  );
   const fibers = tapConst(() => new Map<string | number, FiberState>(), []);
 
   const getElementsMemo = getElementsDeps
@@ -65,9 +74,11 @@ export function tapResources<E extends ResourceElement<any, any>>(
 
       let state = fibers.get(elementKey);
       if (!state) {
-        const fiber = createResourceFiber(element.type, (callback) => {
-          if (callback()) rerender();
-        });
+        const fiber = createResourceFiber(
+          element.type,
+          parentFiber.root,
+          markDirty,
+        );
         const result = renderResourceFiber(fiber, element.props);
         state = {
           fiber,
@@ -77,9 +88,11 @@ export function tapResources<E extends ResourceElement<any, any>>(
         fibers.set(elementKey, state);
         results.push(result.output);
       } else if (state.fiber.type !== element.type) {
-        const fiber = createResourceFiber(element.type, (callback) => {
-          if (callback()) rerender();
-        });
+        const fiber = createResourceFiber(
+          element.type,
+          parentFiber.root,
+          markDirty,
+        );
         const result = renderResourceFiber(fiber, element.props);
         state.next = [fiber, result];
         results.push(result.output);

package/src/hooks/tap-state.ts

@@ -1,94 +1,19 @@
-import { isDevelopment } from "../core/env";
-import { getCurrentResourceFiber } from "../core/execution-context";
-import { Cell, ResourceFiber } from "../core/types";
+import { tapReducer } from "./tap-reducer";
 
 export namespace tapState {
   export type StateUpdater<S> = S | ((prev: S) => S);
 }
 
-const dispatchOnFiber = (
-  fiber: ResourceFiber<any, any>,
-  callback: () => boolean,
-) => {
-  if (fiber.renderContext) {
-    throw new Error("Resource updated during render");
-  }
+const stateReducer = <S>(
+  state: S | undefined,
+  action: tapState.StateUpdater<S>,
+): S =>
+  typeof action === "function"
+    ? (action as (prev: S | undefined) => S)(state)
+    : action;
 
-  if (fiber.isMounted) {
-    // Only schedule rerender if currently mounted
-    fiber.dispatchUpdate(callback);
-  } else if (fiber.isNeverMounted) {
-    throw new Error("Resource updated before mount");
-  }
-
-  // TODO mark dirty
-};
-
-function getStateCell<T>(
-  initialValue: T | (() => T),
-): Cell & { type: "state" } {
-  const fiber = getCurrentResourceFiber();
-  const index = fiber.currentIndex++;
-
-  // Check if we're trying to use more hooks than in previous renders
-  if (!fiber.isFirstRender && index >= fiber.cells.length) {
-    throw new Error(
-      "Rendered more hooks than during the previous render. " +
-        "Hooks must be called in the exact same order in every render.",
-    );
-  }
-
-  const cell = fiber.cells[index];
-  if (cell) {
-    if (cell.type !== "state")
-      throw new Error("Hook order changed between renders");
-
-    return cell as Cell & { type: "state" };
-  }
-
-  const value =
-    typeof initialValue === "function"
-      ? (initialValue as () => T)()
-      : initialValue;
-
-  if (
-    isDevelopment &&
-    fiber.devStrictMode &&
-    typeof initialValue === "function"
-  ) {
-    void (initialValue as () => T)();
-  }
-
-  const newCell: Cell & { type: "state" } = {
-    type: "state",
-    value,
-    set: (updater: tapState.StateUpdater<T>) => {
-      dispatchOnFiber(fiber, () => {
-        const currentValue = newCell.value;
-        const nextValue =
-          typeof updater === "function"
-            ? (updater as (prev: T) => T)(currentValue)
-            : updater;
-
-        if (
-          isDevelopment &&
-          fiber.devStrictMode &&
-          typeof updater === "function"
-        ) {
-          void (updater as (prev: T) => T)(currentValue);
-        }
-
-        if (Object.is(currentValue, nextValue)) return false;
-
-        newCell.value = nextValue;
-        return true;
-      });
-    },
-  };
-
-  fiber.cells[index] = newCell;
-  return newCell;
-}
+const stateInit = <S>(initial: S | (() => S)): S =>
+  typeof initial === "function" ? (initial as () => S)() : initial;
 
 export function tapState<S = undefined>(): [
   S | undefined,
@@ -100,7 +25,5 @@ export function tapState<S>(
 export function tapState<S>(
   initial?: S | (() => S),
 ): [S | undefined, (updater: tapState.StateUpdater<S>) => void] {
-  const cell = getStateCell(initial as S | (() => S));
-
-  return [cell.value, cell.set];
+  return tapReducer(stateReducer, initial, stateInit);
 }

package/src/hooks/utils/tapHook.ts

@@ -1,7 +1,7 @@
-import { getCurrentResourceFiber } from "../../core/execution-context";
+import { getCurrentResourceFiber } from "../../core/helpers/execution-context";
 import { Cell } from "../../core/types";
 
-export const tapHook = <T extends "effect" | "memo">(
+export const tapHook = <T extends Cell["type"]>(
   type: T,
   init: () => Cell,
 ): Cell & { type: T } => {
@@ -31,5 +31,5 @@ export const tapHook = <T extends "effect" | "memo">(
 
 export const registerRenderMountTask = (task: () => void) => {
   const fiber = getCurrentResourceFiber();
-  fiber.renderContext!.commitTasks.push(task);
+  fiber.renderContext!.effectTasks.push(task);
 };

package/src/index.ts

@@ -3,6 +3,7 @@ export { withKey } from "./core/withKey";
 
 // primitive hooks
 export { tapState } from "./hooks/tap-state";
+export { tapReducer, tapReducerWithDerivedState } from "./hooks/tap-reducer";
 export { tapEffect } from "./hooks/tap-effect";
 
 // utility hooks
@@ -14,14 +15,13 @@ export { tapEffectEvent } from "./hooks/tap-effect-event";
 
 // resources
 export { tapResource } from "./hooks/tap-resource";
-export { tapInlineResource } from "./hooks/tap-inline-resource";
 export { tapResources } from "./hooks/tap-resources";
 
 // subscribable
-export { tapSubscribableResource } from "./tapSubscribableResource";
+export { tapResourceRoot } from "./tapResourceRoot";
 
 // imperative
-export { createResource } from "./core/createResource";
+export { createResourceRoot } from "./core/createResourceRoot";
 export { flushResourcesSync } from "./core/scheduler";
 
 // context