feat: suspense

mod.ts

@@ -1,17 +1,18 @@
 export { createSignal } from "~/methods/createSignal.ts";
-export { createEffect } from "~/methods/createEffect.ts";
-export { on } from "~/methods/on.ts";
 export { createMemo } from "~/methods/createMemo.ts";
+export { createEffect } from "~/methods/createEffect.ts";
 export { createRoot } from "~/methods/createRoot.ts";
 export { batch } from "~/methods/batch.ts";
+export { withContext } from "~/methods/withContext.ts";
 export { getContext } from "~/methods/getContext.ts";
 export { untrack } from "~/methods/untrack.ts";
 export { onDispose } from "~/methods/onDispose.ts";
 export { catchError } from "~/methods/catchError.ts";
 export { tick } from "~/methods/tick.ts";
+export { on } from "~/methods/on.ts";
 export { createSelector } from "~/methods/createSelector.ts";
 export { withOwner } from "~/methods/withOwner.ts";
-export { withContext } from "~/methods/withContext.ts";
+export { withSuspense } from "~/methods/withSuspense.ts";
 
 export type {
   Accessor,

src/context.ts

@@ -34,3 +34,5 @@ export const EMPTY_CONTEXT: Contexts = {};
 
 export const ERRORHANDLER_SYMBOL = Symbol("Errorhandler");
 export const ERRORTHROWN_SYMBOL = Symbol("Error thrown");
+
+export const SUSPENSE_SYMBOL = Symbol("Suspense");

src/methods/createRoot.ts

@@ -2,5 +2,5 @@ import { Root } from "~/objects/root.ts";
 import type { RootFunction } from "~/types.ts";
 
 export function createRoot<T>(fn: RootFunction<T>): T {
-  return new Root(fn).wrap();
+  return new Root(fn).runWith();
 }

src/methods/getContext.ts

@@ -1,5 +1,5 @@
 import { CURRENTOWNER } from "~/context.ts";
 
 export function getContext<T = unknown>(id: symbol | string): T | undefined {
-  return CURRENTOWNER?.get(id);
+  return CURRENTOWNER?.contexts?.[id] as T | undefined;
 }

src/methods/withSuspense.ts

@@ -0,0 +1,17 @@
+import { Accessor } from "~/types.ts";
+import { Suspense } from "~/objects/suspense.ts";
+import { createEffect } from "~/methods/createEffect.ts";
+
+export function withSuspense<T>(suspended: Accessor<boolean>, fn: () => T): T {
+  const suspense = new Suspense();
+
+  createEffect(
+    () => {
+      suspense.toggle(suspended());
+    },
+    undefined,
+    { sync: true },
+  );
+
+  return suspense.runWith(fn);
+}

src/objects/context.ts

@@ -1,5 +1,6 @@
 import { Owner } from "~/objects/owner.ts";
 import { Contexts } from "~/types.ts";
+import { ERRORTHROWN_SYMBOL } from "~/context.ts";
 
 export class Context extends Owner {
   contexts: Contexts;
@@ -9,4 +10,10 @@ export class Context extends Owner {
 
     this.contexts = { ...this.parentScope?.contexts, ...contexts };
   }
+
+  runWith<T>(fn: () => T): T {
+    const result = Owner.runWithOwner(fn, this, undefined);
+
+    return result === ERRORTHROWN_SYMBOL ? undefined! : result;
+  }
 }

src/objects/effect.ts

@@ -2,10 +2,12 @@ import {
   ASYNCSCHEDULER,
   ERRORTHROWN_SYMBOL,
   STATE_CLEAN,
+  SUSPENSE_SYMBOL,
   SYNCSCHEDULER,
 } from "~/context.ts";
 import { Observer } from "~/objects/observer.ts";
 import type { CacheState, EffectOptions, ObserverFunction } from "~/types.ts";
+import { Suspense } from "~/objects/suspense.ts";
 
 /**
  * An effect is executed immediately on creation and every time again when one of its dependencies changes
@@ -14,28 +16,48 @@ export class Effect<T> extends Observer<T> {
   /** Stores the last return value of the callback */
   prevValue: T | undefined;
   declare readonly sync?: true;
+  init?: boolean;
+  suspense?: Suspense;
 
   constructor(
     fn: ObserverFunction<undefined | T, T>,
-    init?: T,
+    initialValue?: T,
     options?: EffectOptions,
   ) {
     super(fn);
 
-    this.prevValue = init;
+    this.suspense = this.get(SUSPENSE_SYMBOL);
+
+    this.prevValue = initialValue;
 
     if (options?.sync === true) {
       this.sync = true;
     }
 
     if (options?.sync === "init") {
+      // when the effect is suspended on first run, the suspense will run this effect instantly if it gets unsuspended
+      this.init = true;
+
+      // on first run we have to run the effect immediately
       this.update();
     } else {
       this.schedule();
     }
   }
 
+  /**
+   * Updates immediately if the effect is sync or schedules it with  the async scheduler otherwise
+   * This method is used on initialization when the effect is not caused to run by a signal and in suspense when it gets toggled.
+   * @returns
+   */
   schedule(): void {
+    // if there is a suspense (or multiple suspenses) only schedule if `suspended` > 0
+    if (this.suspense?.suspended) {
+      return;
+    }
+
+    // here we don't use the sync scheduler because it gets flushed by the observable
+    // which is not the case when this is the first run or caused by the suspense boundary toggling
     if (this.sync) {
       this.update();
     } else {
@@ -47,6 +69,11 @@ export class Effect<T> extends Observer<T> {
    * Just runs the callback with this effect as scope
    */
   override update(): void {
+    // if there is a suspense (or multiple suspenses) only schedule if `suspended` > 0
+    if (this.suspense?.suspended) {
+      return;
+    }
+
     const result = super.run(this.prevValue);
 
     if (result !== ERRORTHROWN_SYMBOL) {
@@ -66,8 +93,10 @@ export class Effect<T> extends Observer<T> {
     // if the state is not clean, it already has been added to execution queue
     if (this.state === STATE_CLEAN) {
       if (this.sync) {
+        // the scheduler will be run by the observable after marking all observers
         SYNCSCHEDULER.schedule(this);
       } else {
+        // the scheduler will be run on the next microtask
         ASYNCSCHEDULER.schedule(this);
       }
     }

src/objects/observer.ts

@@ -74,7 +74,7 @@ export abstract class Observer<T> extends Owner {
         if ((this.state as CacheState) === STATE_DIRTY) {
           // Stop the loop here so we won't trigger updates on other parents unnecessarily
           // If our computation changes to no longer use some sources, we don't
-          // want to update() a source we used last time, but now don't use.
+          // want to update a source we used last time, but now don't use.
           break;
         }
       }

src/objects/owner.ts

@@ -1,14 +1,22 @@
 import {
   CURRENTOBSERVER,
   CURRENTOWNER,
+  ERRORHANDLER_SYMBOL,
   ERRORTHROWN_SYMBOL,
   setObserver,
   setOwner,
   STATE_CLEAN,
   STATE_DISPOSED,
 } from "~/context.ts";
 import { handleError } from "~/utils/handleError.ts";
-import type { CacheState, CleanupFunction, Contexts } from "~/types.ts";
+import type {
+  CacheState,
+  CleanupFunction,
+  Contexts,
+  ErrorFunction,
+} from "~/types.ts";
+import { SUSPENSE_SYMBOL } from "~/context.ts";
+import { Suspense } from "~/objects/suspense.ts";
 
 /**
  * A scope is the abstraction over roots and computations. It provides contexts and can own other scopes
@@ -85,6 +93,10 @@ export class Owner {
     // this.contexts = {};
   }
 
+  get(id: typeof SUSPENSE_SYMBOL): Suspense | undefined;
+
+  get(id: typeof ERRORHANDLER_SYMBOL): ErrorFunction | undefined;
+
   /**
    * Searches for the context registered under the given id. If it is not found it searches recursively up the scope-tree
    * @param id The ID under which the context is registered

src/objects/root.ts

@@ -19,7 +19,7 @@ export class Root<T = unknown> extends Owner {
   /**
    * Executes the provided callback with the root as scope
    */
-  wrap(): T {
+  runWith(): T {
     const result = Owner.runWithOwner(
       () => this.fn(this.dispose.bind(this)),
       this,

src/objects/suspense.ts

@@ -0,0 +1,85 @@
+import { Owner } from "~/objects/owner.ts";
+import {
+  ERRORTHROWN_SYMBOL,
+  STATE_CHECK,
+  STATE_DIRTY,
+  SUSPENSE_SYMBOL,
+} from "~/context.ts";
+import { Contexts } from "~/types.ts";
+import { Effect } from "~/objects/effect.ts";
+
+export class Suspense extends Owner {
+  /**
+   * A counter that keeps track of the number of suspensions
+   * It is needed for nested suspense boundaries because parent suspenses suspend there children suspenses as well
+   * So if a children suspense gets unsuspended it should not necesarily run the effects as there could be a parent that still suspends
+   * In that case, `suspended` will still be > 0
+   */
+  public suspended: number;
+
+  contexts: Contexts = {
+    ...this.parentScope?.contexts,
+    [SUSPENSE_SYMBOL]: this,
+  };
+
+  constructor() {
+    super();
+
+    // get the state of a maybe existing parent suspense
+    this.suspended = this.parentScope?.get(SUSPENSE_SYMBOL)?.suspended ?? 0;
+  }
+
+  toggle(suspended: boolean) {
+    // suspended is already 0 meaning it is not suspended, we don't want it to be negative
+    if (!this.suspended && !suspended) {
+      return;
+    }
+
+    const prevSuspended = this.suspended;
+    const nextSuspended = this.suspended + (suspended ? 1 : -1);
+
+    this.suspended = nextSuspended;
+
+    // if the overall state hasn't changed, we don't need to do anything
+    // this can happen if the parent suspense boundary is still suspending
+    if (Boolean(prevSuspended) === Boolean(nextSuspended)) {
+      return;
+    }
+
+    notifyChildrenScopes(this, suspended);
+  }
+
+  runWith<T>(fn: () => T): T {
+    const result = Owner.runWithOwner(fn, this, undefined);
+
+    return result === ERRORTHROWN_SYMBOL ? undefined! : result;
+  }
+}
+
+function notifyChildrenScopes(scope: Owner, suspended: boolean) {
+  for (const childrenScope of scope.childrenScopes) {
+    // notify children suspenses
+    if (childrenScope instanceof Suspense) {
+      childrenScope.toggle(suspended);
+    }
+
+    if (childrenScope instanceof Effect) {
+      // if the effect should have run if there was no suspense boundary, we now update / schedule it
+      if (
+        childrenScope.state === STATE_CHECK ||
+        childrenScope.state === STATE_DIRTY
+      ) {
+        // If `init` is true it means that this effect should be sync on first run
+        if (childrenScope.init) {
+          // childrenScope.init = false;
+
+          childrenScope.update();
+        } else {
+          childrenScope.schedule();
+        }
+      }
+    }
+
+    notifyChildrenScopes(childrenScope, suspended);
+  }
+}

src/scheduler.ts

@@ -38,20 +38,28 @@ export abstract class Scheduler {
   runTop(node: Effect<any>) {
     const ancestors = [node];
 
+    // we traverse up the owner tree and find all effects with a `check` or `dirty` state and collect them
     while ((node = node.parentScope as Effect<any>)) {
       if (node?.state !== STATE_CLEAN) {
         ancestors.push(node);
       }
     }
 
+    // we run the effects in reverse order, meaning that the uppermost effect will be run first
     for (let i = ancestors.length - 1; i >= 0; i--) {
-      ancestors[i].updateIfNecessary();
+      // if the effect is currently suspended we don't run it
+      // TODO: Should we do this at a different place? Maybe in `updateIfNecessary`?
+      if (!ancestors[i].suspense?.suspended) {
+        ancestors[i].updateIfNecessary();
+      }
     }
   }
 }
 
-export class SyncScheduler extends Scheduler {
-}
+/**
+ * the sync scheduler is for `sync` effects. It gets flushed by an observable after all its observers have been marked
+ */
+export class SyncScheduler extends Scheduler {}
 
 /**
  * The async scheduler allows automatic batching by deferring the execution of effects to the next microtask