PoC of immutable state manipulation via simple lenses (and React integration for the lolz)

lenses.js

import { useState, useMemo, useCallback, Dispatch, SetStateAction } from "react";

type InitialState<S> = (() => S) | S;


type Nullable<T> = T|null|undefined;

type Lens<Store = unknown, Value = unknown> = (store: Nullable<Store>) => ({
    get(): Nullable<Value>;
    set(newValue: Nullable<Value>): Nullable<Store>;
});

const prop = <Store extends object, Key extends keyof Store>(key: Key) => (store: Nullable<Store>) => ({
    get: () => store?.[key] ?? null,
    set: (newValue: Store[Key]) => store ? {...store, [key]: newValue} : store,
});

const atIndex = (index: number) => <T extends unknown>(store: Nullable<T[]>) => ({
    get: () => store?.[index] ?? null,
    set: (newValue: T) => {
        if(!store) {
            return store;
        }

        const { length } = store;

        if(index > length) {
            return store;
        } else if(index === length) {
            return [...store, newValue];
        } else {
            const copy = [...store];
            copy[index] = newValue;
            return copy;
        }
    },
});

const compose = <Store extends object>(lenses: Lens[]) => {
    return (store: Store) => {
        return {
            get() {
                return lenses.reduce((value: unknown, lens: Lens) => {
                    return lens(value).get();
                }, store);
            },
            set(newValue: unknown) {
                if(lenses.length === 0) {
                    return store;
                }

                const lens = lenses[0];
                const firstLens = lens(store);
                const firstState = firstLens.get();

                const pairs: ([lens: ReturnType<Lens>, state: unknown])[] = [
                    [firstLens, firstState],
                ];

                for(let i = 1 ; i < lenses.length ; i += 1) {
                    const lens = lenses[i];
                    const [, state] = pairs[i - 1];

                    const newLens = lens(state);
                    const newState = newLens.get();
                    pairs.push([newLens, newState]);
                }

                return pairs.reduceRight((changedValue, [lens]) => {
                    return lens.set(changedValue);
                }, newValue);
            },
        };
    };
};

const store = {
    deep: {
        nested: {
            props: [
                {
                    key: "A",
                    value: "alpha",
                },
                {
                    key: "B",
                    value: "beta",
                },
            ],
        },
    },
};

const useLensState = <RootState extends object, Value>(lens: Lens<RootState, Value>, rootState: Nullable<RootState>, setRootState: Dispatch<SetStateAction<Nullable<RootState>>>) => {
    const lensObj = useMemo(() => lens(rootState), [lens, rootState]);
    const value = useMemo(() => lensObj.get(), [lensObj]);
    const setValue = useCallback((newValue: Nullable<Value>) => {
        const newRootState = lensObj.set(newValue);
        setRootState(newRootState);
    }, [lensObj, setRootState]);

    return [value, setValue];
};

const useLensable = <RootState extends object>(initialRootState: InitialState<Nullable<RootState>>) => {
    const [rootState, setRootState] = useState(initialRootState);
    const useLens = useCallback(<Value extends unknown>(lens: Lens<RootState, Value>) => {
        return useLensState(lens, rootState, setRootState);
    }, [rootState, setRootState]);
    return {
        rootState,
        setRootState,
        useLens
    };
};

const { useLens } = useLensable({
    deep: {
        nested: {
            props: [
                {
                    key: "A",
                    value: "alpha",
                },
                {
                    key: "B",
                    value: "beta",
                },
            ],
        },
    },
});

const alphaLens = compose([
    prop("deep") as Lens,
    prop("nested") as Lens,
    prop("props") as Lens,
    atIndex(0) as Lens,
    prop("value") as Lens,
]);

const [valueA, setValueA] = useLens<string>(alphaLens as Lens<typeof store, string>);

console.log("get(): ", alphaLens(store).get(), "store: ", store);
console.log("set(): ", alphaLens(store).set("megalpha"), "store: ", store);