An effection state persistence utility

state.ts

import { type Provide, ensure, resource } from "effection";
import { type Draft, type Patches, create } from "mutative";
import { match } from "ts-pattern";

export interface KVStore {
	get(key: string): Operation<unknown>;
	put(key: string, value: unknown): Operation<void>;
	putBatch(entries: Record<string, unknown>): Operation<void>;
	delete(key: string): Operation<void>;
	deleteBatch(keys: string[]): Operation<void>;
	list(options: { prefix?: string }): Operation<Map<string, unknown>>;
}
export const KVContext = createContext<KVStore>("kv-store");

function useImmutableState<T>(
	initial: T,
	onStateFinalized: (
		ops: [
			state: T,
			patches: Patches<{ pathAsArray: false }>,
			inversePatches: Patches<{ pathAsArray: false }>,
		],
	) => void,
) {
	return resource(function* (provide: Provide<Draft<T>>) {
		const [draft, finalize] = create(initial, {
			enablePatches: { pathAsArray: false },
		});
		yield* ensure(() => onStateFinalized(finalize()));
		yield* provide(draft);
	});
}

export function* withState<S extends Record<string, unknown>>(
	state: S,
	stateChanged?: () => void,
) {
	return yield* useImmutableState<S>(state, function* ([newState, patches]) {
		const kv = yield* KVContext.get();
		if (!kv) {
			throw new Error("KV store not found");
		}
		const putEntries: Record<string, unknown> = {};
		const deleteKeys: string[] = [];

		for (const patch of patches) {
			match(patch as PatchOp)
				.with({ op: "add" }, ({ path, value }) => {
					putEntries[`:state:${path}`] = value;
				})
				.with({ op: "remove" }, ({ path }) => {
					deleteKeys.push(`:state:${path}`);
				})
				.with({ op: "replace" }, ({ path, value }) => {
					putEntries[`:state:${path}`] = value;
				});
		}

		let hasStateChanged = false;
		if (Object.keys(putEntries).length > 0) {
			yield* kv.putBatch(putEntries);
			hasStateChanged = true;
		}
		if (deleteKeys.length > 0) {
			yield* kv.deleteBatch(deleteKeys);
			hasStateChanged = true;
		}
		if (hasStateChanged && stateChanged) {
			stateChanged();
		}

		// Assign new state properties
		Object.assign(state, newState);
		// Remove keys that don't exist in newState
		for (const key of Object.keys(state)) {
			if (!(key in newState)) {
				delete state[key];
			}
		}
	});
}