Recursive AST atoms [Question]

[loganvolkers]

At SaaSquatch we build a lot of AST editors libraries that edit a homogeneous tree of AST nodes of arbitrary depth. Jotai seems to be a good fit for a lot of the state management use cases, but I'm looking for some guidance on the best way of dealing with recursion.

Example JSON for an HTML AST

```json { "start": 0, "end": 108, "type": "Fragment", "children": [ { "start": 0, "end": 15, "type": "Element", "name": "!DOCTYPE", "attributes": [ { "start": 10, "end": 14, "type": "Attribute", "name": "html", "value": true } ], "children": [] }, { "start": 15, "end": 16, "type": "Text", "raw": "\n", "data": "\n" }, { "start": 16, "end": 108, "type": "Element", "name": "html", "attributes": [], "children": [ { "start": 22, "end": 24, "type": "Text", "raw": "\n\n", "data": "\n\n" }, { "start": 24, "end": 99, "type": "Element", "name": "body", "attributes": [], "children": [ { "start": 30, "end": 35, "type": "Text", "raw": "\n ", "data": "\n " }, { "start": 35, "end": 60, "type": "Element", "name": "h1", "attributes": [], "children": [ { "start": 39, "end": 55, "type": "Text", "raw": "My First Heading", "data": "My First Heading" } ] }, { "start": 60, "end": 65, "type": "Text", "raw": "\n ", "data": "\n " }, { "start": 65, "end": 91, "type": "Element", "name": "p", "attributes": [], "children": [ { "start": 68, "end": 87, "type": "Text", "raw": "My first paragraph.", "data": "My first paragraph." } ] }, { "start": 91, "end": 92, "type": "Text", "raw": "\n", "data": "\n" } ] }, { "start": 99, "end": 101, "type": "Text", "raw": "\n\n", "data": "\n\n" } ] } ] } ```

Background

Option 1 - Build and pass atoms

My current approach is based on the large objects guide. Each node is a React component that is responsible for creating and memoizing sub-atoms.

function NodeEditorPassDown({ nodeAtom }) {
  const childrenAtom = useMemo(
    () => atom((get) => get(nodeAtom)?.children ?? []),
    [nodeAtom]
  );
  const childrenAtomsAtom = splitAtom(childrenAtom);
  const [childrenAtoms] = useAtom(childrenAtomsAtom);
  const [node] = useAtom(nodeAtom);
  return (
    <div>
      Type: {node.type}
      <hr />
      {childrenAtoms.map((a) => (
        <NodeEditorPassDown nodeAtom={a} key={a} />
      ))}
    </div>
  );
}

Option 2 - Using scope / provider

I haven't been able to figure out if this is a good idea or not. The idea is to simplify working on nested editors, since they should be able to know about atoms in their context. The idea here is that other things can connect into the sub-atom state, without having to touch the React hooks.

My big unanswered questions are:

• Is this an abuse of Provider?
• Is React's Context a better fit?
• Is there a way of scoping an atom's get method to read from a different store?
• Is there another way of scoping and connecting stores?

Here's a non-working example that's trying to demonstrate the scoping/context idea.

const NodeScope = Symbol();

const rootAtom = atom(root);
const scopeAtom = atom(rootAtom);

function NodeEditorScopeAware() {
  const nodeAtom = useAtom(scopeAtom, NodeScope);
  if (!nodeAtom) throw new Error("Scope missing");
  const childrenAtom = useMemo(
    () => atom((get) => get(nodeAtom)?.children ?? []),
    [nodeAtom]
  );
  const childrenAtomsAtom = splitAtom(childrenAtom);
  const [childrenAtoms] = useAtom(childrenAtomsAtom);
  const node = useAtom(nodeAtom);
  return (
    <>
      Type: {node.type}
      {childrenAtoms.map((a) => (
        <Provider scope={NodeScope} initialValues={[[scopeAtom, a]]}>
          <NodeEditorScopeAware />
        </Provider>
      ))}
    </>
  );
}
const Component = () => {
  return (
    <Provider scope={NodeScope}>
      <NodeEditor />
    </Provider>
  );
};

[dai-shi]

Yeah, this is still a field we don't have specific recommendations. Happy to have you on board.

First of all, option 1 is good as written in a guide. I'm not very satisfied with it either because it doesn't feel like jotai's way, and technically splitAtom is super performant. (Don't get it wrong. It should work great. My perspective is from the library implementation.)

Quoting your questions:

• Is this an abuse of Provider?

Using nested Providers is a valid pattern, which can be said underrated.
So, it looks like a good approach, if it works.
Whether you need scope is questionable. (The provider scope is basically creating a new React Context. Is it what you would do with pure React Context?)

• Is React's Context a better fit?

Jotai is absolutely developed on top of Context. Basically, all Context features are available in jotai.
So, if you can solve your use case with React Context, that's good.
You can keep using pure React Context, or we can try improving it with jotai atoms.

