Solid.js-like signals on top of the TC39 Signals proposal

solid-signals.ts

import { Signal as WebSignal } from 'signal-polyfill';

export type Accessor<T> = () => T;
export type Setter<in out T> = {
	<U extends T>(...args: undefined extends T ? [] : [value: (prev: T) => U]): undefined extends T
		? undefined
		: U;
	<U extends T>(value: (prev: T) => U): U;
	<U extends T>(value: Exclude<U, Function>): U;
	<U extends T>(value: Exclude<U, Function> | ((prev: T) => U)): U;
};

export type Signal<T> = [get: Accessor<T>, set: Setter<T>];

export interface SignalOptions<T> {
	equals?: false | ((prev: T, next: T) => boolean);
}

export type EffectFunction<Prev, Next extends Prev = Prev> = (v: Prev) => Next;

export type RootFunction<T> = (dispose: () => void) => T;

export type ErrorHandler = (err: unknown) => void;

export interface Owner {
	owner: null | Owner;
	cleanups: null | (() => void)[];
	catch: null | ErrorHandler;
	context: null | Record<string | symbol, unknown>;
}

const __untrack = /* #__PURE__ */ WebSignal.subtle.untrack;
const __currentComputed = /* #__PURE__ */ WebSignal.subtle.currentComputed;

const __State = /* #__PURE__ */ WebSignal.State;
const __Computed = /* #__PURE__ */ WebSignal.Computed;
const __Watcher = /* #__PURE__ */ WebSignal.subtle.Watcher;

const __State_read = /* #__PURE__ */ __State.prototype.get;
const __Computed_read = /* #__PURE__ */ __Computed.prototype.get;

export const untrack: <T>(fn: Accessor<T>) => T = __untrack;

let currentOwner: null | Owner = null;
let batchedEffects: null | WebSignal.Computed<void>[] = null;

function alwaysInvalidate(): false {
	return false;
}

export function batch<T>(fn: Accessor<T>): T {
	return runComputation(fn) as T;
}

export function createSignal<T>(): Signal<T | undefined>;
export function createSignal<T>(value: T, options?: SignalOptions<T>): Signal<T>;
export function createSignal<T>(value?: T, options?: SignalOptions<T | undefined>): Signal<T | undefined> {
	const equals = options?.equals;

	const backing = new __State(value, {
		equals: equals === false ? alwaysInvalidate : equals,
	});

	const setter = (next?: T) => {
		if (typeof next === 'function') {
			next = next(value);
		}

		runComputation(() => backing.set((value = next)));
		return value;
	};

	// @ts-expect-error
	return [__State_read.bind(backing), setter];
}

export function createMemo<Next extends Prev, Prev = Next>(
	fn: EffectFunction<undefined | NoInfer<Prev>, Next>,
): Accessor<Next>;
export function createMemo<Next extends Prev, Init = Next, Prev = Next>(
	fn: EffectFunction<Init | Prev, Next>,
	value: Init,
	options?: SignalOptions<Next>,
): Accessor<Next>;
export function createMemo<Next extends Prev, Init, Prev>(
	fn: EffectFunction<Init | Prev, Next>,
	value?: Init,
	options?: SignalOptions<Next>,
): Accessor<Next> {
	const equals = options?.equals;

	// @ts-expect-error
	const backing = new __Computed(() => (value = fn(value)), {
		equals: equals === false ? alwaysInvalidate : equals,
	});

	return __Computed_read.bind(backing);
}

export function isListening(): boolean {
	return __currentComputed() !== undefined;
}

export function getOwner(): Owner | null {
	return currentOwner;
}

export function runWithOwner<T>(owner: typeof currentOwner, fn: Accessor<T>): T | undefined {
	const previousOwner = currentOwner;

	try {
		currentOwner = owner;
		return isListening() ? runComputation(() => untrack(fn)) : runComputation(fn);
	} catch (err) {
		handleError(err);
	} finally {
		currentOwner = previousOwner;
	}
}

function runComputation<T>(fn: Accessor<T>): T | undefined {
	if (batchedEffects !== null) {
		return fn();
	}

	try {
		batchedEffects = [];

		const result = fn();
		completeComputation();

		return result;
	} catch (err) {
		batchedEffects = null;
		handleError(err);
	}
}

function completeComputation() {
	const effects = batchedEffects!;
	batchedEffects = null;

	if (effects.length > 0) {
		// Solid.js seems to unwatch all subsequent effects if one of them throws
		// an error without a catch handler, this doesn't seem ideal.

		let hasError = false;
		let error: unknown;

		runComputation(() => {
			for (let idx = 0, len = effects.length; idx < len; idx++) {
				const compute = effects[idx];

				// @ts-expect-error
				if (compute.__destroyed) {
					continue;
				}

				try {
					compute.get();
				} catch (e) {
					hasError = true;
					error = e;
				}
			}
		});

		// `runComputation` has the chance to throw an error here.

		if (hasError) {
			throw error;
		}
	}
}

