mobx-keystone 2.0.0 → 2.1.0

package/src/ref/core.ts

@@ -2,16 +2,14 @@ import { action, type ObservableSet, observable, reaction, when } from "mobx"
 import { isModel } from "../model/utils"
 import type { ModelClass } from "../modelShared/BaseModelShared"
 import { model } from "../modelShared/modelDecorator"
-import {
-  getDeepObjectChildren,
-  registerDeepObjectChildrenExtension,
-} from "../parent/coreObjectChildren"
+import { registerDeepObjectChildrenExtension } from "../parent/coreObjectChildren"
 import { fastGetRoot } from "../parent/path"
-import { type ComputedWalkTreeAggregate, computedWalkTreeAggregate } from "../parent/walkTree"
 import { assertTweakedObject } from "../tweaker/core"
 import { assertIsObject, failure, mobxAction } from "../utils"
 import { getOrCreate } from "../utils/mapUtils"
+import { modelRefIdNotIndexed, resolveModelRefId } from "./modelRefIdIndex"
 import { Ref, type RefConstructor } from "./Ref"
+import { resolveRefIdByIndex } from "./refIdIndex"
 
 interface BackRefs<T extends object> {
   all: ObservableSet<Ref<T>>
@@ -192,12 +190,6 @@ export function getModelRefId(target: object): string | undefined {
   return undefined
 }
 
-// one computed id tree per id function
-const computedIdTrees = new WeakMap<
-  (node: object) => string | undefined,
-  ComputedWalkTreeAggregate<string>
->()
-
 /**
  * Resolves a node given its ID.
  *
@@ -212,13 +204,20 @@ export function resolveId<T extends object>(
   id: string,
   getId: RefIdResolver = getModelRefId
 ): T | undefined {
-  // cache/reuse computedIdTrees for same getId function
-  const computedIdTree = getOrCreate(computedIdTrees, getId, () =>
-    computedWalkTreeAggregate<string>((node) => getId(node))
-  )
+  // parent resolutions have priority over child ones
+  if (getId(root) === id) {
+    return root as T
+  }
+
+  if (getId === getModelRefId) {
+    const found = resolveModelRefId(root, id)
+    if (found !== modelRefIdNotIndexed) {
+      return found as T | undefined
+    }
+  }
 
-  const idMap = computedIdTree.walk(root)
-  return idMap ? (idMap.get(id) as T | undefined) : undefined
+  // custom ids (or models with a custom `getRefId`)
+  return resolveRefIdByIndex(root, id, getId) as T | undefined
 }
 
 function getBackRefs<T extends object>(
@@ -285,14 +284,25 @@ export function getRefsResolvingTo<T extends object>(
     oldBackRefs.forEach(updateRef)
   }
 
-  const refsChildrenOfRoot = getDeepChildrenRefs(getDeepObjectChildren(fastGetRoot(target, true)))
-  let refs: Set<Ref<object>> | undefined
-  if (refType) {
-    refs = refsChildrenOfRoot.byType.get(refType.refClass)
+  const refsChildrenOfRoot = getDeepChildrenRefs.get(fastGetRoot(target, true))
+  if (shouldForceUpdateTrackedRefs) {
+    const refs = refType ? refsChildrenOfRoot.byType.get(refType.refClass) : refsChildrenOfRoot.all
+    refs?.forEach(updateRef)
   } else {
-    refs = refsChildrenOfRoot.all
+    // tracked refs keep their back-references up to date by themselves
+    const untrackedRefs = refType
+      ? refsChildrenOfRoot.untrackedByType.get(refType.refClass)
+      : refsChildrenOfRoot.untracked
+    if (untrackedRefs && untrackedRefs.size > 0) {
+      // copy, since updating back-references might change the tree
+      for (const ref of Array.from(untrackedRefs)) {
+        updateRef(ref)
+        if (trackedRefs.has(ref)) {
+          deleteRef(refsChildrenOfRoot.untracked, refsChildrenOfRoot.untrackedByType, ref)
+        }
+      }
+    }
   }
-  refs?.forEach(updateRef)
 
   return getBackRefs(target, refType)
 }
@@ -332,9 +342,24 @@ function isReactionDelayed() {
   return reactionDelayed
 }
 
+type RefsByType = WeakMap<ModelClass<Ref<any>>, Set<Ref<any>>>
+
 interface DeepChildrenRefs {
   all: Set<Ref<any>>
-  byType: WeakMap<ModelClass<Ref<any>>, Set<Ref<any>>>
+  byType: RefsByType
+  /** Refs whose back-references might not be tracked yet. */
+  untracked: Set<Ref<any>>
+  untrackedByType: RefsByType
+}
+
+function addRef(all: Set<Ref<any>>, byType: RefsByType, ref: Ref<any>) {
+  all.add(ref)
+  getOrCreate(byType, ref.constructor as ModelClass<Ref<any>>, () => new Set()).add(ref)
+}
+
+function deleteRef(all: Set<Ref<any>>, byType: RefsByType, ref: Ref<any>) {
+  all.delete(ref)
+  byType.get(ref.constructor as ModelClass<Ref<any>>)?.delete(ref)
 }
 
 const getDeepChildrenRefs = registerDeepObjectChildrenExtension<DeepChildrenRefs>({
@@ -342,14 +367,25 @@ const getDeepChildrenRefs = registerDeepObjectChildrenExtension<DeepChildrenRefs
     return {
       all: new Set(),
       byType: new WeakMap(),
+      untracked: new Set(),
+      untrackedByType: new WeakMap(),
     }
   },
 
   addNode(node, data) {
     if (node instanceof Ref) {
-      data.all.add(node)
-      const refsByThisType = getOrCreate(data.byType, node.constructor, () => new Set())
-      refsByThisType.add(node)
+      addRef(data.all, data.byType, node)
+      if (!trackedRefs.has(node)) {
+        addRef(data.untracked, data.untrackedByType, node)
+      }
+    }
+  },
+
+  removeNode(node, data) {
+    if (node instanceof Ref) {
+      deleteRef(data.all, data.byType, node)
+      deleteRef(data.untracked, data.untrackedByType, node)
     }
+    return true
   },
 })

