Deduplicate multi-map entry classes globally

UCNot · Jul 17, 2023 · 355d1fc · 355d1fc
1 parent 004fda2
commit 355d1fc
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 13 deletions.
diff --git a/src/compiler/deserialization/map.ucrx.class.ts b/src/compiler/deserialization/map.ucrx.class.ts
@@ -57,7 +57,6 @@ export class MapUcrxClass<
   readonly #collector: MapUcrxClass$Collector;
   readonly #store: MapUcrxStore;
   readonly #slot: EsFieldHandle;
-  readonly #rxClasses = new Map<UcrxClass, UcrxClass>();
 
   constructor(
     lib: UcrxLib,
@@ -414,18 +413,7 @@ export class MapUcrxClass<
   entryUcrxFor(key: string | null, schema: UcSchema): UcrxClass {
     const entryClass = MapUcrxEntry.ucrxClass(this.#lib, this.schema, key, schema);
 
-    if (!this.#collector.rxs) {
-      return entryClass;
-    }
-
-    let rxClass = this.#rxClasses.get(entryClass);
-
-    if (!rxClass) {
-      rxClass = new MultiEntryUcrxClass(entryClass);
-      this.#rxClasses.set(entryClass, rxClass);
-    }
-
-    return rxClass;
+    return this.#collector.rxs ? entryClass.associate(MultiEntryUcrxClass$associate) : entryClass;
   }
 
   createEntry(key: string | null, schema: UcSchema): MapUcrxEntry {
@@ -549,3 +537,11 @@ class MultiEntryUcrxClass<
   }
 
 }
+
+function MultiEntryUcrxClass$associate<
+  TArgs extends UcrxSignature.Args = UcrxSignature.Args,
+  T = unknown,
+  TSchema extends UcSchema<T> = UcSchema<T>,
+>(target: UcrxClass<TArgs, T, TSchema>): MultiEntryUcrxClass<TArgs, T, TSchema> {
+  return new MultiEntryUcrxClass(target);
+}
diff --git a/src/compiler/rx/ucrx.class.ts b/src/compiler/rx/ucrx.class.ts
@@ -15,6 +15,7 @@ export abstract class UcrxClass<
   readonly #typeName: string;
   readonly #methodMods = new Map<UcrxMethod<EsSignature.Args, any>, unknown[]>();
   #supportedTypes?: ReadonlySet<string>;
+  #associations?: Map<UcrxClass.Association<unknown, TArgs, T, TSchema, this>, unknown>;
 
   constructor(init: UcrxClass.Init<TArgs, T, TSchema>) {
     const { schema, typeName = ucSchemaTypeSymbol(schema), declare = { at: 'bundle' } } = init;
@@ -87,6 +88,22 @@ export abstract class UcrxClass<
     }
   }
 
+  associate<TAssoc>(associate: UcrxClass.Association<TAssoc, TArgs, T, TSchema, this>): TAssoc {
+    let association: TAssoc | undefined;
+
+    if (!this.#associations) {
+      this.#associations = new Map();
+    } else {
+      association = this.#associations.get(associate) as TAssoc | undefined;
+    }
+    if (!association) {
+      association = associate(this);
+      this.#associations.set(associate, association);
+    }
+
+    return association;
+  }
+
   protected discoverTypes(types: Set<string>): void {
     for (const memberRef of this.members()) {
       const { member, declared } = memberRef;
@@ -127,6 +144,14 @@ export namespace UcrxClass {
     readonly schema: TSchema;
     readonly typeName?: string | undefined;
   };
+
+  export type Association<
+    out TAssoc = unknown,
+    out TArgs extends UcrxSignature.Args = UcrxSignature.Args,
+    out T = unknown,
+    out TSchema extends UcSchema<T> = UcSchema<T>,
+    in TTarget extends UcrxClass<TArgs, T, TSchema> = UcrxClass<TArgs, T, TSchema>,
+  > = { associate(target: TTarget): TAssoc }['associate'];
 }
 
 export type UcrxProto<out T = unknown, out TSchema extends UcSchema<T> = UcSchema<T>> = {