export function createRoot<T>(fn: RootFunction<T>, parentOwner = currentOwner): T {
	const previousOwner = currentOwner;

	const owner: Owner = {
		owner: parentOwner,
		cleanups: null,
		context: parentOwner ? parentOwner.context : null,
		catch: parentOwner ? parentOwner.catch : null,
	};

	try {
		currentOwner = owner;

		// @ts-expect-error: We're just gonna pretend that this will always return
		// even if there's a catch handler set up at the parent owner, for now.
		return runComputation(() => {
			return fn(() => {
				return untrack(() => {
					const cleanups = owner.cleanups;

					if (cleanups !== null) {
						owner.cleanups = null;

						for (let i = 0, il = cleanups.length; i < il; i++) {
							(0, cleanups[i])();
						}
					}
				});
			});
		});
	} finally {
		currentOwner = previousOwner;
	}
}

export function onCleanup(fn: () => void): void {
	if (currentOwner) {
		const cleanups = currentOwner.cleanups;

		if (cleanups !== null) {
			cleanups.push(fn);
		} else {
			currentOwner.cleanups = [fn];
		}
	}
}

function createBackingEffect(fn: EffectFunction<any, any>, value: any, defer: boolean): void {
	const owner: Owner = {
		owner: currentOwner,
		cleanups: null,
		context: currentOwner ? currentOwner.context : null,
		catch: currentOwner ? currentOwner.catch : null,
	};

	const cleanup = () => {
		const cleanups = owner.cleanups;

		if (cleanups !== null) {
			owner.cleanups = null;

			for (let idx = 0, len = cleanups.length; idx < len; idx++) {
				(0, cleanups[idx])();
			}
		}
	};

	const compute = new __Computed(() => {
		if (owner.cleanups !== null) {
			untrack(cleanup);
		}

		const previousOwner = currentOwner;

		try {
			currentOwner = owner;
			value = fn(value);
		} catch (err) {
			handleError(err);
		} finally {
			currentOwner = previousOwner;
		}
	});

	// Watchers don't tell us which computed are now marked dirty, it has to go
	// through `.getPending()`, which gives us an array of computed signals, where
	// it'll also include ones signals we've already queued anyway.

	// Hence each computed signal having their own watcher.
	const watcher = new __Watcher(() => {
		// Watchers cannot directly read or write into signals. which is
		// problematic, as we need synchronous reactions.

		// This seems to only be the case for as long as they're *inside* this
		// callback however, so thankfully it isn't much of a problem so long as
		// we're dealing with our own signals.

		// Our own signals are wrapped in `runComputation` which sets up a global
		// variable for us to dump these dirty computed signals onto.

		// We'll queue a microtask otherwise.

		if (batchedEffects) {
			batchedEffects.push(compute);
		} else {
			queueMicrotask(() => {
				// @ts-expect-error
				if (compute.__destroyed) {
					return;
				}

				runComputation(() => compute.get());
			});
		}

		// Make this watcher callable again. I believe it shouldn't fire so long as
		// these computed signals are still marked as dirty (we haven't called
		// `.get()` on the computed signal.)
		watcher.watch();
	});

	onCleanup(() => {
		// @ts-expect-error
		compute.__destroyed = true;

		watcher.unwatch(compute);
		cleanup();
	});

	// @ts-expect-error
	compute.__destroyed = false;

	if (defer && batchedEffects) {
		batchedEffects.push(compute);
	} else {
		compute.get();
		watcher.watch(compute);
	}
}

export function createRenderEffect<Next>(fn: EffectFunction<undefined | NoInfer<Next>, Next>): void;
export function createRenderEffect<Next, Init = Next>(
	fn: EffectFunction<Init | Next, Next>,
	value: Init,
): void;
export function createRenderEffect<Next, Init>(fn: EffectFunction<Init | Next, Next>, value?: Init): void {
	createBackingEffect(fn, value, true);
}

export function createEffect<Next>(fn: EffectFunction<undefined | NoInfer<Next>, Next>): void;
export function createEffect<Next, Init = Next>(fn: EffectFunction<Init | Next, Next>, value: Init): void;
export function createEffect<Next, Init>(fn: EffectFunction<Init | Next, Next>, value?: Init): void {
	createBackingEffect(fn, value, true);
}

export function catchError<T>(fn: Accessor<T>, handler: (err: unknown) => void) {
	const previousOwner = currentOwner;

	try {
		currentOwner = {
			owner: currentOwner,
			cleanups: null,
			context: currentOwner ? currentOwner.context : null,
			catch: handler,
		};

		return fn();
	} catch (err) {
		handleError(err);
	} finally {
		currentOwner = previousOwner;
	}
}

function handleError(err: unknown, owner = currentOwner) {
	const handler = owner?.catch;
	if (!handler) {
		throw err;
	}

	runErrorHandler(err, handler, owner);
}

function runErrorHandler(err: unknown, handler: ErrorHandler, owner: Owner | null) {
	try {
		handler(err);
	} catch (e) {
		handleError(e, (owner && owner.owner) || null);
	}
}