mobx-keystone 1.25.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>>
@@ -88,10 +86,26 @@ export function internalCustomRef<T extends object>(
       trackedRefs.add(this as Ref<object>)
       this.internalForceUpdateBackRefs(initialTarget)
 
+      // MobX 4/5 do not pass the previous value to reaction effects.
+      // The reaction is usually created inside an action, so its first tracking run is deferred
+      // to the end of the batch and `initialTarget` may be stale by then; capture the value of
+      // that first run instead, since that is the one the effect will be compared against.
+      let previousTarget = initialTarget
+      let firstTrackingRun = true
       // according to mwestrate this won't leak as long as we don't keep the disposer around
       reaction(
-        () => this.maybeCurrent,
-        (newTarget, oldTarget) => {
+        () => {
+          const target = this.maybeCurrent
+          if (firstTrackingRun) {
+            firstTrackingRun = false
+            previousTarget = target
+            this.internalForceUpdateBackRefs(target)
+          }
+          return target
+        },
+        (newTarget) => {
+          const oldTarget = previousTarget
+          previousTarget = newTarget
           this.internalForceUpdateBackRefs(newTarget)
 
           if (onResolvedValueChange && newTarget !== oldTarget) {
@@ -176,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.
  *
@@ -196,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
+  }
 
-  const idMap = computedIdTree.walk(root)
-  return idMap ? (idMap.get(id) as T | undefined) : undefined
+  if (getId === getModelRefId) {
+    const found = resolveModelRefId(root, id)
+    if (found !== modelRefIdNotIndexed) {
+      return found as T | undefined
+    }
+  }
+
+  // custom ids (or models with a custom `getRefId`)
+  return resolveRefIdByIndex(root, id, getId) as T | undefined
 }
 
 function getBackRefs<T extends object>(
@@ -269,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)
 }
@@ -316,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>({
@@ -326,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/rootStore/attachDetach.ts

@@ -4,9 +4,10 @@ import { HookAction } from "../action/hookActions"
 import { wrapInAction, wrapModelMethodInActionIfNeeded } from "../action/wrapInAction"
 import { type AnyModel, BaseModel } from "../model/BaseModel"
 import { WalkTreeMode, walkTree } from "../parent/walkTree"
+import { forEachWithDelayedThrow } from "../utils/forEachWithDelayedThrow"
 
 const onAttachedDisposers = new WeakMap<object, () => void>()
-const attachedToRootStore = new WeakSet<object>()
+const attachedToRootStore = new WeakMap<object, object>()
 
 /**
  * @internal
@@ -16,6 +17,7 @@ export const attachToRootStore = action(
   (rootStore: object, child: object): void => {
     // we use an array to ensure they will get called even if the actual hook modifies the tree
     const childrenToCall: AnyModel[] = []
+    const attachment = {}
 
     walkTree(
       child,
@@ -25,7 +27,7 @@ export const attachToRootStore = action(
         if (attachedToRootStore.has(ch)) {
           return
         }
-        attachedToRootStore.add(ch)
+        attachedToRootStore.set(ch, attachment)
 
         if (ch instanceof BaseModel && (ch as any).onAttachedToRootStore) {
           wrapModelMethodInActionIfNeeded(
@@ -40,15 +42,23 @@ export const attachToRootStore = action(
       WalkTreeMode.ParentFirst
     )
 
-    const childrenToCallLen = childrenToCall.length
-    for (let i = 0; i < childrenToCallLen; i++) {
-      const ch = childrenToCall[i]
+    forEachWithDelayedThrow(childrenToCall, (ch) => {
+      // An earlier hook may have detached or reattached this node already.
+      if (attachedToRootStore.get(ch) !== attachment) {
+        return
+      }
 
       const disposer = (ch as any).onAttachedToRootStore(rootStore)
       if (disposer) {
-        onAttachedDisposers.set(ch, disposer)
+        if (attachedToRootStore.get(ch) === attachment) {
+          onAttachedDisposers.set(ch, disposer)
+        } else {
+          // The hook detached or reattached this node before returning its cleanup.
+          // Dispose this attachment without overwriting a newer attachment's cleanup.
+          wrapDisposer(ch, disposer)()
+        }
       }
-    }
+    })
   }
 )
 
@@ -69,11 +79,7 @@ export const detachFromRootStore = action("detachFromRootStore", (child: object)
       const disposer = onAttachedDisposers.get(ch)
       if (disposer) {
         // wrap disposer in action
-        const disposerAction = wrapInAction({
-          nameOrNameFn: HookAction.OnAttachedToRootStoreDisposer,
-          fn: disposer,
-          actionType: ActionContextActionType.Sync,
-        }).bind(ch)
+        const disposerAction = wrapDisposer(ch, disposer)
         onAttachedDisposers.delete(ch)
 
         disposersToCall.push(disposerAction)
@@ -82,8 +88,13 @@ export const detachFromRootStore = action("detachFromRootStore", (child: object)
     WalkTreeMode.ChildrenFirst
   )
 
-  const disposersToCallLen = disposersToCall.length
-  for (let i = 0; i < disposersToCallLen; i++) {
-    disposersToCall[i]()
-  }
+  forEachWithDelayedThrow(disposersToCall, (dispose) => dispose())
 })
+
+function wrapDisposer(node: object, disposer: () => void): () => void {
+  return wrapInAction({
+    nameOrNameFn: HookAction.OnAttachedToRootStoreDisposer,
+    fn: disposer,
+    actionType: ActionContextActionType.Sync,
+  }).bind(node)
+}

package/src/rootStore/rootStore.ts

@@ -45,9 +45,11 @@ export const registerRootStore: <T extends object>(node: T) => T = action(
 
     entry.is = true
 
-    attachToRootStore(node, node)
-
-    entry.atom?.reportChanged()
+    try {
+      attachToRootStore(node, node)
+    } finally {
+      entry.atom?.reportChanged()
+    }
     return node
   }
 )
@@ -65,9 +67,11 @@ export const unregisterRootStore: (node: object) => void = action("unregisterRoo
   const entry = getOrCreateRootStoreEntry(node)
   entry.is = false
 
-  detachFromRootStore(node)
-
-  entry.atom?.reportChanged()
+  try {
+    detachFromRootStore(node)
+  } finally {
+    entry.atom?.reportChanged()
+  }
 })
 
 /**