Incomplete state manager test inspired by Jotai

atoms.ts

interface Atom<State> {
	listeners: Set<(value: State) => void>;
	get<_Result>(selector?: undefined): State;
	get<Result>(selector: (state: State) => Result): Result;
	set(state: State): void;
	set(updater: (previous: State) => State): void;
	memoize(equality: (previous: State, current: State) => boolean): Atom<State>;
}

interface DerivedAtom<State> extends Atom<State> {
	set: never;
}

type AtomStates<T extends Atom<any>[]> = {
	[Key in keyof T]: T[Key] extends Atom<infer State> ? State : never;
};

function weakReference<T extends defined>(value: T): () => T | undefined {
	const weakSet = new WeakSet([value]);
	return () => next(weakSet)[0];
}

function atom<State>(state: State): Atom<State> {
	let equals = (previous: State, current: State) => previous === current;

	return {
		listeners: new Set(),
		get(selector) {
			return selector ? selector(state) : state;
		},
		set(nextState) {
			nextState = typeIs(nextState, "function") ? nextState(state) : nextState;

			if (!equals(state, nextState)) {
				state = nextState;
				table.clone(this.listeners).forEach((fn) => fn(state));
			}
		},
		memoize(equality) {
			equals = equality;
			return this;
		},
	};
}

function derive<Atoms extends Atom<any>[], Result>(
	...args: [...atoms: Atoms, combiner: (...atoms: AtomStates<Atoms>) => Result]
): DerivedAtom<Result>;

function derive<Atoms extends Atom<defined>[], Result>(
	...args: [...atoms: Atoms, combiner: (...atoms: AtomStates<Atoms>) => Result]
): DerivedAtom<Result> {
	const combiner = args.pop() as (...atoms: AtomStates<Atoms>) => Result;
	const atoms = args as unknown as Atoms;

	const compute = () => {
		const inputs = atoms.map((atom) => atom.get()) as AtomStates<Atoms>;
		return combiner(...inputs);
	};

	const combinerAtom = atom(compute());
	const getCombinerAtom = weakReference(combinerAtom);

	const update = () => {
		const atom = getCombinerAtom();

		if (atom) {
			atom.set(compute());
		} else {
			atoms.forEach((atom) => atom.listeners.delete(update));
		}
	};

	for (const atom of atoms) {
		atom.listeners.add(update);
	}

	return combinerAtom as DerivedAtom<Result>;
}

function subscribe<State>(atom: Atom<State>, callback: (state: State) => void): () => void {
	let pending: thread | undefined;

	const update = () => {
		pending ??= task.defer(() => {
			pending = undefined;
			callback(atom.get());
		});
	};

	atom.listeners.add(update);

	return () => {
		atom.listeners.delete(update);
		if (pending) task.cancel(pending);
	};
}

function observe<K, V>(
	atom: Atom<Map<K, V> | ReadonlyMap<K, V>>,
	factory: (value: V, key: K) => () => void,
): () => void;

function observe(atom: Atom<any>, factory: (value: unknown, key: unknown) => () => void): () => void {
	const disposers = new Map<unknown, () => void>();

	const update = (state: Map<unknown, unknown>) => {
		for (const [key, dispose] of pairs(disposers)) {
			if (!state.has(key)) {
				disposers.delete(key);
				dispose();
			}
		}

		for (const [key, value] of pairs(state)) {
			if (!disposers.has(key)) {
				disposers.set(key, factory(value, key));
			}
		}
	};

	const dispose = subscribe(atom, update);

	update(atom.get());

	return () => {
		dispose();
		disposers.forEach((dispose) => dispose());
	};
}

// Demo

const roundStatusAtom = atom("waiting");
const roundCounterAtom = atom(0);
const roundAtom = derive(roundStatusAtom, roundCounterAtom, (status, counter) => ({
	status,
	counter,
}));

function setRoundStatus(status: string) {
	roundStatusAtom.set(status);
}

function incrementRoundCounter() {
	roundCounterAtom.set((counter) => counter + 1);
}

subscribe(roundAtom, (round) => {
	print(`Round: ${round.status} (${round.counter})`);
});

setRoundStatus("running");
incrementRoundCounter();