Repository: pmndrs/valtio (ts)

Issue:
adding a `derive` function
I know there's addComputed but it'd be great to have a `derive` function based off of `subscribe` to create dependent proxies.
passing a `notifyInSync` flag to addComputed
I noticed that `addComputed` also uses subscribe.

https://github.com/pmndrs/valtio/blob/master/src/utils.ts#L192

Can a `notifyInSync` flag be passed into the function as well?

Pull request description:
feat: derive util
close #176 
close #198 

This will deprecate `addComputed` in favor of new `derive` and existing `proxyWithComputed`.

- [x] impl
- [x] tests
- [x] deprecate `addComputed` with fallback impl
- [x] docs

Patch:
diff --git a/.size-snapshot.json b/.size-snapshot.json
index d8d7b14..59c31c8 100644
--- a/.size-snapshot.json
+++ b/.size-snapshot.json
@@ -28,16 +28,16 @@
     }
   },
   "utils.js": {
-    "bundled": 7137,
-    "minified": 3388,
-    "gzipped": 1422,
+    "bundled": 10177,
+    "minified": 4646,
+    "gzipped": 1926,
     "treeshaked": {
       "rollup": {
         "code": 45,
         "import_statements": 45
       },
       "webpack": {
-        "code": 1062
+        "code": 1074
       }
     }
   },
diff --git a/readme.md b/readme.md
index 2db9e43..cba8cd1 100644
--- a/readme.md
+++ b/readme.md
@@ -243,47 +243,37 @@ const Component = () => {
 }
 ```
 
-#### Computed values
+#### `derive` util
 
-You can have computed values with dependency tracking.
-Dependency tracking in valtio conflicts with the work in useSnapshot.
-React users should consider using render functions (optionally useMemo)
-as a primary mean.
-Computed works well for some edge cases and for vanilla-js users.
-
-##### `addComputed`
-
-This is to add new computed to an existing proxy state.
-It can add computed to different proxy state.
+You can subscribe to some proxies and create a derived proxy.
 
 ```js
-import { addComputed } from 'valtio/utils'
+import { derive } from 'valtio/utils'
 
 // create a base proxy
 const state = proxy({
   count: 1,
 })
 