package/src/ref/modelRefIdIndex.ts

@@ -0,0 +1,131 @@
+import { createAtom, type IAtom } from "mobx"
+import { type AnyModel, BaseModel } from "../model/BaseModel"
+import { modelIdKey } from "../model/metadata"
+import { registerDeepObjectChildrenExtension } from "../parent/coreObjectChildren"
+import { NodesById } from "./nodesById"
+
+interface ModelRefIdIndex {
+  readonly byId: NodesById<AnyModel>
+  /** Id each model is indexed under, since ids can change later. */
+  readonly indexedIds: Map<AnyModel, string>
+  /** Models with a custom `getRefId`, which cannot be indexed. */
+  customRefIdModels: number
+  /** Changed when models with a custom `getRefId` appear or disappear. */
+  readonly customRefIdAtom: IAtom
+  /** Changed when the models with a given id change. Only kept while observed. */
+  readonly idAtoms: Map<string, IAtom>
+}
+
+function hasCustomRefId(model: AnyModel) {
+  return model.getRefId !== BaseModel.prototype.getRefId
+}
+
+function reportIdChanged(index: ModelRefIdIndex, id: string) {
+  index.idAtoms.get(id)?.reportChanged()
+}
+
+const modelRefIdIndex = registerDeepObjectChildrenExtension<ModelRefIdIndex>({
+  initData() {
+    return {
+      byId: new NodesById(),
+      indexedIds: new Map(),
+      customRefIdModels: 0,
+      customRefIdAtom: createAtom("customRefIdModels"),
+      idAtoms: new Map(),
+    }
+  },
+
+  addNode(node, index) {
+    if (!(node instanceof BaseModel)) {
+      return
+    }
+    if (hasCustomRefId(node)) {
+      index.customRefIdModels++
+      if (index.customRefIdModels === 1) {
+        index.customRefIdAtom.reportChanged()
+      }
+      return
+    }
+
+    const id = node[modelIdKey]
+    if (id === undefined) {
+      return
+    }
+    index.byId.add(id, node)
+    index.indexedIds.set(node, id)
+    reportIdChanged(index, id)
+  },
+
+  removeNode(node, index) {
+    if (!(node instanceof BaseModel)) {
+      return true
+    }
+    if (hasCustomRefId(node)) {
+      index.customRefIdModels--
+      if (index.customRefIdModels === 0) {
+        index.customRefIdAtom.reportChanged()
+      }
+      return true
+    }
+
+    const id = index.indexedIds.get(node)
+    if (id === undefined) {
+      return true
+    }
+    index.indexedIds.delete(node)
+    index.byId.delete(id, node)
+    reportIdChanged(index, id)
+    return true
+  },
+
+  dependsOnModelIds: true,
+})
+
+function observeId(index: ModelRefIdIndex, id: string) {
+  const existing = index.idAtoms.get(id)
+  if (existing) {
+    existing.reportObserved()
+    return
+  }
+  const atom = createAtom("modelRefId", undefined, () => {
+    if (index.idAtoms.get(id) === atom) {
+      index.idAtoms.delete(id)
+    }
+  })
+  // only keep atoms something observes, since nothing else would remove them
+  if (atom.reportObserved()) {
+    index.idAtoms.set(id, atom)
+  }
+}
+
+/**
+ * Result of `resolveModelRefId` when the tree has models with a custom
+ * `getRefId`, which it cannot index.
+ *
+ * @internal
+ */
+export const modelRefIdNotIndexed = Symbol("modelRefIdNotIndexed")
+
+/**
+ * Resolves a model below the root (excluding it) by its default reference id
+ * (`getModelRefId`), reactively, in O(1), or O(models × depth) when several
+ * models have that id.
+ * Returns `modelRefIdNotIndexed` when the tree has models with a custom
+ * `getRefId`.
+ *
+ * @internal
+ */
+export function resolveModelRefId(
+  root: object,
+  id: string
+): object | undefined | typeof modelRefIdNotIndexed {
+  const index = modelRefIdIndex.get(root, false)
+
+  index.customRefIdAtom.reportObserved()
+  if (index.customRefIdModels > 0) {
+    return modelRefIdNotIndexed
+  }
+
+  observeId(index, id)
+  return index.byId.resolve(root, id)
+}

