enhance: Hoist getEntityCaches computation to MemoCache init

ntucker · ntucker · commit f229a2871f3d · 2025-05-29T09:10:11.000+01:00
diff --git a/packages/normalizr/src/denormalize/localCache.ts b/packages/normalizr/src/denormalize/localCache.ts
@@ -15,7 +15,7 @@ export default class LocalCache implements Cache {
     if (!this.localCache.has(key)) {
       this.localCache.set(key, new Map<string, any>());
     }
-    const localCacheKey = this.localCache.get(key) as Map<string, any>;
+    const localCacheKey = this.localCache.get(key)!;
 
     if (!localCacheKey.get(pk)) {
       computeValue(localCacheKey);
diff --git a/packages/normalizr/src/memo/MemoCache.ts b/packages/normalizr/src/memo/MemoCache.ts
@@ -3,6 +3,7 @@ import WeakDependencyMap from './WeakDependencyMap.js';
 import buildQueryKey from '../buildQueryKey.js';
 import { getDependency, type BaseDelegate } from './BaseDelegate.js';
 import { PlainDelegate } from './Delegate.js';
+import { GetEntityCache, getEntityCaches } from './entitiesCache.js';
 import { getEntities } from '../denormalize/getEntities.js';
 import getUnvisit from '../denormalize/unvisit.js';
 import type {
@@ -12,7 +13,7 @@ import type {
   Schema,
 } from '../interface.js';
 import type { DenormalizeNullable, NormalizeNullable } from '../types.js';
-import { EndpointsCache, EntityCache } from './types.js';
+import { EndpointsCache } from './types.js';
 import type { INVALID } from '../denormalize/symbol.js';
 
 type DelegateClass = new (v: { entities: any; indexes: any }) => BaseDelegate;
@@ -22,15 +23,17 @@ type DelegateClass = new (v: { entities: any; indexes: any }) => BaseDelegate;
 /** Singleton to store the memoization cache for denormalization methods */
 export default class MemoCache {
   /** Cache for every entity based on its dependencies and its own input */
-  protected entities: EntityCache = new Map();
+  declare protected _getCache: GetEntityCache;
   /** Caches the final denormalized form based on input, entities */
   protected endpoints: EndpointsCache = new WeakDependencyMap<EntityPath>();
   /** Caches the queryKey based on schema, args, and any used entities or indexes */
   protected queryKeys: Map<string, WeakDependencyMap<QueryPath>> = new Map();
+
   declare protected Delegate: DelegateClass;
 
   constructor(D: DelegateClass = PlainDelegate) {
     this.Delegate = D;
+    this._getCache = getEntityCaches(new Map());
   }
 
   /** Compute denormalized form maintaining referential equality for same inputs */
@@ -58,7 +61,7 @@ export default class MemoCache {
 
     return getUnvisit(
       getEntity,
-      new GlobalCache(getEntity, this.entities, this.endpoints),
+      new GlobalCache(getEntity, this._getCache, this.endpoints),
       args,
     )(schema, input);
   }
diff --git a/packages/normalizr/src/memo/entitiesCache.ts b/packages/normalizr/src/memo/entitiesCache.ts
@@ -0,0 +1,33 @@
+import { EntityInterface, EntityPath } from '../interface.js';
+import { EntityCache } from './types.js';
+import WeakDependencyMap from './WeakDependencyMap.js';
+
+export type GetEntityCache = (
+  pk: string,
+  schema: EntityInterface,
+) => WeakDependencyMap<EntityPath, object, any>;
+
+export const getEntityCaches = (entityCache: EntityCache): GetEntityCache => {
+  return (pk: string, schema: EntityInterface) => {
+    const key = schema.key;
+    // collections should use the entities they collect over
+    // TODO: this should be based on a public interface
+    const entityInstance: EntityInterface = (schema.cacheWith as any) ?? schema;
+
+    if (!entityCache.has(key)) {
+      entityCache.set(key, new Map());
+    }
+    const entityCacheKey = entityCache.get(key)!;
+
+    if (!entityCacheKey.has(pk)) entityCacheKey.set(pk, new WeakMap());
+    const entityCachePk = entityCacheKey.get(pk)!;
+
+    let wem = entityCachePk.get(entityInstance);
+    if (!wem) {
+      wem = new WeakDependencyMap<EntityPath, object, any>();
+      entityCachePk.set(entityInstance, wem);
+    }
+
+    return wem;
+  };
+};
diff --git a/packages/normalizr/src/memo/globalCache.ts b/packages/normalizr/src/memo/globalCache.ts
@@ -1,32 +1,30 @@
-import { EndpointsCache, EntityCache } from './types.js';
+import { EndpointsCache } from './types.js';
 import WeakDependencyMap, { type Dep } from './WeakDependencyMap.js';
 import type Cache from '../denormalize/cache.js';
 import type { GetEntity } from '../denormalize/getEntities.js';
 import type { INVALID } from '../denormalize/symbol.js';
 import type { EntityInterface, EntityPath } from '../interface.js';
+import type { GetEntityCache } from './entitiesCache.js';
 
 export default class GlobalCache implements Cache {
   private dependencies: Dep<EntityPath>[] = [];
   private cycleCache: Map<string, Map<string, number>> = new Map();
   private cycleIndex = -1;
   private localCache: Map<string, Map<string, any>> = new Map();
 
-  declare private getCache: (
-    pk: string,
-    schema: EntityInterface,
-  ) => WeakDependencyMap<EntityPath, object, any>;
+  declare private _getCache: GetEntityCache;
 
   declare private _getEntity: GetEntity;
-  declare private resultCache: EndpointsCache;
+  declare private _resultCache: EndpointsCache;
 
   constructor(
     getEntity: GetEntity,
-    entityCache: EntityCache,
+    getCache: GetEntityCache,
     resultCache: EndpointsCache,
   ) {
     this._getEntity = getEntity;
-    this.getCache = getEntityCaches(entityCache);
-    this.resultCache = resultCache;
+    this._getCache = getCache;
+    this._resultCache = resultCache;
   }
 
   getEntity(
@@ -43,7 +41,7 @@ export default class GlobalCache implements Cache {
         EntityPath,
         object,
         EntityCacheValue
-      > = this.getCache(pk, schema);
+      > = this._getCache(pk, schema);
       const [cacheValue, cachePath] = globalCache.get(entity, this._getEntity);
       // TODO: what if this just returned the deps - then we don't need to store them
 
@@ -118,15 +116,15 @@ export default class GlobalCache implements Cache {
       return { data: computeValue(), paths: this.paths() };
     }
 
-    let [data, paths] = this.resultCache.get(input, this._getEntity);
+    let [data, paths] = this._resultCache.get(input, this._getEntity);
 
     if (paths === undefined) {
       data = computeValue();
       // we want to do this before we add our 'input' entry
       paths = this.paths();
       // for the first entry, `path` is ignored so empty members is fine
       this.dependencies.unshift({ path: { key: '', pk: '' }, entity: input });
-      this.resultCache.set(this.dependencies, data);
+      this._resultCache.set(this.dependencies, data);
     } else {
       paths.shift();
     }
@@ -142,41 +140,3 @@ interface EntityCacheValue {
   dependencies: Dep<EntityPath>[];
   value: object | typeof INVALID | undefined;
 }
-
-const getEntityCaches = (entityCache: EntityCache) => {
-  return (pk: string, schema: EntityInterface) => {
-    const key = schema.key;
-    // collections should use the entities they collect over
-    // TODO: this should be based on a public interface
-    const entityInstance: EntityInterface = (schema.cacheWith as any) ?? schema;
-
-    if (!entityCache.has(key)) {
-      entityCache.set(key, new Map());
-    }
-    const entityCacheKey = entityCache.get(key)!;
-    if (!entityCacheKey.get(pk))
-      entityCacheKey.set(
-        pk,
-        new WeakMap<
-          EntityInterface,
-          WeakDependencyMap<EntityPath, object, any>
-        >(),
-      );
-
-    const entityCachePk = entityCacheKey.get(pk) as WeakMap<
-      EntityInterface,
-      WeakDependencyMap<EntityPath, object, any>
-    >;
-    let wem = entityCachePk.get(entityInstance) as WeakDependencyMap<
-      EntityPath,
-      object,
-      any
-    >;
-    if (!wem) {
-      wem = new WeakDependencyMap<EntityPath, object, any>();
-      entityCachePk.set(entityInstance, wem);
-    }
-
-    return wem;
-  };
-};
diff --git a/packages/normalizr/src/normalize/getCheckLoop.ts b/packages/normalizr/src/normalize/getCheckLoop.ts
@@ -5,15 +5,13 @@ export function getCheckLoop() {
     if (!visitedEntities.has(entityKey)) {
       visitedEntities.set(entityKey, new Map<string, object[]>());
     }
-    // we have to tell typescript this can't be undefined (due to line above)
-    const entitiesOneType: Map<string, object[]> = visitedEntities.get(
-      entityKey,
-    ) as Map<string, object[]>;
+    const entitiesOneType = visitedEntities.get(entityKey)!;
 
     if (!entitiesOneType.has(pk)) {
       entitiesOneType.set(pk, []);
     }
-    const visitedEntityList = entitiesOneType.get(pk) as object[];
+    const visitedEntityList = entitiesOneType.get(pk)!;
+
     if (visitedEntityList.some((entity: any) => entity === input)) {
       return true;
     }
diff --git a/website/src/components/Playground/editor-types/@data-client/core.d.ts b/website/src/components/Playground/editor-types/@data-client/core.d.ts
@@ -49,19 +49,23 @@ interface NormalizedIndex {
         };
     };
 }
-type EntityPath = [key: string, pk: string];
+/** Used in denormalize. Lookup to find an entity in the store table */
+interface EntityPath {
+    key: string;
+    pk: string;
+}
 type IndexPath = [key: string, index: string, value: string];
 type EntitiesPath = [key: string];
-type QueryPath = IndexPath | EntityPath | EntitiesPath;
+type QueryPath = IndexPath | [key: string, pk: string] | EntitiesPath;
 /** Get all normalized entities of one type from store */
 interface GetEntities {
-    (...path: EntitiesPath): {
+    (key: string): {
         readonly [pk: string]: any;
     } | undefined;
 }
 /** Get normalized Entity from store */
 interface GetEntity {
-    (...path: EntityPath): any;
+    (key: string, pk: string): any;
 }
 /** Get PK using an Entity Index */
 interface GetIndex {
@@ -226,24 +230,24 @@ declare abstract class BaseDelegate {
         indexes: any;
     });
     abstract getEntities(...path: EntitiesPath): object | undefined;
-    abstract getEntity(...path: EntityPath): object | undefined;
+    abstract getEntity(key: string, pk: string): object | undefined;
     abstract getIndex(...path: IndexPath): object | undefined;
     abstract getIndexEnd(entity: any, value: string): string | undefined;
     tracked(schema: any): [delegate: IQueryDelegate, dependencies: Dep<QueryPath>[]];
 }
 
-interface EntityCache extends Map<string, Map<string, WeakMap<EntityInterface, WeakDependencyMap<EntityPath, object, any>>>> {
-}
 type EndpointsCache = WeakDependencyMap<EntityPath, object, any>;
 
+type GetEntityCache = (pk: string, schema: EntityInterface) => WeakDependencyMap<EntityPath, object, any>;
+
 type DelegateClass = new (v: {
     entities: any;
     indexes: any;
 }) => BaseDelegate;
 /** Singleton to store the memoization cache for denormalization methods */
 declare class MemoCache {
     /** Cache for every entity based on its dependencies and its own input */
-    protected entities: EntityCache;
+    protected _getCache: GetEntityCache;
     /** Caches the final denormalized form based on input, entities */
     protected endpoints: EndpointsCache;
     /** Caches the queryKey based on schema, args, and any used entities or indexes */
diff --git a/website/src/components/Playground/editor-types/@data-client/endpoint.d.ts b/website/src/components/Playground/editor-types/@data-client/endpoint.d.ts
@@ -253,7 +253,11 @@ interface Visit {
     (schema: any, value: any, parent: any, key: any, args: readonly any[]): any;
     creating?: boolean;
 }
-type EntityPath = [key: string, pk: string];
+/** Used in denormalize. Lookup to find an entity in the store table */
+interface EntityPath {
+    key: string;
+    pk: string;
+}
 type IndexPath = [key: string, index: string, value: string];
 type EntitiesPath = [key: string];
 /** Returns true if a circular reference is found */
@@ -262,13 +266,13 @@ interface CheckLoop {
 }
 /** Get all normalized entities of one type from store */
 interface GetEntities {
-    (...path: EntitiesPath): {
+    (key: string): {
         readonly [pk: string]: any;
     } | undefined;
 }
 /** Get normalized Entity from store */
 interface GetEntity {
-    (...path: EntityPath): any;
+    (key: string, pk: string): any;
 }
 /** Get PK using an Entity Index */
 interface GetIndex {
diff --git a/website/src/components/Playground/editor-types/@data-client/graphql.d.ts b/website/src/components/Playground/editor-types/@data-client/graphql.d.ts
@@ -253,7 +253,11 @@ interface Visit {
     (schema: any, value: any, parent: any, key: any, args: readonly any[]): any;
     creating?: boolean;
 }
-type EntityPath = [key: string, pk: string];
+/** Used in denormalize. Lookup to find an entity in the store table */
+interface EntityPath {
+    key: string;
+    pk: string;
+}
 type IndexPath = [key: string, index: string, value: string];
 type EntitiesPath = [key: string];
 /** Returns true if a circular reference is found */
@@ -262,13 +266,13 @@ interface CheckLoop {
 }
 /** Get all normalized entities of one type from store */
 interface GetEntities {
-    (...path: EntitiesPath): {
+    (key: string): {
         readonly [pk: string]: any;
     } | undefined;
 }
 /** Get normalized Entity from store */
 interface GetEntity {
-    (...path: EntityPath): any;
+    (key: string, pk: string): any;
 }
 /** Get PK using an Entity Index */
 interface GetIndex {
diff --git a/website/src/components/Playground/editor-types/@data-client/normalizr.d.ts b/website/src/components/Playground/editor-types/@data-client/normalizr.d.ts
@@ -58,23 +58,27 @@ interface EntityTable {
 interface Visit {
     (schema: any, value: any, parent: any, key: any, args: readonly any[]): any;
 }
-type EntityPath = [key: string, pk: string];
+/** Used in denormalize. Lookup to find an entity in the store table */
+interface EntityPath {
+    key: string;
+    pk: string;
+}
 type IndexPath = [key: string, index: string, value: string];
 type EntitiesPath = [key: string];
-type QueryPath = IndexPath | EntityPath | EntitiesPath;
+type QueryPath = IndexPath | [key: string, pk: string] | EntitiesPath;
 /** Returns true if a circular reference is found */
 interface CheckLoop {
     (entityKey: string, pk: string, input: object): boolean;
 }
 /** Get all normalized entities of one type from store */
 interface GetEntities {
-    (...path: EntitiesPath): {
+    (key: string): {
         readonly [pk: string]: any;
     } | undefined;
 }
 /** Get normalized Entity from store */
 interface GetEntity {
-    (...path: EntityPath): any;
+    (key: string, pk: string): any;
 }
 /** Get PK using an Entity Index */
 interface GetIndex {
@@ -266,24 +270,24 @@ declare abstract class BaseDelegate {
         indexes: any;
     });
     abstract getEntities(...path: EntitiesPath): object | undefined;
-    abstract getEntity(...path: EntityPath): object | undefined;
+    abstract getEntity(key: string, pk: string): object | undefined;
     abstract getIndex(...path: IndexPath): object | undefined;
     abstract getIndexEnd(entity: any, value: string): string | undefined;
     tracked(schema: any): [delegate: IQueryDelegate, dependencies: Dep<QueryPath>[]];
 }
 
-interface EntityCache extends Map<string, Map<string, WeakMap<EntityInterface, WeakDependencyMap<EntityPath, object, any>>>> {
-}
 type EndpointsCache = WeakDependencyMap<EntityPath, object, any>;
 
+type GetEntityCache = (pk: string, schema: EntityInterface) => WeakDependencyMap<EntityPath, object, any>;
+
 type DelegateClass = new (v: {
     entities: any;
     indexes: any;
 }) => BaseDelegate;
 /** Singleton to store the memoization cache for denormalization methods */
 declare class MemoCache {
     /** Cache for every entity based on its dependencies and its own input */
-    protected entities: EntityCache;
+    protected _getCache: GetEntityCache;
     /** Caches the final denormalized form based on input, entities */
     protected endpoints: EndpointsCache;
     /** Caches the queryKey based on schema, args, and any used entities or indexes */
diff --git a/website/src/components/Playground/editor-types/@data-client/rest.d.ts b/website/src/components/Playground/editor-types/@data-client/rest.d.ts
diff --git a/website/src/components/Playground/editor-types/globals.d.ts b/website/src/components/Playground/editor-types/globals.d.ts

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ export default class LocalCache implements Cache {`
`15`	`15`	`if (!this.localCache.has(key)) {`
`16`	`16`	`this.localCache.set(key, new Map<string, any>());`
`17`	`17`	`}`
`18`		`- const localCacheKey = this.localCache.get(key) as Map<string, any>;`
	`18`	`+ const localCacheKey = this.localCache.get(key)!;`
`19`	`19`
`20`	`20`	`if (!localCacheKey.get(pk)) {`
`21`	`21`	`computeValue(localCacheKey);`
Original file line number	Diff line number	Diff line change
`@@ -5,15 +5,13 @@ export function getCheckLoop() {`
`5`	`5`	`if (!visitedEntities.has(entityKey)) {`
`6`	`6`	`visitedEntities.set(entityKey, new Map<string, object[]>());`
`7`	`7`	`}`
`8`		`- // we have to tell typescript this can't be undefined (due to line above)`
`9`		`- const entitiesOneType: Map<string, object[]> = visitedEntities.get(`
`10`		`- entityKey,`
`11`		`- ) as Map<string, object[]>;`
	`8`	`+ const entitiesOneType = visitedEntities.get(entityKey)!;`
`12`	`9`
`13`	`10`	`if (!entitiesOneType.has(pk)) {`
`14`	`11`	`entitiesOneType.set(pk, []);`
`15`	`12`	`}`
`16`		`- const visitedEntityList = entitiesOneType.get(pk) as object[];`
	`13`	`+ const visitedEntityList = entitiesOneType.get(pk)!;`
	`14`	`+`
`17`	`15`	`if (visitedEntityList.some((entity: any) => entity === input)) {`
`18`	`16`	`return true;`
`19`	`17`	`}`