[v2] Refactor: Move Form Groups to trie (#2326)

* refactor: use stop symbol instead of boolean for traversal

* refactor: move form groups into trie structure

* add changeset

* fix: apply suggestions from coderabbit

* fix: do second pass of fixes

Author: LeCarbonator

.changeset/tasty-humans-joke.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/form-core': patch
+---
+
+Refactor: Store Form Groups on trie nodes instead of the form instance

packages/form-core/src/FieldApi/FieldApi.lib.ts

@@ -50,6 +50,7 @@ import type {
 import type { ResolvedInternalFieldUpdateOptions } from '../types.lib'
 import type { FieldUpdateOptions, Updater } from '../types.public'
 import type { AnyInternalFormApi } from '../FormApi/FormApi.lib'
+import type { AnyInternalFormGroupApi } from '../FormGroupApi/FormGroupApi.lib'
 import type { ReadonlyAtom } from '@tanstack/store'
 import type { FieldApi, FieldApiOptions } from './FieldApi.public'
 import type {
@@ -310,6 +311,8 @@ export class InternalFieldApi<
   _listeners: Array<AnyFieldListener> | null
   _errorVisibility: ErrorVisibility<any, any> | undefined
   _errorBoundary: boolean
+  /** The form group occupying this trie node. */
+  _formGroup: AnyInternalFormGroupApi | null = null
 
   // TODO implement
   /**
@@ -651,15 +654,15 @@ export class InternalFieldApi<
 
     const seenValidatorFields = new WeakSet<AnyInternalFieldApi>()
 
-    visitFieldAndAncestors(this, (current) => {
-      if (current._isKilled) return false
+    visitFieldAndAncestors(this, (current, stop) => {
+      if (current._isKilled) return stop
 
       current._runFieldValidation(event)
       current._notifyValidator(event, seenValidatorFields)
       return undefined
     })
 
-    const group = this.form._getNearestFormGroupForField(this.name)
+    const group = this._getFormGroup()
     if (group) {
       group.validate(event, { triggerFieldApi: this })
       return
@@ -677,6 +680,7 @@ export class InternalFieldApi<
     options?: {
       onResult?: boolean
       onlyRunValidatorIndeces?: Array<number> | null
+      _startValidation?: () => () => void
     },
   ): Promise<FieldValidatorPipelineResult> {
     if (this._isKilled)
@@ -695,7 +699,14 @@ export class InternalFieldApi<
         thrownError: null,
       }
 
-    this._setValidationCount((count) => count + 1)
+    let finishValidation = options?._startValidation?.()
+    if (!finishValidation) {
+      this._setValidationCount((count) => count + 1)
+      finishValidation = () => {
+        this._setValidationCount((count) => Math.max(0, count - 1))
+      }
+    }
+
     try {
       return await runFieldValidatorPipeline({
         pipeline: validators,
@@ -712,7 +723,7 @@ export class InternalFieldApi<
         validatorIndecesToRun: options?.onlyRunValidatorIndeces ?? null,
       })
     } finally {
-      this._setValidationCount((count) => Math.max(0, count - 1))
+      finishValidation()
     }
   }
 
@@ -848,7 +859,7 @@ export class InternalFieldApi<
     batch(() => {
       const seenListenerFields = new WeakSet<AnyInternalFieldApi>()
 
-      visitFieldAndAncestors(this, (currNode) => {
+      visitFieldAndAncestors(this, (currNode, stop) => {
         const isOriginalField = currNode === originalField
         const { isSelfDirty, isSelfTouched, isBlurred } = currNode.meta
         const shouldUpdateDirty = isOriginalField && markAsDirty && !isSelfDirty
@@ -868,7 +879,7 @@ export class InternalFieldApi<
 
         currNode._notifyListener(event, seenListenerFields)
 
-        if (!doPropagate) return false
+        if (!doPropagate) return stop
         return undefined
       })
     })
@@ -1084,6 +1095,30 @@ export class InternalFieldApi<
     pruneFieldIfUnused(this)
   }
 
+  /** Updates the form group occupying this trie node. */
+  _setFormGroup(formGroup: AnyInternalFormGroupApi | null): void {
+    if (this._formGroup === formGroup) return
+
+    this._formGroup = formGroup
+    devtools().updateField?.(this)
+
+    if (!formGroup) this._pruneIfUnused()
+  }
+
+  /** Returns the form group containing this trie node, if one exists. */
+  _getFormGroup(): AnyInternalFormGroupApi | null {
+    let formGroup: AnyInternalFormGroupApi | null = null
+
+    visitFieldAndAncestors(this, (field, stop) => {
+      if (!field._formGroup) return
+
+      formGroup = field._formGroup
+      return stop
+    })
+
+    return formGroup
+  }
+
   _getValue(): any {
     return this.form.getFieldValue(this.name)
   }

packages/form-core/src/FieldApi/fieldState.lib.ts

@@ -29,7 +29,7 @@ export const childContributionKeys: Array<ChildContributionKey> = [
 interface MetaExtension {
   _formValidatorErrors: Array<Array<ValidationIssue>>
   _formValidatorErrorSourceEvents: Array<string | null>
-  _formGroupValidatorErrors: Map<object, FormGroupFieldErrorMeta>
+  _formGroupValidatorErrors: FormGroupFieldErrorMeta | null
   _fieldValidatorErrors: Array<Array<ValidationIssue>>
   _fieldValidatorErrorSourceEvents: Array<string | null>
   childContributionCounts: ChildContributionCounts
@@ -87,7 +87,7 @@ export const defaultInternalBaseFieldMeta: InternalBaseFieldMeta = {
   _fieldValidatorErrorSourceEvents: [],
   _formValidatorErrors: [],
   _formValidatorErrorSourceEvents: [],
-  _formGroupValidatorErrors: new Map(),
+  _formGroupValidatorErrors: null,
   _arrayVersion: 0,
 }
 
@@ -232,7 +232,7 @@ function shouldDisplayErrors(
   isDefaultValue = true,
 ): boolean {
   if (!field || !errorVisibility) return true
-  const group = field.form._getNearestFormGroupForField(field.name)
+  const group = field._getFormGroup()
   const stateOverrides = group?._getScopedFormStateOverrides()
 
   return errorVisibility({
@@ -295,8 +295,10 @@ export function getChildContributionStates(
   }
 }
 
-function hasValidatorErrors(errors: Array<Array<ValidationIssue>>): boolean {
-  return errors.some((validatorErrors) => validatorErrors.length > 0)
+function hasValidatorErrors(
+  errors: Array<Array<ValidationIssue>> | undefined,
+): boolean {
+  return errors?.some((validatorErrors) => validatorErrors.length > 0) ?? false
 }
 
 export function isPrunableMeta(meta: InternalBaseFieldMeta): boolean {
@@ -307,7 +309,7 @@ export function isPrunableMeta(meta: InternalBaseFieldMeta): boolean {
   if (meta._validationCount !== 0) return false
   if (meta._arrayVersion !== 0) return false
   if (hasValidatorErrors(meta._fieldValidatorErrors)) return false
-  if (hasFormGroupValidatorErrors(meta._formGroupValidatorErrors)) return false
+  if (hasValidatorErrors(meta._formGroupValidatorErrors?.errors)) return false
   if (hasValidatorErrors(meta._formValidatorErrors)) return false
 
   return childContributionKeys.every(
@@ -329,11 +331,7 @@ function getErrorsFromBaseMeta(
     result = previousMeta.original.errors
   } else {
     result = baseMeta._fieldValidatorErrors
-      .concat(
-        Array.from(baseMeta._formGroupValidatorErrors.values()).flatMap(
-          (groupErrors) => groupErrors.errors,
-        ),
-      )
+      .concat(baseMeta._formGroupValidatorErrors?.errors ?? [])
       .concat(baseMeta._formValidatorErrors)
       // ValidationError is OneOrMany, TypeScript doesn't realize that
       // flat also takes care of that
@@ -342,15 +340,6 @@ function getErrorsFromBaseMeta(
   return result
 }
 
-export function hasFormGroupValidatorErrors(
-  groupErrors: Map<object, FormGroupFieldErrorMeta>,
-): boolean {
-  for (const { errors } of groupErrors.values()) {
-    if (hasValidatorErrors(errors)) return true
-  }
-  return false
-}
-
 export function hasFieldMetaErrors(meta: InternalBaseFieldMeta): boolean {
   return (
     getErrorsFromBaseMeta(meta).length > 0 ||

packages/form-core/src/FieldApi/fieldTraversal.lib.ts

@@ -1,11 +1,18 @@
 import type { AnyInternalFieldApi } from './FieldApi.lib'
 import type { InternalRootFieldApi } from './RootFieldApi.lib'
 
-type FieldVisitor = (field: AnyInternalFieldApi) => void | false
+const stop = Symbol('stop field traversal')
+type FieldTraversalStop = typeof stop
+
+type FieldVisitor = (
+  field: AnyInternalFieldApi,
+  stop: FieldTraversalStop,
+) => void | FieldTraversalStop
 
 /**
  * Visits a field node followed by each of its ancestors, stopping before the
- * synthetic root node. Return `false` from the visitor to stop the traversal.
+ * synthetic root node. Return the visitor's `stop` argument to stop the
+ * traversal.
  *
  * The next parent is captured before the visitor runs, so removing or
  * reparenting the current node does not change the ancestor chain being walked.
@@ -18,7 +25,7 @@ export function visitFieldAndAncestors(
 
   while (!current._isRoot) {
     const parent: AnyInternalFieldApi | InternalRootFieldApi = current._parent
-    if (visitor(current) === false) return
+    if (visitor(current, stop) === stop) return
     current = parent
   }
 }
@@ -39,7 +46,7 @@ function visitFields(
 
   while (stack.length > 0) {
     const field = stack.pop()!
-    if (visitor(field) === false) return
+    if (visitor(field, stop) === stop) return
 
     const children = field._children
     for (let index = children.length - 1; index >= 0; index--) {
@@ -50,7 +57,7 @@ function visitFields(
 
 /**
  * Visits a field node and its descendants in insertion-order preorder. The
- * starting field is included. Return `false` from the visitor to stop the
+ * starting field is included. Return the visitor's `stop` argument to stop the
  * entire traversal, not only the current branch.
  *
  * A node's children are read after its visitor runs, so structural mutations
@@ -66,7 +73,7 @@ export function visitFieldSubtree(
 /**
  * Visits every field in a form trie in insertion-order preorder. The synthetic
  * root node is excluded; traversal starts at each of its field children. Return
- * `false` from the visitor to stop the entire traversal.
+ * the visitor's `stop` argument to stop the entire traversal.
  *
  * A node's children are read after its visitor runs, so structural mutations
  * made by the visitor affect which descendants are visited.

packages/form-core/src/FieldApi/fieldTree.lib.ts

@@ -262,6 +262,41 @@ function notifyFieldSubtreeListeners(
   }
 }
 
+function prepareFormGroupsForFieldReplacement(
+  fields: ReadonlyArray<AnyInternalFieldApi>,
+): () => void {
+  const fieldsToReplace = new Set(fields)
+  const formGroups = fields.flatMap((field) =>
+    field._formGroup ? [{ group: field._formGroup, name: field.name }] : [],
+  )
+  const affectedFormGroups = new Set(formGroups.map(({ group }) => group))
+
+  const replacementRoot = fields[0]
+  if (replacementRoot) {
+    visitFieldAndAncestors(replacementRoot, (field) => {
+      if (field._formGroup) affectedFormGroups.add(field._formGroup)
+    })
+  }
+
+  for (const group of affectedFormGroups) {
+    group._removeRoutedErrorFields(fieldsToReplace)
+  }
+
+  for (const { group } of formGroups) {
+    group._cancelValidation()
+  }
+
+  return () => {
+    if (formGroups.length === 0) return
+
+    batch(() => {
+      for (const { group, name } of formGroups) {
+        group._attachToFieldTrie(name)
+      }
+    })
+  }
+}
+
 export function killField(
   field: AnyInternalFieldApi,
   options: {
@@ -276,6 +311,7 @@ export function killField(
     field: AnyInternalFieldApi
     previousPath: string
   }> = []
+  let reattachFormGroups = () => {}
 
   batch(() => {
     const nodesToKill = collectFieldSubtree(field)
@@ -286,6 +322,8 @@ export function killField(
     const nodesToKillSet = new Set(nodesToKill)
     const fieldsToPruneAfterKill = new Set<AnyInternalFieldApi>()
 
+    reattachFormGroups = prepareFormGroupsForFieldReplacement(nodesToKill)
+
     if (options.listenerEvent) {
       notifyFieldSubtreeListeners(field, options.listenerEvent)
     }
@@ -330,6 +368,7 @@ export function killField(
 
       node._isKilled = true
       node._refCount = 0
+      node._formGroup = null
       node._defaultValueCache = null
       node._atoms.store = undefined
       if (node._pipelineCache) {
@@ -403,12 +442,14 @@ export function killField(
   if (dependencyChanges && dependencyChanges.length > 0) {
     bridge.fieldDependenciesChanged?.(dependencyChanges)
   }
+  reattachFormGroups()
 }
 
 export function canPruneField(field: AnyInternalFieldApi): boolean {
   if (field._isKilled) return false
 
   if (field._refCount > 0) return false
+  if (field._formGroup) return false
   if (field._childrenMap.size > 0) return false
   if (field._watchingFields) return false
   if (field._watchingValidatorFields) return false
@@ -428,8 +469,8 @@ export function pruneFieldIfUnused(field: AnyInternalFieldApi): void {
     ? new Array<{ field: AnyInternalFieldApi; previousPath: string }>()
     : null
 
-  visitFieldAndAncestors(field, (node) => {
-    if (!canPruneField(node)) return false
+  visitFieldAndAncestors(field, (node, stop) => {
+    if (!canPruneField(node)) return stop
 
     removedFields?.push({ field: node, previousPath: node.name })
     node._parent._removeChild(node._segment)