New implementation of proxyMap and proxySet with performance improvem…

…ents (#965)

* pushing up to get more eyes on it

* added test for snapshot

* taking a break

* taking a break

* removing extra set methods for now

* sealed proxied obects

* got all tests to pass

* got rid of comma

* got rid of set tests in utils folder

* update version after update

* updated to allow values to be keys

* updated to allow values to be keys

* updated package.json scripts

* 2.0.1

* changing brances so storing attempts

* updates

* fixed get

* added return type to fix types on snapshot

* added this.data.length in has method

* updated has to only show this.data.map if key doesn't exist

* updated has to only show this.data.map if key doesn't exist

* changed has to check if it's in a snapshot

* remove getVersion - doesn't work

* added back in canProxy

* imported canProxy from valtio

* removed unneeded functions and checks that are done in valtio core

* changed canProxy

* pushing up for proxyMap testing

* removed unused file

* updated proxyMap to use sparse array

* removed use of internal map

* added proxySet

* moved proxySet to be the main one

* added old proxySet in for reference

* fixed proxySet

* made proxyMap get set and delete consistent

* better syntax for maybeProxify

* added benchmarks for proxyMap

* added benchmarks for proxySet

* removed deno from lock file

* removed all test files

* updated maybeproxify in proxySet

* pnpm frozen lockfile

* change autoinstall peers to true in pnpm-lock

* changed vitest.config.mts to ts

* changed version number back and updated import paths to be relative

* fixed toJSON to remain in order and proxified keys in .has() for proxyMap

* tested set vs filter on empty slots

* fixed pnpm lock file

* ran pnpm install

* auto-install-peers fix

* added isSnapshot function for better readability

* removed unused package

* removed linked package in package.json

* fixed lock file

* changed to isProxy with this context

* Update src/vanilla/utils/proxyMap.ts

Co-authored-by: Daishi Kato <dai-shi@users.noreply.github.com>

* pushing up for testing

* pushing up for testing

* fixed delete, has and get

* fixed delete, has and get

* uploading newest changes

* uploading newest changes

* reverting

* fixed import in test file

* updated to use nextIndex

* put nextIndex on the object itself

* added undefined values to initialization

* updated entries check

* changed nextIndex to use local var

* cached indexMap.has to index

* removed unneeded index.has call

* added this.nextIndex

* replaced this.data.length with this.nextIndex

* refactor

* updated initial data and initial next data

* added filled proxyMap and new benchmark

* benchmark minSize

* Revert "benchmark minSize"

This reverts commit 91206eb.

* compare 3

* compare 3

* keyVal

* new keyval

* minor optimization

* added MIN_DATA_SIZE constant

* move keyval

* revert some changes

* rename

* removed old benchmarks and added vitest benchmarks for proxyMap

* fixed lock file

* added ts extension for older versions of ts in test file

* changed MapIterator to IterableIterator

* no longer importing with .ts extension

* no longer importing with .ts extension

* proxySet changes and added keyvals solution

* removed ts extension from imports in proxySet and proxyMap-indexMap-keyvals

* fixed toJSON in proxySet

* change MapIterator to IterableIterator

* add tree1 impl

* got lock file from main

* fixing lock file

* added ts extensions backf

* removed test example

* extensions added back

* recreating lockfile

* reverting changes for website lockfile and todo-with-proxyMap lockfile

* removed ts extension from proxyMap.bench.ts

* reverted changes for lockfile and package.json

* removed old btree implementation

* changed MapIterator to IterableIterator

* incomplete rawMap impl

* a minor fix

* fix?

* export them properly

* rename a little bit

* revert default

* fixed size issue

* removed unneeded copy of key-vals

* snap map

* just to avoid warnings

* an improvement to rawMap1, not sure if it works

* enumerable size

* Revert "enumerable size"

This reverts commit 5acaed0.

* hack with _registerSnap

* added benchmarks

* do not enumerate index

* added snap registry to keyval

* fixed index check

* added snap awareness to proxySet

* added snap test to proxySet

* fix

* removed variations

* removed subscribes in proxyMap

* recover tree1 -> tree2

* fix a bug

* added old file for reference

* tree2 impl

* refactor

* bench tree2

* tests

* updated proxySet

* refactor proxyMap&Set

* removed test file

* merged branch and got to merge status

* removed todo comment

* added review suggestions

* added review suggestions for proxySet as well

* avoid maybeProxify

* no need to add key in this data

* removed isProxy check on has calls

* removed maybeProxify from proxySet

* added back in maybeProxify to proxySet

* refactor

---------

Co-authored-by: Daishi Kato <dai-shi@users.noreply.github.com>
Co-authored-by: daishi <daishi@axlight.com>

src/vanilla/utils/proxyMap.ts

@@ -1,111 +1,149 @@
-import { proxy } from '../../vanilla.ts'
+import { proxy, unstable_getInternalStates } from '../../vanilla.ts'
 
-type KeyValRecord<K, V> = [key: K, value: V]
+const { proxyStateMap, snapCache } = unstable_getInternalStates()
+const isProxy = (x: any) => proxyStateMap.has(x)
 
-type InternalProxyMap<K, V> = Map<K, V> & {
-  data: KeyValRecord<K, V>[]
-  toJSON: object
+type InternalProxyObject<K, V> = Map<K, V> & {
+  data: Array<V>
+  index: number
+  toJSON: () => Map<K, V>
 }
 
-/**
- * proxyMap
- *
- * This is to create a proxy which mimic the native Map behavior.
- * The API is the same as Map API
- *
- * @example
- * import { proxyMap } from 'valtio/utils'
- * const state = proxyMap([["key", "value"]])
- *
- * //can be used inside a proxy as well
- * const state = proxy({
- *   count: 1,
- *   map: proxyMap()
- * })
- *
- * // When using an object as a key, you can wrap it with `ref` so it's not proxied
- * // this is useful if you want to preserve the key equality
- * import { ref } from 'valtio'
- *
- * const key = ref({})
- * state.set(key, "value")
- * state.get(key) //value
- *
- * const key = {}
- * state.set(key, "value")
- * state.get(key) //undefined
- */
-export function proxyMap<K, V>(entries?: Iterable<readonly [K, V]> | null) {
-  const map: InternalProxyMap<K, V> = proxy({
-    data: Array.from(entries || []) as KeyValRecord<K, V>[],
-    has(key) {
-      return this.data.some((p) => p[0] === key)
+export function proxyMap<K, V>(entries?: Iterable<[K, V]> | undefined | null) {
+  const initialData: Array<V> = []
+  let initialIndex = 0
+  const indexMap = new Map<K, number>()
+
+  const snapMapCache = new WeakMap<object, Map<K, number>>()
+  const registerSnapMap = () => {
+    const cache = snapCache.get(vObject)
+    const latestSnap = cache?.[1]
+    if (latestSnap && !snapMapCache.has(latestSnap)) {
+      const clonedMap = new Map(indexMap)
+      snapMapCache.set(latestSnap, clonedMap)
+    }
+  }
+  const getMapForThis = (x: any) => snapMapCache.get(x) || indexMap
+
+  if (entries) {
+    if (typeof entries[Symbol.iterator] !== 'function') {
+      throw new TypeError(
+        'proxyMap:\n\tinitial state must be iterable\n\t\ttip: structure should be [[key, value]]',
+      )
+    }
+    for (const [key, value] of entries) {
+      indexMap.set(key, initialIndex)
+      initialData[initialIndex++] = value
+    }
+  }
+
+  const vObject: InternalProxyObject<K, V> = {
+    data: initialData,
+    index: initialIndex,
+    get size() {
+      if (!isProxy(this)) {
+        registerSnapMap()
+      }
+      const map = getMapForThis(this)
+      return map.size
+    },
+    get(key: K) {
+      const map = getMapForThis(this)
+      const index = map.get(key)
+      if (index === undefined) {
+        // eslint-disable-next-line @typescript-eslint/no-unused-expressions
+        this.index // touch property for tracking
+        return undefined
+      }
+      return this.data[index]
+    },
+    has(key: K) {
+      const map = getMapForThis(this)
+      const exists = map.has(key)
+      if (!exists) {
+        // eslint-disable-next-line @typescript-eslint/no-unused-expressions
+        this.index // touch property for tracking
+      }
+      return exists
     },
-    set(key, value) {
-      const record = this.data.find((p) => p[0] === key)
-      if (record) {
-        record[1] = value
+    set(key: K, value: V) {
+      if (!isProxy(this)) {
+        throw new Error('Cannot perform mutations on a snapshot')
+      }
+      const index = indexMap.get(key)
+      if (index === undefined) {
+        indexMap.set(key, this.index)
+        this.data[this.index++] = value
       } else {
-        this.data.push([key, value])
+        this.data[index] = value
       }
       return this
     },
-    get(key) {
-      return this.data.find((p) => p[0] === key)?.[1]
-    },
-    delete(key) {
-      const index = this.data.findIndex((p) => p[0] === key)
-      if (index === -1) {
+    delete(key: K) {
+      if (!isProxy(this)) {
+        throw new Error('Cannot perform mutations on a snapshot')
+      }
+      const index = indexMap.get(key)
+      if (index === undefined) {
         return false
       }
-      this.data.splice(index, 1)
+      delete this.data[index]
+      indexMap.delete(key)
       return true
     },
     clear() {
-      this.data.splice(0)
-    },
-    get size() {
-      return this.data.length
-    },
-    toJSON() {
-      return new Map(this.data)
-    },
-    forEach(cb) {
-      this.data.forEach((p) => {
-        cb(p[1], p[0], this)
+      if (!isProxy(this)) {
+        throw new Error('Cannot perform mutations on a snapshot')
+      }
+      this.data.length = 0 // empty array
+      this.index = 0
+      indexMap.clear()
+    },
+    forEach(cb: (value: V, key: K, map: Map<K, V>) => void) {
+      const map = getMapForThis(this)
+      map.forEach((index, key) => {
+        cb(this.data[index]!, key, this)
       })
     },
-    keys() {
-      return this.data.map((p) => p[0]).values()
-    },
-    values() {
-      return this.data.map((p) => p[1]).values()
+    *entries(): MapIterator<[K, V]> {
+      const map = getMapForThis(this)
+      for (const [key, index] of map) {
+        yield [key, this.data[index]!]
+      }
     },
-    entries() {
-      return new Map(this.data).entries()
+    *keys(): IterableIterator<K> {
+      const map = getMapForThis(this)
+      for (const key of map.keys()) {
+        yield key
+      }
     },
-    get [Symbol.toStringTag]() {
-      return 'Map'
+    *values(): IterableIterator<V> {
+      const map = getMapForThis(this)
+      for (const index of map.values()) {
+        yield this.data[index]!
+      }
     },
     [Symbol.iterator]() {
       return this.entries()
     },
-  })
-
-  Object.defineProperties(map, {
-    data: {
-      enumerable: false,
-    },
-    size: {
-      enumerable: false,
+    get [Symbol.toStringTag]() {
+      return 'Map'
     },
-    toJSON: {
-      enumerable: false,
+    toJSON(): Map<K, V> {
+      return new Map(this.entries())
     },
+  }
+
+  const proxiedObject = proxy(vObject)
+  Object.defineProperties(proxiedObject, {
+    size: { enumerable: false },
+    index: { enumerable: false },
+    data: { enumerable: false },
+    toJSON: { enumerable: false },
   })
-  Object.seal(map)
+  Object.seal(proxiedObject)
 
-  return map as unknown as Map<K, V> & {
+  return proxiedObject as unknown as Map<K, V> & {
     $$valtioSnapshot: Omit<Map<K, V>, 'set' | 'delete' | 'clear'>
   }
 }