• Is there a way of scoping an atom's get method to read from a different store?

I'm not sure if I understand the goal. So, you want to avoid reading a node that is not in the children?

• Is there another way of scoping and connecting stores?

I would explore atoms-in-atom pattern.

[loganvolkers]

Thanks for the quick response!

technically splitAtom is super performant

I assume you meant "isn't super performant"? (Not that I have noticed any performance problems)

Using nested Providers is a valid pattern, which can be said underrated. So, it looks like a good approach, if it works.

I've investigated this, and I don't think this can work:

• There are "node-scoped" atoms, state that is derived from each node
• There are "global-scoped" atoms, such as the root object and selection tracking
• If I used a Provider for this, then there is no way to connect atoms across scopes (without writing a lot of hooks to be the glue)

Here's a diagram showing where the connection would break.

I'm not sure if I understand the goal. So, you want to avoid reading a node that is not in the children?

This would be a solution to the above problem, which could looks something like this.

const nodeSelected = atom(get => get(nodeAtom, nodeScope) === get(selectedAtom, globalScope))

I don't see this being a path forward for Jotai, since it seems like it would add bloat and complication to the core, but it would allow for "cross-scope dependencies".

I would explore atoms-in-atom pattern.

This is super useful, and is what I'm already doing in Option 1 (e.g. splitAtom). This would be good for a "edit selected node", but can't really work with recursion.

Jotai is absolutely developed on top of Context. Basically, all Context features are available in jotai.
So, if you can solve your use case with React Context, that's good.
You can keep using pure React Context, or we can try improving it with jotai atoms.

There are a bunch of helpers here, but they all require custom hook code. Ideally I could use context directly in the atom tree, but that's not possible since even a derived atom based on the atoms-in-atom pattern and react context can still only have a single value at a time (e.g. edit selected node), but we need multiple children to be viewable and editable simultaneously (e.g. edit all nodes).

Yeah, this is still a field we don't have specific recommendations. Happy to have you on board.

Here are my takeways and my current working theory:

• Using React hooks to dynamically create, cache, and prop-drill atoms is the most obvious way to power recursion
• React Context can simplify the amount of prop drilling required
• createMemoizeAtom can reduce the amount of custom useMemo hooks are required (side note that createMemoizeAtom should be exported)
• scope can't be used for recursion if you need to connect atoms across scopes, at least not with the public core API

[dai-shi]

I assume you meant "isn't super performant"?

Oops, you are right. Reading this again, it was a bit misleading.
I would rephrase it as splitAtom is a complicated util implementation-wise.

there is no way to connect atoms across scopes

right, scopes are not designed for this use cases.

I hope this should be solved in atom-in-atom pattern with or without nested providers.

<Provider key={parentAtom} initialValues={[[parentAtomAtom, parentAtom]]}>

On second thought, as your AST already exists outside React, option 1 seems fairly reasonable.
Hm, option 2 still uses splitAtom, so what's the goal of option 2 that can't be accomplished by option 1?

[loganvolkers]

what's the goal of option 2 that can't be accomplished by option 1?

Option 1 is working, but the boilerplate is a bit cumbersome.

Option 2 would be about reducing prop drilling by being able to access your "in scope" node atom as a statically referenceable constant. It's about simplifying and making the node atom tree (do we say molecule?) statically defined instead of created dynamically.

This is possible for the global scoped atoms. This is possible for node scoped atoms. But if you need to connect these two molecules, you need hooks to allow recursion.

[dai-shi]

Option 2 would be about reducing prop drilling by being able to access your "in scope" node atom as a statically referenceable constant.

Then, I'm pretty sure one good solution is: <Provider initialValues={...}>.
It's like what you wanted to do with "scope", but with "atom-in-atom" pattern.

[loganvolkers]

The problem with Provider is that it creates a new store, so it doesn't allow connecting atoms across the boundary. This prevents edits propagating up, and derived atoms that use both local and root atoms.

To be fair, hooks don't support recursion without adding some sort of wrapping component to act as the glue, too.

[loganvolkers]

I've got an example here: https://codesandbox.io/s/recursive-jotai-pfv5i?file=/src/useScopedAtom.tsx

I believe that since the atoms from splitAtom are stable, then there shouldn't be re-renders on the useContext and useMemo glue code, and it doesn't require createMemoizedAtom from weakCache.

The basic interface is the ScopedAtomCreator that allows nested editors to be implemented without prop drilling.

type ScopedAtomCreator<R, W> = (
  atom: Atom<unknown>
) => WritableAtom<R, W> | Atom<R>;

I'll probably want to remove the word "scope" and rename to something like "ContextualAtomCreator" and "useContextualAtom".

Is there something performance-wise that might be missing here? Is there an reference equality landmine? Any ideas on how to improve the API surface area?