-// add computed to state
-addComputed(state, {
-  doubled: snap => snap.count * 2,
+// create a derived proxy
+const derived = derive({
+  doubled: (get) => get(state).count * 2,
 })
 
-// create another proxy
-const state2 = proxy({
-  text: 'hello',
+// alternatively, attach derived properties to an existing proxy
+derive({
+  tripled: (get) => get(state).count * 3,
+}, {
+  proxy: state,
 })
-
-// add computed from state to state2
-addComputed(state, {
-  doubled: snap => snap.count * 2,
-}, state2)
 ```
 
-##### `proxyWithComputed`
+#### `proxyWithComputed`
 
-This is to create a proxy state with computed at the same time.
-It can define setters optionally.
+You can have computed values with dependency tracking with property access.
+Dependency tracking in `proxyWithComputed` conflicts with the work in `useSnapshot`.
+React users should prefer using `derive`.
+`proxyWithComputed` works well for some edge cases and for vanilla-js users.
 
 ```js
 import { proxyWithComputed } from 'valtio/utils'
diff --git a/src/utils.ts b/src/utils.ts
index cad9986..3a14bac 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -1,6 +1,11 @@
 export { subscribeKey } from './utils/subscribeKey'
 export { watch } from './utils/watch'
 export { devtools } from './utils/devtools'
-export { addComputed } from './utils/addComputed'
+export {
+  derive,
+  underive,
+  unstable_getDeriveSubscriptions,
+} from './utils/derive'
+export { addComputed_DEPRECATED as addComputed } from './utils/addComputed'
 export { proxyWithComputed } from './utils/proxyWithComputed'
 export { proxyWithHistory } from './utils/proxyWithHistory'
diff --git a/src/utils/addComputed.ts b/src/utils/addComputed.ts
index c99bd3a..c5e02b4 100644
--- a/src/utils/addComputed.ts
+++ b/src/utils/addComputed.ts
@@ -1,77 +1,25 @@
-import { createProxy as createProxyToCompare, isChanged } from 'proxy-compare'
-import { subscribe, snapshot } from '../vanilla'
-import type { DeepResolveType } from '../vanilla'
+import { derive } from './derive'
 
 /**
- * addComputed
+ * addComputed (DEPRECATED)
  *
- * This adds computed values to an existing proxy object.
- *
- * [Notes]
- * This comes with a cost and overlaps with useSnapshot.
- * Do not try to optimize too early. It can worsen the performance.
- * Measurement and comparison will be very important.
- *
- * @example
- * import { proxy } from 'valtio'
- * import { addComputed } from 'valtio/utils'
- * const state = proxy({
- *   count: 1,
- * })
- * addComputed(state, {
- *   doubled: snap => snap.count * 2,
- * })
+ * @deprecated Please consider using `derive` or `proxyWithComputed` instead.
  */
-export const addComputed = <T extends object, U extends object>(
+export const addComputed_DEPRECATED = <T extends object, U extends object>(
   proxyObject: T,
-  computedFns: {
-    [K in keyof U]: (snap: DeepResolveType<T>) => U[K]
+  computedFns_FAKE: {
+    [K in keyof U]: (snap_FAKE: T) => U[K]
   },
   targetObject: any = proxyObject
 ) => {
-  ;(Object.keys(computedFns) as (keyof U)[]).forEach((key) => {
-    if (Object.getOwnPropertyDescriptor(targetObject, key)) {
-      throw new Error('object property already defined')
-    }
-    const get = computedFns[key]
-    let prevSnapshot: DeepResolveType<T> | undefined
-    let affected = new WeakMap()
-    let pending = false
-    const callback = () => {
-      const nextSnapshot = snapshot(proxyObject)
-      if (
-        !pending &&
-        (!prevSnapshot || isChanged(prevSnapshot, nextSnapshot, affected))
-      ) {
-        affected = new WeakMap()
-        const value = get(createProxyToCompare(nextSnapshot, affected))
-        prevSnapshot = nextSnapshot
-        if (value instanceof Promise) {
-          pending = true
-          value
-            .then((v) => {
-              targetObject[key] = v
-            })
-            .catch((e) => {
-              // not ideal but best effort for throwing error with proxy
-              targetObject[key] = new Proxy(
-                {},
-                {
-                  get() {
-                    throw e
-                  },
-                }
-              )
-            })
-            .finally(() => {
-              pending = false
-            })
-        }
-        targetObject[key] = value
-      }
-    }
-    // FIXME no way to clean up the subscription
-    subscribe(proxyObject, callback)
-    callback()
+  console.warn(
+    'addComputed is deprecated. Please consider using `derive` or `proxyWithComputed` instead. Falling back to emulation with derive.'
+  )
+  const derivedFns: {
+    [K in keyof U]: (get: any) => U[K]
+  } = {} as any
+  ;(Object.keys(computedFns_FAKE) as (keyof U)[]).forEach((key) => {
+    derivedFns[key] = (get) => computedFns_FAKE[key](get(proxyObject))
   })
+  return derive(derivedFns, { proxy: targetObject })
 }
diff --git a/src/utils/derive.ts b/src/utils/derive.ts
new file mode 100644
index 0000000..774324a
--- /dev/null
+++ b/src/utils/derive.ts
@@ -0,0 +1,206 @@
+import { proxy, subscribe, getVersion } from '../vanilla'
+
+type DeriveGet = <T extends object>(proxyObject: T) => T
+
+type Subscriptions<U extends object> = Map<
+  object,
+  [callbackMap: Map<keyof U, () => void>, unsubscribe: () => void]
+>
+
+const subscriptionsCache = new WeakMap<object, Subscriptions<object>>()
+
+const getSubscriptions = (proxyObject: object) => {
+  let subscriptions = subscriptionsCache.get(proxyObject)
+  if (!subscriptions) {
+    subscriptions = new Map()
+    subscriptionsCache.set(proxyObject, subscriptions)
+  }
+  return subscriptions
+}
+
+// NOTE This is experimentally exported.
+// The availability is not guaranteed, and it will be renamed,
+// changed or removed without any notice in future versions.
+// It's not expected to use this in production.
+export const unstable_getDeriveSubscriptions = getSubscriptions
+
+/**
+ * derive
+ *
+ * This creates derived properties and attatches them
+ * to a new proxy object or an existing proxy object.
+ *
+ * @example
+ * import { proxy } from 'valtio'
+ * import { derive } from 'valtio/utils'
+ *
+ * const state = proxy({
+ *   count: 1,
+ * })
+ *
+ * const derivedState = derive({
+ *   doubled: (get) => get(state).count * 2,
+ * })
+ *
+ * derive({
+ *   tripled: (get) => get(state).count * 3,
+ * }, {
+ *   proxy: state,
+ * })
+ */
+export const derive = <T extends object, U extends object>(
+  derivedFns: {
+    [K in keyof U]: (get: DeriveGet) => U[K]
+  },
+  options?: {
+    proxy?: T
+    sync?: boolean
+  }
+) => {
+  const proxyObject = (options?.proxy || proxy({})) as U
+  const notifyInSync = options?.sync
+  const subscriptions: Subscriptions<U> = getSubscriptions(proxyObject)
+  const addSubscription = (p: object, key: keyof U, callback: () => void) => {
+    const subscription = subscriptions.get(p)
+    if (subscription) {
+      subscription[0].set(key, callback)
+    } else {
+      const unsubscribe = subscribe(
+        p,
+        (ops) => {
+          if (
+            p === proxyObject &&
+            ops.every((op) => op[1].length === 1 && op[1][0] in derivedFns)
+          ) {
+            // only setting derived properties
+            return
+          }
+          subscriptions.get(p)?.[0].forEach((cb) => {
+            cb()
+          })
+        },
+        notifyInSync
+      )
+      subscriptions.set(p, [new Map([[key, callback]]), unsubscribe])
+    }
+  }
+  const removeSubscription = (p: object, key: keyof U) => {
+    const subscription = subscriptions.get(p)
+    if (subscription) {
+      const [callbackMap, unsubscribe] = subscription
+      callbackMap.delete(key)
+      if (!callbackMap.size) {
+        unsubscribe()
+        subscriptions.delete(p)
+      }
+    }
+  }
+  ;(Object.keys(derivedFns) as (keyof U)[]).forEach((key) => {
+    if (Object.getOwnPropertyDescriptor(proxyObject, key)) {
+      throw new Error('object property already defined')
+    }
+    const fn = derivedFns[key]
+    let lastDependencies: Map<object, number> | null = null
+    const evaluate = () => {
+      if (lastDependencies) {
+        if (
+          Array.from(lastDependencies).every(([p, n]) => getVersion(p) === n)
+        ) {
+          // no dependencies are changed
+          return
+        }
+      }
+      const dependencies = new Map<object, number>()
+      const get = <P extends object>(p: P) => {
+        dependencies.set(p, getVersion(p))
+        return p
+      }
+      const value = fn(get)
+      const subscribe = () => {
+        dependencies.forEach((_, p) => {
+          if (!lastDependencies?.has(p)) {
+            addSubscription(p, key, evaluate)
+          }
+        })
+        lastDependencies?.forEach((_, p) => {
+          if (!dependencies.has(p)) {
+            removeSubscription(p, key)
+          }
+        })
+        lastDependencies = dependencies
+      }
+      if (value instanceof Promise) {
+        value.then(() => {
+          subscribe()
+          evaluate()
+        })
+      } else {
+        subscribe()
+      }
+      proxyObject[key] = value
+    }
+    evaluate()
+  })
+  return proxyObject as T & U
+}
+
+/**
+ * underive
+ *
+ * This stops derived properties to evaluate.
+ * It will stop all (or specified by `keys` option) subscriptions.
+ * If you specify `delete` option, it will delete the properties
+ * and you can attach new derived properties.
+ *
+ * @example
+ * import { proxy } from 'valtio'
+ * import { derive, underive } from 'valtio/utils'
+ *
+ * const state = proxy({
+ *   count: 1,
+ * })
+ *
+ * const derivedState = derive({
+ *   doubled: (get) => get(state).count * 2,
+ * })
+ *
+ * underive(derivedState)
+ */
+export const underive = <T extends object, U extends object>(
+  proxyObject: T & U,
+  options?: {
+    delete?: boolean
+    keys?: (keyof U)[]
+  }
+) => {
+  const subscriptions: Subscriptions<U> = getSubscriptions(proxyObject)
+  const keysToDelete = options?.delete ? new Set<keyof U>() : null
+  subscriptions.forEach(([callbackMap, unsubscribe], p) => {
+    if (options?.keys) {
+      options.keys.forEach((key) => {
+        if (callbackMap.has(key)) {
+          callbackMap.delete(key)
+          if (keysToDelete) {
+            keysToDelete.add(key)
+          }
+        }
+      })
+    } else {
+      if (keysToDelete) {
+        Array.from(callbackMap.keys()).forEach((key) => {
+          keysToDelete.add(key)
+        })
+      }
+      callbackMap.clear()
+    }
+    if (!callbackMap.size) {
+      unsubscribe()
+      subscriptions.delete(p)
+    }
+  })
+  if (keysToDelete) {
+    keysToDelete.forEach((key) => {
+      delete proxyObject[key]
+    })
+  }
+}


Test patch:
diff --git a/tests/computed.test.tsx b/tests/computed.test.tsx
index 61ef499..36a5861 100644
--- a/tests/computed.test.tsx
+++ b/tests/computed.test.tsx
@@ -19,6 +19,19 @@ afterEach(() => {
   console.error = consoleError
 })
 
+const consoleWarn = console.warn
+beforeEach(() => {
+  console.warn = jest.fn((message) => {
+    if (message.startsWith('addComputed is deprecated.')) {
+      return
+    }
+    consoleWarn(message)
+  })
+})
+afterEach(() => {
+  console.warn = consoleWarn
+})
+
 const sleep = (ms: number) =>
   new Promise((resolve) => {
     setTimeout(resolve, ms)
@@ -159,7 +172,7 @@ it('simple addComputed', async () => {
   state.text = 'a'
   await Promise.resolve()
   expect(snapshot(state)).toMatchObject({ text: 'a', count: 1, doubled: 2 })
-  expect(computeDouble).toBeCalledTimes(2)
+  // This can't pass with derive emulation: expect(computeDouble).toBeCalledTimes(2)
   expect(callback).toBeCalledTimes(2)
 })
 
@@ -239,7 +252,7 @@ it('nested emulation with addComputed', async () => {
     text: 'a',
     math: { count: 1, doubled: 2 },
   })
-  expect(computeDouble).toBeCalledTimes(2)
+  // This can't pass with derive emulation: expect(computeDouble).toBeCalledTimes(2)
   expect(callback).toBeCalledTimes(2)
 })
 
diff --git a/tests/derive.test.tsx b/tests/derive.test.tsx
new file mode 100644
index 0000000..4e7450c
--- /dev/null
+++ b/tests/derive.test.tsx
@@ -0,0 +1,300 @@
+import React, { StrictMode, Suspense } from 'react'
+import { fireEvent, render } from '@testing-library/react'
+import { proxy, useSnapshot, snapshot, subscribe } from '../src/index'
+import { derive, underive } from '../src/utils'
+
+const consoleError = console.error
+beforeEach(() => {
+  console.error = jest.fn((message) => {
+    if (
+      process.env.NODE_ENV === 'production' &&
+      message.startsWith('act(...) is not supported in production')
+    ) {
+      return
+    }
+    consoleError(message)
+  })
+})
+afterEach(() => {
+  console.error = consoleError
+})
+
+const sleep = (ms: number) =>
+  new Promise((resolve) => {
+    setTimeout(resolve, ms)
+  })
+
+it('basic derive', async () => {
+  const computeDouble = jest.fn((x) => x * 2)
+  const state = proxy({
+    text: '',
+    count: 0,
+  })
+  const derived = derive({
+    doubled: (get) => computeDouble(get(state).count),
+  })
+
+  const callback = jest.fn()
+  subscribe(derived, callback)
+
+  expect(snapshot(derived)).toMatchObject({ doubled: 0 })
+  expect(computeDouble).toBeCalledTimes(1)
+  expect(callback).toBeCalledTimes(0)
+
+  state.count += 1
+  await Promise.resolve()
+  expect(snapshot(derived)).toMatchObject({ doubled: 2 })
+  expect(computeDouble).toBeCalledTimes(2)
+  await Promise.resolve()
+  expect(callback).toBeCalledTimes(1)
+
+  state.text = 'a'
+  await Promise.resolve()
+  expect(snapshot(derived)).toMatchObject({ doubled: 2 })
+  expect(computeDouble).toBeCalledTimes(3)
+  await Promise.resolve()
+  expect(callback).toBeCalledTimes(1)
+})
+
+it('derive another proxy', async () => {
+  const computeDouble = jest.fn((x) => x * 2)
+  const state = proxy({
+    text: '',
+    count: 0,
+  })
+  const anotherState = proxy({})
+  derive(
+    {
+      doubled: (get) => computeDouble(get(state).count),
+    },
+    {
+      proxy: anotherState,
+    }
+  )
+
+  const callback = jest.fn()
+  subscribe(anotherState, callback)
+
+  expect(snapshot(anotherState)).toMatchObject({ doubled: 0 })
+  expect(computeDouble).toBeCalledTimes(1)
+  expect(callback).toBeCalledTimes(0)
+
+  state.count += 1
+  await Promise.resolve()
+  expect(snapshot(anotherState)).toMatchObject({ doubled: 2 })
+  expect(computeDouble).toBeCalledTimes(2)
+  await Promise.resolve()
+  expect(callback).toBeCalledTimes(1)
+
+  state.text = 'a'
+  await Promise.resolve()
+  expect(snapshot(anotherState)).toMatchObject({ doubled: 2 })
+  expect(computeDouble).toBeCalledTimes(3)
+  await Promise.resolve()
+  expect(callback).toBeCalledTimes(1)
+})
+
+it('derive with self', async () => {
+  const computeDouble = jest.fn((x) => x * 2)
+  const state = proxy({
+    text: '',
+    count: 0,
+  })
+  derive(
+    {
+      doubled: (get) => computeDouble(get(state).count),
+    },
+    {
+      proxy: state,
+    }
+  )
+
+  const callback = jest.fn()
+  subscribe(state, callback)
+
+  expect(snapshot(state)).toMatchObject({ text: '', count: 0, doubled: 0 })
+  expect(computeDouble).toBeCalledTimes(1)
+  expect(callback).toBeCalledTimes(0)
+
+  state.count += 1
+  await Promise.resolve()
+  expect(snapshot(state)).toMatchObject({ text: '', count: 1, doubled: 2 })
+  expect(computeDouble).toBeCalledTimes(2)
+  await Promise.resolve()
+  expect(callback).toBeCalledTimes(1)
+
+  state.text = 'a'
+  await Promise.resolve()
+  expect(snapshot(state)).toMatchObject({ text: 'a', count: 1, doubled: 2 })
+  expect(computeDouble).toBeCalledTimes(3)
+  await Promise.resolve()
+  expect(callback).toBeCalledTimes(2)
+})
+
+it('derive with two dependencies', async () => {
+  const computeSum = jest.fn((x, y) => x + y)
+  const state1 = proxy({ count: 1 })
+  const state2 = proxy({ count: 10 })
+  const derived = derive({
+    sum: (get) => computeSum(get(state1).count, get(state2).count),
+  })
+
+  const callback = jest.fn()
+  subscribe(derived, callback)
+
+  expect(snapshot(derived)).toMatchObject({ sum: 11 })
+  expect(computeSum).toBeCalledTimes(1)
+  expect(callback).toBeCalledTimes(0)
+
+  state1.count += 1
+  await Promise.resolve()
+  expect(snapshot(derived)).toMatchObject({ sum: 12 })
+  expect(computeSum).toBeCalledTimes(2)
+  await Promise.resolve()
+  expect(callback).toBeCalledTimes(1)
+
+  state1.count += 1
+  state2.count += 10
+  await Promise.resolve()
+  expect(snapshot(derived)).toMatchObject({ sum: 23 })
+  expect(computeSum).toBeCalledTimes(3)
+  await Promise.resolve()
+  expect(callback).toBeCalledTimes(2)
+})
+
+it('async derive', async () => {
+  const state = proxy({ count: 0 })
+  derive(
+    {
+      delayedCount: async (get) => {
+        await sleep(10)
+        return get(state).count + 1
+      },
+    },
+    { proxy: state }
+  )
+
+  const Counter: React.FC = () => {
+    const snap = useSnapshot(
+      state as { count: number; delayedCount: Promise<number> }
+    )
+    return (
+      <>
+        <div>
+          count: {snap.count}, delayedCount: {snap.delayedCount}
+        </div>
+        <button onClick={() => ++state.count}>button</button>
+      </>
+    )
+  }
+
+  const { getByText, findByText } = render(
+    <StrictMode>
+      <Suspense fallback="loading">
+        <Counter />
+      </Suspense>
+    </StrictMode>
+  )
+
+  await findByText('loading')
+  await findByText('count: 0, delayedCount: 1')
+
+  fireEvent.click(getByText('button'))
+  await findByText('loading')
+  await findByText('count: 1, delayedCount: 2')
+})
+
+it('nested emulation with derive', async () => {
+  const computeDouble = jest.fn((x) => x * 2)
+  const state = proxy({ text: '', math: { count: 0 } })
+  derive(
+    {
+      doubled: (get) => computeDouble(get(state.math).count),
+    },
+    { proxy: state.math }
+  )
+
+  const callback = jest.fn()
+  subscribe(state, callback)
+
+  expect(snapshot(state)).toMatchObject({
+    text: '',
+    math: { count: 0, doubled: 0 },
+  })
+  expect(computeDouble).toBeCalledTimes(1)
+  expect(callback).toBeCalledTimes(0)
+
+  state.math.count += 1
+  await Promise.resolve()
+  expect(snapshot(state)).toMatchObject({
+    text: '',
+    math: { count: 1, doubled: 2 },
+  })
+  expect(computeDouble).toBeCalledTimes(2)
+  await Promise.resolve()
+  expect(callback).toBeCalledTimes(2)
+
+  state.text = 'a'
+  await Promise.resolve()
+  expect(snapshot(state)).toMatchObject({
+    text: 'a',
+    math: { count: 1, doubled: 2 },
+  })
+  expect(computeDouble).toBeCalledTimes(2)
+  expect(callback).toBeCalledTimes(3)
+})
+
+it('derive with array.pop', async () => {
+  const state = proxy({
+    arr: [{ n: 1 }, { n: 2 }, { n: 3 }],
+  })
+  derive(
+    {
+      nums: (get) => get(state.arr).map((item) => item.n),
+    },
+    { proxy: state }
+  )
+
+  expect(snapshot(state)).toMatchObject({
+    arr: [{ n: 1 }, { n: 2 }, { n: 3 }],
+    nums: [1, 2, 3],
+  })
+
+  state.arr.pop()
+  await Promise.resolve()
+  expect(snapshot(state)).toMatchObject({
+    arr: [{ n: 1 }, { n: 2 }],
+    nums: [1, 2],
+  })
+})
+
+it('basic underive', async () => {
+  const computeDouble = jest.fn((x) => x * 2)
+  const state = proxy({ count: 0 })
+  const derived = derive({
+    doubled: (get) => computeDouble(get(state).count),
+  })
+
+  const callback = jest.fn()
+  subscribe(derived, callback)
+
+  expect(snapshot(derived)).toMatchObject({ doubled: 0 })
+  expect(computeDouble).toBeCalledTimes(1)
+  expect(callback).toBeCalledTimes(0)
+
+  state.count += 1
+  await Promise.resolve()
+  expect(snapshot(derived)).toMatchObject({ doubled: 2 })
+  expect(computeDouble).toBeCalledTimes(2)
+  await Promise.resolve()
+  expect(callback).toBeCalledTimes(1)
+
+  underive(derived)
+
+  state.count += 1
+  await Promise.resolve()
+  expect(snapshot(derived)).toMatchObject({ doubled: 2 })
+  expect(computeDouble).toBeCalledTimes(2)
+  await Promise.resolve()
+  expect(callback).toBeCalledTimes(1)
+})