package/src/ref/nodesById.ts

@@ -0,0 +1,60 @@
+import { compareTreePaths, getPathFromAncestor } from "../parent/treeOrder"
+
+/**
+ * The nodes below a root indexed by id, where several nodes may share an id.
+ *
+ * @internal
+ */
+export class NodesById<T extends object> {
+  // (tree nodes are never sets)
+  private readonly byId = new Map<string, T | Set<T>>()
+
+  add(id: string, node: T): void {
+    const existing = this.byId.get(id)
+    if (existing === undefined) {
+      this.byId.set(id, node)
+    } else if (existing instanceof Set) {
+      existing.add(node)
+    } else {
+      this.byId.set(id, new Set([existing, node]))
+    }
+  }
+
+  delete(id: string, node: T): void {
+    const existing = this.byId.get(id)
+    if (existing instanceof Set) {
+      existing.delete(node)
+      if (existing.size === 1) {
+        this.byId.set(id, existing.values().next().value!)
+      }
+    } else if (existing === node) {
+      this.byId.delete(id)
+    }
+  }
+
+  /** Resolves the node with the id, picking one when several share it. */
+  resolve(root: object, id: string): T | undefined {
+    const found = this.byId.get(id)
+    return found instanceof Set ? resolveDuplicatedId(root, found) : found
+  }
+}
+
+/**
+ * Picks which of several nodes below the root sharing an id it resolves to:
+ * an ancestor wins over its descendants, and otherwise the node below the
+ * later attached child of their closest common ancestor wins. That is the last
+ * one a post-order walk of the tree would find. Reactive, in
+ * O(nodes × depth).
+ */
+function resolveDuplicatedId<T extends object>(root: object, nodes: Iterable<T>): T {
+  let winner: T | undefined
+  let winnerPath: object[] | undefined
+  for (const node of nodes) {
+    const path = getPathFromAncestor(root, node, true)!
+    if (!winnerPath || compareTreePaths(path, winnerPath, false) > 0) {
+      winner = node
+      winnerPath = path
+    }
+  }
+  return winner!
+}

