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 = WebSignal.subtle.untrack;
const __currentComputed = WebSignal.subtle.currentComputed;

const __State = WebSignal.State;
const __Computed = WebSignal.Computed;
const __Watcher = WebSignal.subtle.Watcher;

const __State_read = __State.prototype.get;
const __Computed_read = __Computed.prototype.get;

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

let currentOwner: null | Owner = null;
let batchedEffects: null | { compute: WebSignal.Computed<void>; watcher?: WebSignal.subtle.Watcher }[] = 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) {
		runComputation(() => {
			for (let idx = 0, len = effects.length; idx < len; idx++) {
				const { compute: e, watcher: w } = effects[idx];
				e.get();
				w?.watch(e);
			}
		});
	}
}

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
		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) {
	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;
		}
	});

	const watcher = new __Watcher(() => {
		if (batchedEffects) {
			batchedEffects.push({ compute, watcher });
		} else {
			// The only way this could happen is if we're listening to a signal that
			// we don't manage.
			queueMicrotask(() => {
				compute.get();
				watcher.watch(compute);
			});
		}
	});

	onCleanup(() => {
		watcher.unwatch(compute);
		cleanup();
	});

	watcher.watch(compute);

	return 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 {
	const compute = createBackingEffect(fn, value);

	// Run it now.
	compute.get();
}

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 {
	const compute = createBackingEffect(fn, value);

	if (batchedEffects) {
		batchedEffects.push({ compute });
	} else {
		compute.get();
	}
}

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);
	}
}