Idea: lazy initialization for primitive atoms

[dai-shi]

Idea: lazy initialization for primitive atoms

[dmaskasky]

Interesting idea. How about atomWithLazy?

[dai-shi]

I was thinking about using a LAZY symbol, so atomWithLazy would be possible too. The point is that we need a new capability in core store impl.

[iwoplaza]

Would it involve changes to the store implementation, or would it be a handy wrapper for the existing derived workaround?

Here is what I mean by a wrapper:

export function atomWithLazy<Value>(makeInitial: () => Value) {
  const wrappedAtom = atom(() => {
    // an atom is created per store
    return atom(makeInitial())
  })

  return atom((get) => get(get(wrappedAtom)))
}

Since wrappedAtom does not have any dependencies, it will compute its value only when first used, and the value happens to be a primitive atom that can then be written to and updated at will.

Then we could use it like initially proposed:

import { atomWithLazy } from 'jotai/vanilla'; // or 'jotai'

const a = atomWithLazy(() => Date.now()); // will be evaluated on the first use in the store.

[dai-shi]

Very interesting. We don't need anything new in the store. How nice.

export function atomWithLazy(makeInitial) {
  const wrappedAtom = atom(() => atom(makeInitial()))
  return atom(
    (get) => get(get(wrappedAtom)),
    (get, set, ...args) => set(get(wrappedAtom), ..args)
  )
}

[dai-shi]

@iwoplaza Would you like to open a PR? It should be in utils.

[iwoplaza]

Sure thing, I will do that either later today or tomorrow.

[rothsandro]

The only drawback of the atom in atom approach is that it will break when used in combination with Jotai Scope. Lazy atoms used outside and inside of a Jotai Scope Provider without being a scoped atom will unexpectedly be created twice.

[dai-shi]

Could it solved with unstable_is? 🤔

[iwoplaza]

I wrote a reproduction of the issue with jotai-scope in this CodeSandbox.

It seems that as long as the atom that is returned from atomWithLazy thinks it is writing to itself, instead of the wrappedAtom, it works as intended.

export function atomWithLazy<Value>(makeInitial: () => Value) {
  const wrappedAtom: Atom<PrimitiveAtom<Value>> & { init?: Atom<undefined> } =
    atom(() => atom(makeInitial()));

  const proxyAtom = atom(
    (get) => get(get(wrappedAtom)),
    (get, set, value: SetStateAction<Value>) => set(get(wrappedAtom), value)
  );

  wrappedAtom.init = atom(undefined);
  // when writing to wrappedAtom through proxyAtom, the store
  // thinks we are actually storing the value of `proxyAtom`.
  wrappedAtom.unstable_is = (a) =>
    a.unstable_is?.(proxyAtom) ?? a === proxyAtom;

  return proxyAtom;
}

[dai-shi]

Wow, that seems like a great hack.

Some minor suggestions:

// does this work?
wrappedAtom.init = undefined;
// just a syntax preference
wrappedAtom.unstable_is = (a) =>
    a.unstable_is ? a.unstable_is(proxyAtom) : a === proxyAtom;

[iwoplaza]

Setting wrappedAtom.init to undefined causes a Cannot read properties of undefined (reading 'unstable_is'). I am currently unsure what in the store impl causes this, but it happens in the proxyAtom read function, equivalent to:

(get) => get(undefined)

I updated the lazyAtom implementation (along with the syntax preference) in the PR.

[iwoplaza]

Could we discuss on the possible use cases of this new utility? We can then list them in the documentation.
My first thought was that it could be used for resources available only after app initialization, not during, but it is a rather vague example. Any more concrete ideas?

[iwoplaza]

The current docs for atomWithLazy can be found here. Focused mostly on its ability to post-pone (or remove the need for) initialization of an atom in case the initialization itself is expensive.

[dai-shi]

Setting wrappedAtom.init to undefined causes a Cannot read properties of undefined (reading 'unstable_is').

Hmm, tbh, I'm not sure why .init is required.

I am currently unsure what in the store impl causes this, but it happens in the proxyAtom read function, equivalent to:

Can you reproduce it in the test? I'm fine with .init hack for now, but we may want to remove it in the future.

[dai-shi]

Could we discuss on the possible use cases of this new utility?

I have at least two use cases in my mind.

1. lazy evaluation of heavy computation.
2. avoiding sharing an object between different providers/stores.

[iwoplaza]

After crawling through the store implementation, I figured out an alternative implementation that could be a little more elegant and understandable:

export function atomWithLazy<Value>(makeInitial: () => Value) {
  return {
    // create a regular primitive atom with a placeholder initial value
    ...atom(undefined as unknown as Value),
    // hijack the `init` property into being computed on first use
    get init() {
      return makeInitial();
    },
  };
}

This removes the atom-in-atom problem, but introduces extraneous computations when used alongside jotai-scope. To be more precise, without the use of jotai-scope, the makeInitial is called exactly once per store as expected. When the atom is used within a ScopeProvider, the makeInitial function is called exactly 3 times, no matter how many times we actually scope the atom breadth-wise (with nested ScopeProviders the number of invocations will probably be higher). My current hypothesis is that makeInitial is invoked when a lazy atom gets intercepted (spread into a new object), when the store first initializes the intercepted atom (.init), and when the store first initializes the unintercepted atom (.init). If we are okay with this overhead (3x instead of 1x for the case of using jotai-scope), then that seems to be the only drawback.

[dai-shi]

    get init() {
      return makeInitial();
    },

What a neat solution!

If we are okay with this overhead (3x instead of 1x for the case of using jotai-scope), then that seems to be the only drawback.

I think it's a reasonable limitation. It's not the best for performance, but the capability is preserved, right? If it's noted as a caveat in jotai-scope, it should be fine. @yf-yang

I figured out an alternative implementation that could be a little more elegant and understandable

Yes, let's go with it.