package/src/ref/refIdIndex.ts

@@ -0,0 +1,209 @@
+import { computed, createAtom, type IComputedValue, Reaction } from "mobx"
+import {
+  type DeepObjectChildrenExtensionAccessor,
+  registerDeepObjectChildrenExtension,
+} from "../parent/coreObjectChildren"
+import { ChunkedSet, type SequenceChunk } from "../utils/chunks"
+import { getOrCreate } from "../utils/mapUtils"
+import { NodesById } from "./nodesById"
+
+type GetId = (target: object) => string | undefined
+
+/**
+ * A group of nodes, with a computed map of their current ids, so an id change
+ * only re-reads the ids of the nodes in its chunk.
+ */
+class IdChunk implements SequenceChunk<object> {
+  items: object[]
+  /** Changed when nodes are added or removed. */
+  readonly atom = createAtom("refIdChunk")
+  readonly ids: IComputedValue<Map<string, object | object[]>>
+  /** Ids last copied to the index. */
+  syncedIds: Map<string, object | object[]> | undefined
+
+  constructor(nodes: object[], getId: GetId) {
+    this.items = nodes
+    this.ids = computed(() => {
+      this.atom.reportObserved()
+      const ids = new Map<string, object | object[]>()
+      for (const node of this.items) {
+        const id = getId(node)
+        if (id === undefined) {
+          continue
+        }
+        const existing = ids.get(id)
+        if (existing === undefined) {
+          ids.set(id, node)
+        } else if (Array.isArray(existing)) {
+          existing.push(node)
+        } else {
+          ids.set(id, [existing, node])
+        }
+      }
+      return ids
+    })
+  }
+}
+
+class RefIdIndex extends ChunkedSet<object, IdChunk> {
+  /** Changed when chunks are added or removed. */
+  readonly chunksAtom = createAtom("refIdChunks")
+  readonly byId = new NodesById<object>()
+  /** Copies the chunks' id changes into `byId`. Changes when `byId` does. */
+  readonly synced: IComputedValue<number>
+  /** Set when `byId` changed outside of a sync, so the version must change. */
+  private changedOutsideSync = false
+  keeper: Reaction | undefined
+  private readonly getId: GetId
+
+  constructor(getId: GetId) {
+    super()
+    this.getId = getId
+    let version = 0
+    this.synced = computed(
+      () => {
+        this.chunksAtom.reportObserved()
+        let changed = this.changedOutsideSync
+        this.changedOutsideSync = false
+        for (const chunk of this.chunks) {
+          if (this.syncChunk(chunk)) {
+            changed = true
+          }
+        }
+        if (changed) {
+          version++
+        }
+        return version
+      },
+      { name: "refIdIndex" }
+    )
+  }
+
+  protected createChunk(nodes: object[]): IdChunk {
+    return new IdChunk(nodes, this.getId)
+  }
+
+  protected onChunkChanged(chunk: IdChunk): void {
+    chunk.atom.reportChanged()
+  }
+
+  protected onChunksChanged(): void {
+    this.chunksAtom.reportChanged()
+  }
+
+  protected onChunkDropped(chunk: IdChunk): void {
+    // its nodes (if any) moved to a chunk that will add their ids again
+    const synced = chunk.syncedIds
+    if (synced && synced.size > 0) {
+      for (const [id, value] of synced) {
+        forEachNode(value, (node) => this.byId.delete(id, node))
+      }
+      this.changedOutsideSync = true
+    }
+  }
+
+  private syncChunk(chunk: IdChunk): boolean {
+    const ids = chunk.ids.get()
+    const previous = chunk.syncedIds
+    if (ids === previous) {
+      return false
+    }
+    chunk.syncedIds = ids
+
+    let changed = false
+    if (previous) {
+      for (const [id, value] of previous) {
+        if (!sameNodes(ids.get(id), value)) {
+          forEachNode(value, (node) => this.byId.delete(id, node))
+          changed = true
+        }
+      }
+    }
+    for (const [id, value] of ids) {
+      if (!sameNodes(previous?.get(id), value)) {
+        forEachNode(value, (node) => this.byId.add(id, node))
+        changed = true
+      }
+    }
+    return changed
+  }
+}
+
+function forEachNode(value: object | object[], fn: (node: object) => void) {
+  if (Array.isArray(value)) {
+    value.forEach(fn)
+  } else {
+    fn(value)
+  }
+}
+
+function sameNodes(a: object | object[] | undefined, b: object | object[]) {
+  if (a === b) {
+    return true
+  }
+  if (!Array.isArray(a) || !Array.isArray(b) || a.length !== b.length) {
+    return false
+  }
+  return a.every((node, i) => node === b[i])
+}
+
+function createRefIdIndexExtension(getId: GetId) {
+  return registerDeepObjectChildrenExtension<RefIdIndex>({
+    initData() {
+      return new RefIdIndex(getId)
+    },
+
+    addNode(node, index) {
+      index.add(node)
+    },
+
+    removeNode(node, index) {
+      index.delete(node)
+      return true
+    },
+  })
+}
+
+const refIdIndexes = new WeakMap<GetId, DeepObjectChildrenExtensionAccessor<RefIdIndex>>()
+
+// Outside reactions and actions MobX does not cache computeds, so every
+// untracked resolution would re-read every id. Keep the index observed until
+// any of its ids change, then release it, so ids read from outside the tree
+// cannot keep it alive.
+function keepAliveUntilChanged(index: RefIdIndex) {
+  if (index.keeper) {
+    return
+  }
+  const keeper = new Reaction("refIdIndex keeper", () => {
+    keeper.dispose()
+    if (index.keeper === keeper) {
+      index.keeper = undefined
+    }
+  })
+  index.keeper = keeper
+  keeper.track(() => {
+    try {
+      index.synced.get()
+    } catch {
+      // the caller's own read rethrows the cached error
+    }
+  })
+}
+
+/**
+ * Resolves a node below the root (excluding it) whose `getId` is the given
+ * id, reactively, in O(tree size / 128) plus the cost of re-reading the ids of
+ * any chunk of 128 nodes that changed, plus O(nodes × depth) when several
+ * nodes have that id.
+ *
+ * @internal
+ */
+export function resolveRefIdByIndex(root: object, id: string, getId: GetId): object | undefined {
+  const accessor = getOrCreate(refIdIndexes, getId, () => createRefIdIndexExtension(getId))
+  const index = accessor.get(root, false)
+
+  keepAliveUntilChanged(index)
+  index.synced.get()
+
+  return index.byId.resolve(root, id)
+}

