Support declarative API in React with lazyGet and set methods

[DonnyVerduijn]

Currently, we expose a useLocalObservable hook from mobx-react-lite, but after using it shortly, i realized that it was far from ideal, having to call any defined methods inside the hook, from either arrow functions passed to props or using React.useCallback, to prevent this being bound to the call site. To solve this problem, i decided to autobind methods in a custom hook i called useMobx.

That solved some of the problems, but luckily not all. I soon realized that i wanted to use late dereferencing as much as possible and took some inspiration from the docs, by doing prop drilling "render prop" style, by providing a reference to defined functions in the hook, as props to child components. However, it still required me to define separate functions for every property in this fashion.

So, then it came to my mind that we could get rid of these methods provided to useLocalObservable and attach memoized set and lazyGet methods to the observable returned from the hook, which would provide stable references across re-renders, yet provide the functions to the leaf component to do late dereferencing. And so i was able to create something, from which i believe it deserves a place in mobx-react-lite.

import * as React from 'react';

enum PowerState {
  OFF = 'OFF',
  ON = 'ON',
}

interface LightBulbState {
  status: { power: PowerState };
}

const defaultState: LightBulbState = {
  status: { power: PowerState.OFF };
};

// nothing is dereferenced in this component,
// so we don't need an observer.

const LightBulb: React.FC = () => {
  const bulb = useMobx(() => defaultState);
  return (
    <OnOffSwitch
      getValue={bulb.lazyGet('status.power', (value) =>
        value === PowerState.ON ? true : false
      )}
      onChange={bulb.set('status.power', (value) =>
        value ? PowerState.ON : PowerState.OFF
      )}
    />
  );
};

Note that in the optional map function that we provide as the second argument to lazyGet and set, the type of the parameter is correctly inferred by the associated prop type or property key. If omitted, the identity function is applied here. The property key provided as the first argument, is itself typed as a union of all possible dot-separated paths, which matches the implementation that parses and traverses the path to get the observable value.

In this specific example, OnOffSwitch (or any other component), is assumed to have a getValue prop and is wrapped by observer, calling the function that gets provided to its getValue prop.

import { observer } from 'mobx-react-lite';
import Switch from '@mui/material/Switch';

interface Props {
  readonly getValue: () => boolean;
  readonly onChange: (value: boolean) => void;
}

const OnOffSwitch: React.FC<Props> = observer(({ getValue, onChange }) => {
  const handleChange = React.useCallback<
    (e: React.ChangeEvent<HTMLInputElement>, v: boolean) => void
  >((_, v) => onChange(v), [onChange]);

  return (
    <Switch
       checked={getValue()}
       onChange={handleChange}
    />
  );
});

The exact implementation of the hook could vary, given the dependency on remeda and hotscript in the example below. Memoized calls to lazyGet and set effectively eliminate the need for React.useCallback in components, as we rely on the stable reference returned from the hook.

import React, { useState } from 'react';
import { Call, O, S } from 'hotscript';
import { AnnotationsMap, observable, runInAction } from 'mobx';
import { identity, isFunction } from 'remeda';
import { memoize } from 'common/utils/memoize';

export const useMobx = <T extends Record<string, any>>(
  initialize: () => T,
  annotations?: AnnotationsMap<T, never>
) =>
  useState(() => {
    const obs = observable(initialize(), annotations, {
      autoBind: true,
    });

    // split a dot separated path into its constituents,
    // and type it as a tuple of literals.
    const getKeys = <TPath extends string>(path: TPath) => [
      ...(path.split('.') as Call<S.Split<'.'>, typeof path> & string[]),
    ];

    // given an observable object and array of path constituents,
    // get the parent of the target.
    const getParent = <TParent>(
      observable: Record<string, unknown>,
      keys: string[]
    ): TParent => {
      let target: unknown = observable;
      for (let i = 0; i < keys.length - 1; i++) {
        if (!Object.prototype.hasOwnProperty.call(target, keys[i])) {
          throw new Error('Invalid path');
        }
        // @ts-expect-error target changes so is unknown
        target = target[keys[i]];
      }
      // target contains the parent reference.
      return target as TParent;
    };

    type LazyGetFn<TType> = <
      P extends Call<O.AllPaths, TType>,
      V extends Call<O.Get<P>, TType>,
      R,
    >(
      path: P,
      map?: (value: V) => R
    ) => () => R;

    const lazyGet: LazyGetFn<T> = memoize((path, map = identity) => {
      // by using map from the closure, V8 will only preparse it on re-renders
      const keys = getKeys(path);
      const parent = getParent<Parameters<typeof map>[0]>(obs, keys);
      const key = keys[keys.length - 1];
      // Note how we keep the parent reference in the closure that is kept alive,
      // until the component that calls useMobx unmounts.
      // This is important, to avoid traversing every re-render.
      return () => map(parent[key] as Call<O.Get<typeof path>, T>);
    });

    type SetFn<TType> = <
      P extends Call<O.AllPaths, TType>,
      S extends Call<O.Get<P>, TType>,
      V,
    >(
      path: P,
      map?: (value: V, state: S) => S
    ) => (value: V) => void;

    const set: SetFn<T> = memoize(
      // we currently don't typecheck without the map function,
      // although it would be great.
      (path, map = (value, state) => value as unknown as typeof state) => {
        const keys = getKeys(path);
        const parent = getParent<Parameters<typeof map>[0]>(obs, keys);
        const key = keys[keys.length - 1];

        return (value) =>
          runInAction(() => {
            parent[key] = map(
              value,
              parent[key] as Call<O.Get<typeof path>, T>
            );
          });
      }
    );

    return Object.assign(obs, {
      lazyGet,
      set,
    });
  })[0];

For reference, here is a memoize implementation to complement the example. Note that it stringifies all of the arguments, so it is not very efficient or safe in all circumstances.

export const memoize = <T extends (...args: any[]) => R | Promise<R>, R>(
  targetFn: T
) => {
  const cache: Map<string, R> = new Map();

  const memoizedFunction = (...args: Parameters<T>) => {
    const key: string = JSON.stringify(args);
    if (cache.has(key)) {
      return cache.get(key)!;
    }

    const result = targetFn.apply(null, args);
    if (result instanceof Promise) {
      return result
        .then((resolvedResult) => {
          cache.set(key, resolvedResult);
          return resolvedResult;
        })
        .catch((error: unknown) => {
          cache.delete(key);
          throw error;
        });
    } else {
      cache.set(key, result);
      return result;
    }
  };
  // Parameters<T> widens the parameter,
  // therefore we need to cast to T
  return memoizedFunction as T;
};

Please give it a try and let me know what you think. If it is not feasible to include it in mobx-react-lite, i would love to have a discussion about the potential problems at scale with this approach. I believe that ideally, we should track if observables stay unobserved after component unmounts, to do some garbage collection along the way. Shamelessly tagging @mweststrate here.

EDIT: After giving it some more thought, i think one of the important considerations is how to deal with arrays. I would say that indexed lookup is out of question and arrays should be considered atomic units. The rationale behind this would be, that we don't want memoized functions for indexes that might be removed. Now the problem that appears, is how we gain access to the same lazyGet and set methods after iterating over an array. Given we deal with observable objects again that aren't extended, the question comes up if we rather need to extend any prototypes with these set and lazyGet methods?

Another question is how we want to support changing multiple properties at once. Personally i don't feel we should recommend this api just for single property lookups/writes if it can't support multiple, so that is also something to think about.