@tldraw/store 3.16.0-internal.a478398270c6 → 3.16.0-internal.f8b97f0c414f

package/src/lib/StoreSchema.ts

@@ -107,6 +107,7 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 
 	readonly migrations: Record<string, MigrationSequence> = {}
 	readonly sortedMigrations: readonly Migration[]
+	private readonly migrationCache = new WeakMap<SerializedSchema, Result<Migration[], string>>()
 
 	private constructor(
 		public readonly types: {
@@ -158,10 +159,17 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 		}
 	}
 
-	// TODO: use a weakmap to store the result of this function
 	public getMigrationsSince(persistedSchema: SerializedSchema): Result<Migration[], string> {
+		// Check cache first
+		const cached = this.migrationCache.get(persistedSchema)
+		if (cached) {
+			return cached
+		}
+
 		const upgradeResult = upgradeSchema(persistedSchema)
 		if (!upgradeResult.ok) {
+			// Cache the error result
+			this.migrationCache.set(persistedSchema, upgradeResult)
 			return upgradeResult
 		}
 		const schema = upgradeResult.value
@@ -178,7 +186,10 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 		}
 
 		if (sequenceIdsToInclude.size === 0) {
-			return Result.ok([])
+			const result = Result.ok([])
+			// Cache the empty result
+			this.migrationCache.set(persistedSchema, result)
+			return result
 		}
 
 		const allMigrationsToInclude = new Set<MigrationId>()
@@ -197,7 +208,10 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 			const idx = this.migrations[sequenceId].sequence.findIndex((m) => m.id === theirVersionId)
 			// todo: better error handling
 			if (idx === -1) {
-				return Result.err('Incompatible schema?')
+				const result = Result.err('Incompatible schema?')
+				// Cache the error result
+				this.migrationCache.set(persistedSchema, result)
+				return result
 			}
 			for (const migration of this.migrations[sequenceId].sequence.slice(idx + 1)) {
 				allMigrationsToInclude.add(migration.id)
@@ -205,7 +219,12 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 		}
 
 		// collect any migrations
-		return Result.ok(this.sortedMigrations.filter(({ id }) => allMigrationsToInclude.has(id)))
+		const result = Result.ok(
+			this.sortedMigrations.filter(({ id }) => allMigrationsToInclude.has(id))
+		)
+		// Cache the result
+		this.migrationCache.set(persistedSchema, result)
+		return result
 	}
 
 	migratePersistedRecord(
@@ -263,7 +282,10 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 		return { type: 'success', value: record }
 	}
 