package/src/snapshot/applySnapshot.ts

@@ -13,6 +13,7 @@ import {
   getModelNotRegisteredErrorMessage,
   modelInfoByClass,
 } from "../modelShared/modelInfo"
+import { withChangeBatch } from "../tweaker/changeRollback"
 import { assertTweakedObject } from "../tweaker/core"
 import { assertIsObject, inDevMode, isArray, isMap, isPlainObject, isSet, lazy } from "../utils"
 import { runWithErrorDiagnosticsContext } from "../utils/errorDiagnostics"
@@ -24,6 +25,7 @@ import { SnapshotProcessingError } from "./SnapshotProcessingError"
 
 /**
  * Applies a full snapshot over an object, reconciling it with the current contents of the object.
+ * If it throws, the changes it made are rolled back (errors thrown by listeners do not roll them back).
  *
  * @template T Object type.
  * @param node Target object (model object, object or array).
@@ -52,27 +54,33 @@ export function internalApplySnapshot<T extends object>(
 ): void {
   const obj = this
 
-  const reconcile = () => {
-    // Avoid building a model pool (and its deep children traversal) when the
-    // node already holds this exact snapshot. Child updates may still await
-    // propagation within an action, so flush before trusting the reference.
-    flushInternalSnapshot(obj, false)
-    if (getInternalSnapshot(obj)?.transformed === sn) {
-      return
-    }
+  const reconcile = () =>
+    withChangeBatch(() => {
+      // Avoid building a model pool (and its deep children traversal) when the
+      // node already holds this exact snapshot. Child updates may still await
+      // propagation within an action, so flush before trusting the reference.
+      flushInternalSnapshot(obj, false)
+      if (getInternalSnapshot(obj)?.transformed === sn) {
+        return
+      }
 
-    const modelPool = new ModelPool(obj)
-    const ret = reconcileSnapshot(obj, sn, modelPool, undefined)
+      const modelPool = new ModelPool(obj)
+      let ret: unknown
+      try {
+        ret = reconcileSnapshot(obj, sn, modelPool, undefined)
+      } finally {
+        modelPool.release()
+      }
 
-    if (inDevMode) {
-      if (ret !== obj) {
-        throw new SnapshotProcessingError({
-          message: "assertion failed: reconciled object has to be the same",
-          actualSnapshot: sn,
-        })
+      if (inDevMode) {
+        if (ret !== obj) {
+          throw new SnapshotProcessingError({
+            message: "assertion failed: reconciled object has to be the same",
+            actualSnapshot: sn,
+          })
+        }
       }
-    }
-  }
+    })
 
   if (isArray(sn)) {
     if (!isArray(obj)) {

package/src/snapshot/fromPlainObjectSnapshot.ts

@@ -1,13 +1,8 @@
 import { observable } from "mobx"
 import { tweakPlainObject } from "../tweaker/tweakPlainObject"
-import { isPlainObject, setProtoProp } from "../utils"
+import { isPlainObject, nodeObservableOptions, setProtoProp } from "../utils"
 import { getCurrentErrorDiagnosticsContext } from "../utils/errorDiagnostics"
-import {
-  type FromSnapshotContext,
-  internalFromSnapshot,
-  observableOptions,
-  registerSnapshotter,
-} from "./fromSnapshot"
+import { type FromSnapshotContext, internalFromSnapshot, registerSnapshotter } from "./fromSnapshot"
 import type { SnapshotInOfObject } from "./SnapshotOf"
 import { SnapshotterAndReconcilerPriority } from "./SnapshotterAndReconcilerPriority"
 
@@ -37,7 +32,7 @@ function fromPlainObjectSnapshot(sn: SnapshotInOfObject<any>, ctx: FromSnapshotC
     }
   }
   return tweakPlainObject(
-    observable.object(plainObj, undefined, observableOptions),
+    observable.object(plainObj, undefined, nodeObservableOptions),
     undefined,
     undefined,
     true,

package/src/snapshot/fromSnapshot.ts

@@ -5,7 +5,7 @@ import { resolveStandardTypeNoThrow, resolveTypeChecker } from "../types/resolve
 import type { AnyType, TypeToData, TypeToSnapshotIn } from "../types/schemas"
 import { isLateTypeChecker, TypeChecker } from "../types/TypeChecker"
 import { resolveCodecSupport } from "../types/utility/typesCodec"
-import { isMap, isPrimitive, isSet, setProtoProp } from "../utils"
+import { isMap, isPrimitive, isSet, nodeObservableOptions, setProtoProp } from "../utils"
 import {
   getCurrentErrorDiagnosticsContext,
   runWithErrorDiagnosticsContext,
@@ -210,13 +210,9 @@ function snapshotToInitialData(
       }
     }
   }
-  const observableInitialData = observable.object(initialData, undefined, observableOptions)
+  const observableInitialData = observable.object(initialData, undefined, nodeObservableOptions)
   if (allValuesPrimitive) {
     setModelInitialDataSnapshot(observableInitialData, initialData)
   }
   return observableInitialData
 }
-
-export const observableOptions = {
-  deep: false,
-}