Reforest jotai sketch

atoms.tsx

import * as React from "react"
import { atom, useAtom, useAtomValue } from "jotai"

import "./styles.css"

const isServer = typeof window === "undefined"
const useIsomorphicLayoutEffect = isServer
  ? React.useEffect
  : React.useLayoutEffect

const TreeContext = React.createContext<any>(null)

function Parent({ children }: { children: React.ReactNode }) {
  const [treeMapAtom] = React.useState(() => atom(new Map()))
  const [computedTreeMapAtom] = React.useState(() => atom(new Map()))
  const treeNodeAtoms = React.useMemo(
    () =>
      atom((get) =>
        Array.from(get(treeMapAtom).values()).map((atom) => get(atom)),
      ),
    [treeMapAtom],
  )
  const [treeNodes] = useAtom(treeNodeAtoms)
  const totalDuration = treeNodes.reduce(
    (total, node: any) => total + node.duration,
    0,
  )
  const contextValue = React.useMemo(
    () => ({
      computedTreeMapAtom,
      treeMapAtom,
      treeNodeAtoms,
      totalDuration,
    }),
    [computedTreeMapAtom, treeMapAtom, treeNodeAtoms, totalDuration],
  )

  return (
    <>
      <div>Total Duration: {totalDuration}</div>
      <div style={{ display: "flex", gap: "1rem" }}>
        <TreeContext.Provider value={contextValue}>
          {children}
        </TreeContext.Provider>
      </div>
    </>
  )
}

function Child({ color, duration }: { color: string; duration: number }) {
  const treeContext = React.useContext(TreeContext)
  const [treeMap, setTreeMap] = useAtom(treeContext.treeMapAtom)
  const treeDataAtom = React.useMemo(() => atom({ color, duration }), [
    color,
    duration,
  ])

  useIsomorphicLayoutEffect(() => {
    const id = treeDataAtom.toString()

    setTreeMap((currentMap: Map<string, any>) => {
      const nextMap = new Map(currentMap)
      nextMap.set(id, treeDataAtom)
      return nextMap
    })

    return () => {
      setTreeMap((currentMap: Map<string, any>) => {
        const nextMap = new Map(currentMap)
        nextMap.delete(id)
        return nextMap
      })
    }
  }, [treeDataAtom])

  // useIsomorphicLayoutEffect(() => {
  //   const id = treeDataAtom.toString()

  //   setTreeMap((currentMap: Map<string, any>) => {
  //     const nextMap = new Map(currentMap)
  //     nextMap.set(id, treeDataAtom)
  //     return nextMap
  //   })

  //   return () => {
  //     setTreeMap((currentMap: Map<string, any>) => {
  //       const nextMap = new Map(currentMap)
  //       nextMap.delete(id)
  //       return nextMap
  //     })
  //   }
  // }, [treeDataAtom])

  return (
    <div
      style={{
        display: "grid",
        padding: 16,
        backgroundColor: color,
        color: "white",
      }}
    >
      <div>Duration: {duration}</div>
      <div>Count: {treeMap.size}</div>
    </div>
  )
}

export default function App() {
  const [showChild, setShowChild] = React.useState(false)

  return (
    <>
      <button onClick={() => setShowChild((bool) => !bool)}>
        Toggle Child
      </button>
      <React.Suspense fallback={null}>
        <Parent>
          <Child duration={3} color="green" />
          <Child duration={1.5} color="blue" />
          {showChild ? <Child duration={4} color="pink" /> : null}
          <Child duration={2} color="orange" />
        </Parent>
      </React.Suspense>
    </>
  )
}