[dai-shi]

The problem with Provider is that it creates a new store

Oh, you are absolutely right. I was missing that point. Then, Provider scope might come and play.

[dai-shi]

Thanks for the example!
Here's what I meant by Provider scope.
Basically, there's no difference in the approach. It just replaces React Context with jotai Provider.
https://codesandbox.io/s/recursive-jotai-forked-w7xbc?file=/src/useScopedAtom.tsx

Naming is not good either. Typing could be improved.

Is there something performance-wise that might be missing here? Is there an reference equality landmine?

No, it looks good.

[loganvolkers]

Nice!

I'll play around with the API a bit, but I think this does a good job of removing boilerplate, decoupling components and allowing recursion. Hopefully it's also efficient.

[loganvolkers]

Thanks for your help!

[dai-shi]

Good to hear!

This is probably the first time to see the real Provider scope use case other than the library use case.

[RikuVan]

I had a similar kind of design decision as @loganvolkers with a complex recursive structure: use atoms in atoms or providers. I tried both and this discussion helped. I find the provider with scopes to more flexible because I can then have the atom in module scope and pass it around freely, unlike an atom that is passed as a prop. It might be nice at some point if an example of complex provider usage with scopes was added to the docs.

[RikuVan]

Still one question or problem with this approach is still tripping me up. I have many derived atoms. Now that I have nested scopes, these are not working. So I have get(atomFromOuterScope) and get(atomFromInnerScope) in a derived atom (previously used atoms in atoms so it worked). It seems this is impossible since I cannot give a scope to get and I would need to copy atom values across scopes to use them in derived atoms. Is there any better solution to this?

[dai-shi]

scopes are not tied to atoms. The real issue is that a derived atom can't access another atom in a different scope, by design (otherwise, it's not "scoping"). Now, I know there are some use cases that we want to access atom values across providers, but not it's possible. If your use case falls into this, maybe the scope is not the right tool for you. If we really need to combine two atom values from different scopes, only the solution is doing in a hook (or a component).

[dai-shi]

On a side note, I would like to explore a pattern like atoms-in-atom, that can work for recursive structure, without providers.

[RikuVan]

I have returned to the atoms-in-atom structure so I could get rid of the inner provider and just have a single one around the main component. I now have a list of trees, a single atom, and each tree is an atom and each tree node is now an atom. Scoping seemed easier to reason about at first but the end result was more complex--once I embraced passing atoms as props more, it turned out to be more verbose but familiar and React-ish and it is much easier to use derived atoms and callbacks in atoms this way to handle updates at different levels. Btw one way in Svelte I have handled similar challenges is by passing Svelte stores around via context to avoid prop drilling. Can you put pass an atom itself via React.Context? --don't need this currently really just wondering if it is possible.

[dai-shi]

I'm interested in how you did the atoms-in-atom structure recursively.

Yeah, inner provider can be hard to deal with in this case. It works well if we create an isolated component or a library.

Passing an atom itself via React Context is totally fine (as well as via props).
I believe it's what @loganvolkers originally did.

[RikuVan]

In my case it is quite simple, the container component is recursive but the data is flattened to a tree of atom nodes, each with a parent id.

There is a parent Container component (ul) which is rendered recursively by the Node (li) component. This creates a tree of drag and droppable nodes. There is actually a list of these, so mulitple trees.

<ul ref={ref} className={className} key={parentId}>
      {view.map((node, index) => (
          <React.Fragment key={node.id}>
              <Placeholder depth={depth} listCount={view.length} dropTargetId={parentId} index={index} />
              <Node id={node.id} depth={depth} treeIndex={treeIndex} treeAtom={treeAtom} rootId={rootId} />
          </React.Fragment>
      ))}
      <Placeholder depth={depth} listCount={view.length} dropTargetId={parentId} />
</ul>

 <li ref={nodeRef} className={className}>
        {render ? render(item, params, handleRef) : <DefaultNodeContent {...params} ref={handleRef} />}
        <AnimatePresence initial={false}>
            {Boolean(open && item.hasChildren) && (
                <motion.div
                    key={id}
                    initial="exit"
                    animate="enter"
                    exit="exit"
                    variants={collapsibleVariants}
                    transition={{ ...collapsibleTransition, duration: 0.3 }}
                >
                    <Container parentId={id} depth={depth + 1} treeAtom={treeAtom} treeIndex={treeIndex} />
                </motion.div>
            )}
       </AnimatePresence>
 </li>

[loganvolkers]

Solved this with a new library https://github.com/saasquatch/jotai-molecules

There are a lot of tests and examples. Thanks for the help @dai-shi in figuring out how to create non-global atoms in some good ways.

[dai-shi]

It looks very nice!!

Molecule scopes can be interconnected

You did it! ❤️

cc: @aulneau may find this interesting.