-	migrateStoreSnapshot(snapshot: StoreSnapshot<R>): MigrationResult<SerializedStore<R>> {
+	migrateStoreSnapshot(
+		snapshot: StoreSnapshot<R>,
+		opts?: { mutateInputStore?: boolean }
+	): MigrationResult<SerializedStore<R>> {
 		let { store } = snapshot
 		const migrations = this.getMigrationsSince(snapshot.schema)
 		if (!migrations.ok) {
@@ -276,7 +298,9 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 			return { type: 'success', value: store }
 		}
 
-		store = structuredClone(store)
+		if (!opts?.mutateInputStore) {
+			store = structuredClone(store)
+		}
 
 		try {
 			for (const migration of migrationsToApply) {
@@ -286,13 +310,13 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 						if (!shouldApply) continue
 						const result = migration.up!(record as any)
 						if (result) {
-							store[id as keyof typeof store] = structuredClone(result) as any
+							store[id as keyof typeof store] = result as any
 						}
 					}
 				} else if (migration.scope === 'store') {
 					const result = migration.up!(store)
 					if (result) {
-						store = structuredClone(result) as any
+						store = result as any
 					}
 				} else {
 					exhaustiveSwitchError(migration)
@@ -325,6 +349,7 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 
 	/**
 	 * @deprecated This is only here for legacy reasons, don't use it unless you have david's blessing!
+	 * @internal
 	 */
 	serializeEarliestVersion(): SerializedSchema {
 		return {

package/src/lib/migrate.ts

@@ -2,45 +2,6 @@ import { assert, objectMapEntries } from '@tldraw/utils'
 import { UnknownRecord } from './BaseRecord'
 import { SerializedStore } from './Store'
 
-let didWarn = false
-
-/**
- * @public
- * @deprecated use `createShapePropsMigrationSequence` instead. See [the docs](https://tldraw.dev/docs/persistence#Updating-legacy-shape-migrations-defineMigrations) for how to migrate.
- */
-export function defineMigrations(opts: {
-	firstVersion?: number
-	currentVersion?: number
-	migrators?: Record<number, LegacyMigration>
-	subTypeKey?: string
-	subTypeMigrations?: Record<string, LegacyBaseMigrationsInfo>
-}): LegacyMigrations {
-	const { currentVersion, firstVersion, migrators = {}, subTypeKey, subTypeMigrations } = opts
-	if (!didWarn) {
-		console.warn(
-			`The 'defineMigrations' function is deprecated and will be removed in a future release. Use the new migrations API instead. See the migration guide for more info: https://tldraw.dev/docs/persistence#Updating-legacy-shape-migrations-defineMigrations`
-		)
-		didWarn = true
-	}
-
-	// Some basic guards against impossible version combinations, some of which will be caught by TypeScript
-	if (typeof currentVersion === 'number' && typeof firstVersion === 'number') {
-		if ((currentVersion as number) === (firstVersion as number)) {
-			throw Error(`Current version is equal to initial version.`)
-		} else if (currentVersion < firstVersion) {
-			throw Error(`Current version is lower than initial version.`)
-		}
-	}
-
-	return {
-		firstVersion: (firstVersion as number) ?? 0, // defaults
-		currentVersion: (currentVersion as number) ?? 0, // defaults
-		migrators,
-		subTypeKey,
-		subTypeMigrations,
-	}
-}
-
 function squashDependsOn(sequence: Array<Migration | StandaloneDependsOn>): Migration[] {
 	const result: Migration[] = []
 	for (let i = sequence.length - 1; i >= 0; i--) {
@@ -199,38 +160,126 @@ export interface MigrationSequence {
 	sequence: Migration[]
 }
 
+/**
+ * Sorts migrations using a distance-minimizing topological sort.
+ *
+ * This function respects two types of dependencies:
+ * 1. Implicit sequence dependencies (foo/1 must come before foo/2)
+ * 2. Explicit dependencies via `dependsOn` property
+ *
+ * The algorithm minimizes the total distance between migrations and their explicit
+ * dependencies in the final ordering, while maintaining topological correctness.
+ * This means when migration A depends on migration B, A will be scheduled as close
+ * as possible to B (while respecting all constraints).
+ *
+ * Implementation uses Kahn's algorithm with priority scoring:
+ * - Builds dependency graph and calculates in-degrees
+ * - Uses priority queue that prioritizes migrations which unblock explicit dependencies
+ * - Processes migrations in urgency order while maintaining topological constraints
+ * - Detects cycles by ensuring all migrations are processed
+ *
+ * @param migrations - Array of migrations to sort
+ * @returns Sorted array of migrations in execution order
+ */
 export function sortMigrations(migrations: Migration[]): Migration[] {
-	// we do a topological sort using dependsOn and implicit dependencies between migrations in the same sequence
-	const byId = new Map(migrations.map((m) => [m.id, m]))
-	const isProcessing = new Set<MigrationId>()
+	if (migrations.length === 0) return []
 
-	const result: Migration[] = []
+	// Build dependency graph and calculate in-degrees
+	const byId = new Map(migrations.map((m) => [m.id, m]))
+	const dependents = new Map<MigrationId, Set<MigrationId>>() // who depends on this
+	const inDegree = new Map<MigrationId, number>()
+	const explicitDeps = new Map<MigrationId, Set<MigrationId>>() // explicit dependsOn relationships
 
-	function process(m: Migration) {
-		assert(!isProcessing.has(m.id), `Circular dependency in migrations: ${m.id}`)
-		isProcessing.add(m.id)
+	// Initialize
+	for (const m of migrations) {
+		inDegree.set(m.id, 0)
+		dependents.set(m.id, new Set())
+		explicitDeps.set(m.id, new Set())
+	}
 
+	// Add implicit sequence dependencies and explicit dependencies
+	for (const m of migrations) {
 		const { version, sequenceId } = parseMigrationId(m.id)
-		const parent = byId.get(`${sequenceId}/${version - 1}`)
-		if (parent) {
-			process(parent)
+
+		// Implicit dependency on previous in sequence
+		const prevId = `${sequenceId}/${version - 1}` as MigrationId
+		if (byId.has(prevId)) {
+			dependents.get(prevId)!.add(m.id)
+			inDegree.set(m.id, inDegree.get(m.id)! + 1)
 		}
 
+		// Explicit dependencies
 		if (m.dependsOn) {
-			for (const dep of m.dependsOn) {
-				const depMigration = byId.get(dep)
-				if (depMigration) {
-					process(depMigration)
+			for (const depId of m.dependsOn) {
+				if (byId.has(depId)) {
+					dependents.get(depId)!.add(m.id)
+					explicitDeps.get(m.id)!.add(depId)
+					inDegree.set(m.id, inDegree.get(m.id)! + 1)
+				}
+			}
+		}
+	}
+
+	// Priority queue: migrations ready to process (in-degree 0)
+	const ready = migrations.filter((m) => inDegree.get(m.id) === 0)
+	const result: Migration[] = []
+	const processed = new Set<MigrationId>()
+
+	while (ready.length > 0) {
+		// Calculate urgency scores for ready migrations and pick the best one
+		let bestCandidate: Migration | undefined
+		let bestCandidateScore = -Infinity
+
+		for (const m of ready) {
+			let urgencyScore = 0
+
+			for (const depId of dependents.get(m.id) || []) {
+				if (!processed.has(depId)) {
+					// Priority 1: Count all unprocessed dependents (to break ties)
+					urgencyScore += 1
+
+					// Priority 2: If this migration is explicitly depended on by others, boost priority
+					if (explicitDeps.get(depId)!.has(m.id)) {
+						urgencyScore += 100
+					}
 				}
 			}
+
+			if (
+				urgencyScore > bestCandidateScore ||
+				// Tiebreaker: prefer lower sequence/version
+				(urgencyScore === bestCandidateScore && m.id.localeCompare(bestCandidate?.id ?? '') < 0)
+			) {
+				bestCandidate = m
+				bestCandidateScore = urgencyScore
+			}
 		}
 
-		byId.delete(m.id)
-		result.push(m)
+		const nextMigration = bestCandidate!
+		ready.splice(ready.indexOf(nextMigration), 1)
+
+		// Cycle detection - if we have processed everything and still have items left, there's a cycle
+		// This is handled by Kahn's algorithm naturally - if we finish with items unprocessed, there's a cycle
+
+		// Process this migration
+		result.push(nextMigration)
+		processed.add(nextMigration.id)
+
+		// Update in-degrees and add newly ready migrations
+		for (const depId of dependents.get(nextMigration.id) || []) {
+			if (!processed.has(depId)) {
+				inDegree.set(depId, inDegree.get(depId)! - 1)
+				if (inDegree.get(depId) === 0) {
+					ready.push(byId.get(depId)!)
+				}
+			}
+		}
 	}
 
-	for (const m of byId.values()) {
-		process(m)
+	// Check for cycles - if we didn't process all migrations, there's a cycle
+	if (result.length !== migrations.length) {
+		const unprocessed = migrations.filter((m) => !processed.has(m.id))
+		assert(false, `Circular dependency in migrations: ${unprocessed[0].id}`)
 	}
 
 	return result

package/src/lib/test/AtomMap.test.ts

@@ -1,4 +1,5 @@
 import { react, transaction } from '@tldraw/state'
+import { vi } from 'vitest'
 import { AtomMap } from '../AtomMap'
 
 describe('AtomMap', () => {
@@ -11,7 +12,7 @@ describe('AtomMap', () => {
 	})
 
 	function testReactor(name: string, fn: () => any) {
-		const cb = jest.fn(fn)
+		const cb = vi.fn(fn)
 		const cleanup = react(name, cb)
 		cleanupFns.push(() => cleanup())
 		return cb

package/src/lib/test/dependsOn.test.ts

@@ -26,7 +26,7 @@ describe('dependsOn', () => {
 				}
 			)
 		}).toThrowErrorMatchingInlineSnapshot(
-			`"Migration 'foo/1' depends on missing migration 'bar/1'"`
+			`[Error: Migration 'foo/1' depends on missing migration 'bar/1']`
 		)
 	})
 
@@ -108,7 +108,7 @@ describe('standalone dependsOn', () => {
 				}
 			)
 		}).toThrowErrorMatchingInlineSnapshot(
-			`"Migration 'foo/1' depends on missing migration 'bar/1'"`
+			`[Error: Migration 'foo/1' depends on missing migration 'bar/1']`
 		)
 	})
 

package/src/lib/test/migrationCaching.test.ts

@@ -0,0 +1,209 @@
+/* eslint-disable @typescript-eslint/no-deprecated */
+import { assert } from '@tldraw/utils'
+import {
+	BaseRecord,
+	Migration,
+	RecordId,
+	createMigrationIds,
+	createMigrationSequence,
+	createRecordType,
+} from '../../'
+import { StoreSchema } from '../StoreSchema'
+
+interface TestRecord extends BaseRecord<'test', RecordId<TestRecord>> {
+	name: string
+	version: number
+}
+
+describe('StoreSchema migration caching', () => {
+	// Create migration IDs
+	const TestVersions = createMigrationIds('com.tldraw.test', {
+		AddVersion: 1,
+		UpdateVersion: 2,
+	})
+
+	// Create a simple schema with migrations
+	const createTestSchema = (version: number) => {
+		const TestRecordType = createRecordType<TestRecord>('test', {
+			scope: 'document',
+		})
+
+		const sequence: Migration[] = []
+
+		if (version > 1) {
+			sequence.push({
+				id: TestVersions.AddVersion,
+				scope: 'record',
+				up: (record: any) => {
+					// Mutate the record in place
+					record.version = 2
+					// Don't return anything
+				},
+				down: (record: any) => {
+					record.version = 1
+					// Don't return anything
+				},
+			})
+		}
+
+		if (version > 2) {
+			sequence.push({
+				id: TestVersions.UpdateVersion,
+				scope: 'record',
+				up: (record: any) => {
+					record.version = 3
+					// Don't return anything
+				},
+				down: (record: any) => {
+					record.version = 2
+					// Don't return anything
+				},
+			})
+		}
+
+		const schema = StoreSchema.create(
+			{
+				test: TestRecordType,
+			},
+			{
+				migrations: [createMigrationSequence({ sequenceId: 'com.tldraw.test', sequence })],
+			}
+		)
+
+		return schema
+	}
+
+	it('should cache migration results and return same array reference', () => {
+		const schema = createTestSchema(3)
+		const oldSchema = schema.serializeEarliestVersion()
+
+		// First call should create the migrations array
+		const migrations1 = schema.getMigrationsSince(oldSchema)
+		assert(migrations1.ok)
+		expect(migrations1.value).toHaveLength(2)
+
+		// Second call should return the same array reference (cached)
+		const migrations2 = schema.getMigrationsSince(oldSchema)
+		assert(migrations2.ok)
+		expect(migrations2.value).toBe(migrations1.value) // Same array reference
+
+		// Third call should also return the same array reference
+		const migrations3 = schema.getMigrationsSince(oldSchema)
+		assert(migrations3.ok)
+		expect(migrations3.value).toBe(migrations1.value)
+	})
+
+	it('should not cache when schema versions are different', () => {
+		const schema = createTestSchema(3)
+		const oldSchema = schema.serializeEarliestVersion()
+
+		// Call with original schema
+		const migrations1 = schema.getMigrationsSince(oldSchema)
+		expect(migrations1.ok).toBe(true)
+		if (!migrations1.ok) throw new Error('Expected migrations1 to be ok')
+
+		// Create a different schema version by using a schema with version 2
+		const schemaV2 = createTestSchema(2)
+		const schemaV2Serialized = schemaV2.serializeEarliestVersion()
+		const migrations2 = schema.getMigrationsSince(schemaV2Serialized)
+		expect(migrations2.ok).toBe(true)
+		if (!migrations2.ok) throw new Error('Expected migrations2 to be ok')
+
+		// Should be different arrays (no cache hit)
+		expect(migrations2.value).not.toBe(migrations1.value)
+	})
+
+	it('should handle mutateInputStore: true with migrators that return void', () => {
+		const schema = createTestSchema(3)
+		const oldSchema = schema.serializeEarliestVersion()
+
+		const store = {
+			test1: {
+				id: 'test1',
+				name: 'Test 1',
+				version: 1,
+				typeName: 'test',
+			},
+			test2: {
+				id: 'test2',
+				name: 'Test 2',
+				version: 1,
+				typeName: 'test',
+			},
+		}
+
+		// Test with mutateInputStore: true
+		const result1 = schema.migrateStoreSnapshot(
+			{ store, schema: oldSchema },
+			{ mutateInputStore: true }
+		)
+
+		assert(result1.type === 'success')
+		expect((result1.value as any).test1.version).toBe(3)
+		expect((result1.value as any).test2.version).toBe(3)
+
+		// The input store should be mutated in place
+		expect(result1.value).toBe(store)
+	})
+
+	it('should handle mutateInputStore: false with migrators that return void', () => {
+		const schema = createTestSchema(3)
+		const oldSchema = schema.serializeEarliestVersion()
+
+		const store = {
+			test1: {
+				id: 'test1',
+				name: 'Test 1',
+				version: 1,
+				typeName: 'test',
+			},
+		}
+
+		// Test with mutateInputStore: false (default)
+		const result = schema.migrateStoreSnapshot({ store, schema: oldSchema })
+
+		assert(result.type === 'success')
+		expect((result.value as any).test1.version).toBe(3)
+
+		// The input store should NOT be mutated
+		expect(store.test1.version).toBe(1)
+	})
+
+	it('should handle empty migration list caching', () => {
+		const schema = createTestSchema(1) // No migrations
+		const oldSchema = schema.serializeEarliestVersion()
+
+		// First call
+		const migrations1 = schema.getMigrationsSince(oldSchema)
+		assert(migrations1.ok)
+
+		expect(migrations1.value).toHaveLength(0)
+
+		// Second call should return same array reference
+		const migrations2 = schema.getMigrationsSince(oldSchema)
+		assert(migrations2.ok)
+		expect(migrations2.value).toBe(migrations1.value)
+		expect(migrations2.value).toHaveLength(0)
+	})
+
+	it('should handle incompatible schema caching', () => {
+		const schema = createTestSchema(3)
+		const incompatibleSchema = {
+			schemaVersion: 1 as const,
+			storeVersion: 1,
+			recordVersions: {
+				test: {
+					version: 999, // Much higher version than what we support
+				},
+			},
+		}
+
+		// First call should fail
+		const migrations1 = schema.getMigrationsSince(incompatibleSchema)
+		expect(migrations1.ok).toBe(false)
+
+		// Second call should also fail (but might be cached)
+		const migrations2 = schema.getMigrationsSince(incompatibleSchema)
+		expect(migrations2.ok).toBe(false)
+	})
+})