feat(core): experimental unstable versioned write for limited use case (#854)

* feat(core): experimental version object

* refactor

* fix test name in ci

* revert test/query

* empty commit

* fix workflow file

* remove delay in urql/atomWithQuery

* fix test

* tweak query test

* fix atomWithQuery in urql

* revert again query timing

* wip: non-working commit, trying another path

* another approach with provider which requires pure write

* fix logic, add tests, skip tests

* refactor, simplify

* wip: refactor

* skip a test

* skip one more test

* skip 16.8.6

* only notify changed atom state

* remove extra falsy check

* fix types

* wip: a workaround for now

* fix logic, skip one test for now

* a minor refactor

* add some notes

* suspense promise should not throw

Author: dai-shi

.github/workflows/test-multiple-versions.yml

@@ -24,7 +24,11 @@ jobs:
       - name: Test Default without Provider
         run: yarn test:ci
         env:
-          PROVIDER_LESS_MODE: 'true'
+          PROVIDER_MODE: 'PROVIDER_LESS'
+      - name: Test Default with Versioned Write
+        run: yarn test:ci
+        env:
+          PROVIDER_MODE: 'VERSIONED_WRITE'
 
   test_matrix:
     runs-on: ubuntu-latest
@@ -37,9 +41,10 @@ jobs:
           - 17.0.0
           - 18.0.0-rc.0
           - 0.0.0-experimental-aa8f2bdbc-20211215
-        mode: [withProvider, withoutProvider]
+        mode: [NORMAL, PROVIDER_LESS, VERSIONED_WRITE]
         testing: [default, alpha]
         exclude:
+          - { react: 16.8.6, mode: VERSIONED_WRITE }
           - { react: 16.8.6, testing: alpha }
           - { react: 16.9.0, testing: alpha }
           - { react: 17.0.0, testing: alpha }
@@ -71,4 +76,4 @@ jobs:
           yarn add -D react@${{ matrix.react }} react-dom@${{ matrix.react }}
           yarn test:ci
         env:
-          PROVIDER_LESS_MODE: ${{ matrix.mode == 'withoutProvider' }}
+          PROVIDER_MODE: ${{ matrix.mode }}

package.json

@@ -127,7 +127,7 @@
     "eslint": "eslint --fix '*.{js,json}' '{src,tests,benchmarks}/**/*.{ts,tsx}'",
     "eslint:ci": "eslint '*.{js,json}' '{src,tests,benchmarks}/**/*.{ts,tsx}'",
     "pretest": "tsc --noEmit",
-    "test": "jest && jest --setupFiles ./tests/setProviderLessMode.ts",
+    "test": "jest && jest --setupFiles ./tests/setProviderLessMode.ts && jest --setupFiles ./tests/setVersionedWriteMode.ts",
     "test:ci": "jest",
     "test:dev": "jest --watch --no-coverage",
     "test:coverage:watch": "jest --watch",

src/core/Provider.ts

@@ -10,25 +10,55 @@ import type { Atom, Scope } from './atom'
 import { createScopeContainer, getScopeContext } from './contexts'
 import type { ScopeContainer } from './contexts'
 import {
+  COMMIT_ATOM,
   DEV_GET_ATOM_STATE,
   DEV_GET_MOUNTED,
   DEV_GET_MOUNTED_ATOMS,
   DEV_SUBSCRIBE_STATE,
 } from './store'
-import type { AtomState, Store } from './store'
+import type { AtomState, Store, VersionObject } from './store'
 
 export const Provider = ({
+  children,
   initialValues,
   scope,
-  children,
+  unstable_enableVersionedWrite,
 }: PropsWithChildren<{
   initialValues?: Iterable<readonly [Atom<unknown>, unknown]>
   scope?: Scope
+  /**
+   * This is an unstable experimental feature for React 18.
+   * When this is enabled, a) write function must be pure
+   * (read function must be pure regardless of this),
+   * b) React will show warning in DEV mode,
+   * c) then state branching works.
+   */
+  unstable_enableVersionedWrite?: boolean
 }>) => {
+  const [version, setVersion] = useState<VersionObject>()
+  useEffect(() => {
+    if (version) {
+      ;(scopeContainerRef.current as ScopeContainer).s[COMMIT_ATOM](
+        null,
+        version
+      )
+      delete version.p
+    }
+  }, [version])
+
   const scopeContainerRef = useRef<ScopeContainer>()
   if (!scopeContainerRef.current) {
     // lazy initialization
     scopeContainerRef.current = createScopeContainer(initialValues)
+    if (unstable_enableVersionedWrite) {
+      scopeContainerRef.current.w = (write) => {
+        setVersion((parentVersion) => {
+          const nextVersion = parentVersion ? { p: parentVersion } : {}
+          write(nextVersion)
+          return nextVersion
+        })
+      }
+    }
   }
 
   if (
@@ -79,7 +109,7 @@ const stateToPrintable = ([store, atoms]: [Store, Atom<unknown>[]]) =>
 // We keep a reference to the atoms in Provider's registeredAtoms in dev mode,
 // so atoms aren't garbage collected by the WeakMap of mounted atoms
 const useDebugState = (scopeContainer: ScopeContainer) => {
-  const store = scopeContainer.s
+  const { s: store } = scopeContainer
   const [atoms, setAtoms] = useState<Atom<unknown>[]>([])
   useEffect(() => {
     const callback = () => {

src/core/contexts.ts

@@ -4,8 +4,11 @@ import type { Atom, Scope } from './atom'
 import { createStore } from './store'
 import type { Store } from './store'
 
+type VersionedWrite = (write: (version?: object) => void) => void
+
 export type ScopeContainer = {
   s: Store
+  w?: VersionedWrite
 }
 
 export const createScopeContainer = (