feat(core): unstable_createStore (pmndrs#922)

* wip: unstable_createStore

* get should not return undefined

* get to return a promise

* return promise only if enabled with options

* never throw on suspend

* unstable_createStore prop in Provider

* asyncGet store method

src/core/Provider.ts

@@ -15,17 +15,23 @@ import {
   DEV_GET_MOUNTED,
   DEV_GET_MOUNTED_ATOMS,
   DEV_SUBSCRIBE_STATE,
+  createStoreForExport,
 } from './store'
 import type { AtomState, Store, VersionObject } from './store'
 
 export const Provider = ({
   children,
   initialValues,
   scope,
+  unstable_createStore,
   unstable_enableVersionedWrite,
 }: PropsWithChildren<{
   initialValues?: Iterable<readonly [Atom<unknown>, unknown]>
   scope?: Scope
+  /**
+   * This is an unstable feature to use exported createStore.
+   */
+  unstable_createStore?: typeof createStoreForExport
   /**
    * This is an unstable experimental feature for React 18.
    * When this is enabled, a) write function must be pure
@@ -49,7 +55,10 @@ export const Provider = ({
   const scopeContainerRef = useRef<ScopeContainer>()
   if (!scopeContainerRef.current) {
     // lazy initialization
-    scopeContainerRef.current = createScopeContainer(initialValues)
+    scopeContainerRef.current = createScopeContainer(
+      initialValues,
+      unstable_createStore
+    )
     if (unstable_enableVersionedWrite) {
       scopeContainerRef.current.w = (write) => {
         setVersion((parentVersion) => {

src/core/contexts.ts

@@ -1,7 +1,7 @@
 import { createContext } from 'react'
 import type { Context } from 'react'
 import type { Atom, Scope } from './atom'
-import { createStore } from './store'
+import { createStore, createStoreForExport } from './store'
 import type { Store } from './store'
 
 type VersionedWrite = (write: (version?: object) => void) => void
@@ -12,9 +12,12 @@ export type ScopeContainer = {
 }
 
 export const createScopeContainer = (
-  initialValues?: Iterable<readonly [Atom<unknown>, unknown]>
+  initialValues?: Iterable<readonly [Atom<unknown>, unknown]>,
+  unstable_createStore?: typeof createStoreForExport
 ): ScopeContainer => {
-  const store = createStore(initialValues)
+  const store = unstable_createStore
+    ? unstable_createStore(initialValues).SECRET_INTERNAL_store
+    : createStore(initialValues)
   return { s: store }
 }
 

src/core/store.ts

@@ -614,7 +614,7 @@ export const createStore = (
     version: VersionObject | undefined,
     atom: WritableAtom<Value, Update, Result>,
     update: Update
-  ): void | Promise<void> => {
+  ): Result => {
     let isSync = true
     const writeGetter: WriteGetter = <V>(
       a: Atom<V>,
@@ -687,7 +687,7 @@ export const createStore = (
     writingAtom: WritableAtom<Value, Update, Result>,
     update: Update,
     version?: VersionObject
-  ): void | Promise<void> => {
+  ): Result => {
     const promiseOrVoid = writeAtomState(version, writingAtom, update)
     flushPending(version)
     return promiseOrVoid
@@ -901,3 +901,43 @@ export const createStore = (
 }
 
 export type Store = ReturnType<typeof createStore>
+
+export const createStoreForExport = (
+  initialValues?: Iterable<readonly [AnyAtom, AnyAtomValue]>
+) => {
+  const store = createStore(initialValues)
+  const get = <Value>(atom: Atom<Value>) => {
+    const atomState = store[READ_ATOM](atom)
+    if ('e' in atomState) {
+      throw atomState.e // read error
+    }
+    if ('p' in atomState) {
+      return undefined // suspended
+    }
+    return atomState.v
+  }
+  const asyncGet = <Value>(atom: Atom<Value>) =>
+    new Promise<ResolveType<Value>>((resolve, reject) => {
+      const atomState = store[READ_ATOM](atom)
+      if ('e' in atomState) {
+        reject(atomState.e) // read error
+      } else if ('p' in atomState) {
+        resolve(atomState.p.then(() => asyncGet(atom))) // retry later
+      } else {
+        resolve(atomState.v)
+      }
+    })
+  const set = <Value, Update, Result extends void | Promise<void>>(
+    atom: WritableAtom<Value, Update, Result>,
+    update: Update
+  ) => store[WRITE_ATOM](atom, update)
+  const sub = (atom: AnyAtom, callback: () => void) =>
+    store[SUBSCRIBE_ATOM](atom, callback)
+  return {
+    get,
+    asyncGet,
+    set,
+    sub,
+    SECRET_INTERNAL_store: store,
+  }
+}

src/index.ts

@@ -1,6 +1,7 @@
 export { Provider } from './core/Provider'
 export { atom } from './core/atom'
 export { useAtom } from './core/useAtom'
+export { createStoreForExport as unstable_createStore } from './core/store'
 export type { Atom, WritableAtom, PrimitiveAtom } from './core/atom'
 export type {
   Getter,