API: make object registries project-specific to avoid id collisions on types and signatures (#4341)

Author: piotrtomiak

_packages/native-preview/src/api/async/api.ts

@@ -31,10 +31,6 @@ import {
     readSourceFileHash,
     RemoteSourceFile,
 } from "../node/node.ts";
-import {
-    type ObjectFactories,
-    ObjectRegistry,
-} from "../objectRegistry.ts";
 import type {
     APIOptions,
     LSPConnectionOptions,
@@ -236,9 +232,9 @@ export class Snapshot {
     private projectMap: Map<Path, Project>;
     private toPath: (fileName: string) => Path;
     private client: Client;
-    private objectRegistry: SnapshotObjectRegistry;
     private disposed: boolean = false;
     private onDispose: () => void;
+    private snapshotRegistry: SnapshotObjectRegistry;
 
     constructor(
         data: UpdateSnapshotResponse,
@@ -251,21 +247,11 @@ export class Snapshot {
         this.client = client;
         this.toPath = toPath;
         this.onDispose = onDispose;
+        this.snapshotRegistry = new SnapshotObjectRegistry(client, this.id);
 
-        this.objectRegistry = new SnapshotObjectRegistry(
-            {
-                createSymbol: symbolData => new Symbol(symbolData, this.objectRegistry),
-                createType: typeData => new TypeObject(typeData, this.objectRegistry),
-                createSignature: sigData => new Signature(sigData, this.objectRegistry),
-            },
-            client,
-            this.id,
-        );
-
-        // Create projects
         this.projectMap = new Map();
         for (const projData of data.projects) {
-            const project = new Project(projData, this.id, client, this.objectRegistry, sourceFileCache, toPath);
+            const project = new Project(projData, this.id, client, sourceFileCache, toPath, this.snapshotRegistry);
             this.projectMap.set(toPath(projData.configFileName), project);
         }
     }
@@ -297,7 +283,11 @@ export class Snapshot {
     async dispose(): Promise<void> {
         if (this.disposed) return;
         this.disposed = true;
-        this.objectRegistry.clear();
+        for (const project of this.projectMap.values()) {
+            project.dispose();
+        }
+        this.projectMap.clear();
+        this.snapshotRegistry.clear();
         this.onDispose();
         await this.client.apiRequest("release", { snapshot: this.id });
     }
@@ -313,18 +303,128 @@ export class Snapshot {
     }
 }
 
-class SnapshotObjectRegistry extends ObjectRegistry<Symbol, TypeObject, Signature> {
+class SnapshotObjectRegistry {
+    private readonly symbols: Map<number, Symbol> = new Map();
+    private readonly client: Client;
+    private readonly snapshotId: number;
+
+    constructor(client: Client, snapshotId: number) {
+        this.client = client;
+        this.snapshotId = snapshotId;
+    }
+
+    getOrCreateSymbol(data: SymbolResponse): Symbol {
+        let symbol = this.symbols.get(data.id);
+        if (!symbol) {
+            symbol = new Symbol(data, this);
+            this.symbols.set(data.id, symbol);
+        }
+        return symbol;
+    }
+
+    getSymbol(id: number): Symbol | undefined {
+        return this.symbols.get(id);
+    }
+
+    clear(): void {
+        this.symbols.clear();
+    }
+
+    async fetchSymbol(source: Symbol | Signature | Type, method: string, handle: number | undefined, projectId?: string): Promise<Symbol> {
+        if (!handle) return undefined as unknown as Symbol;
+        const cached = this.getSymbol(handle);
+        if (cached) return cached;
+
+        const data = await this.client.apiRequest<SymbolResponse | null>(method, {
+            snapshot: this.snapshotId,
+            project: projectId,
+            objectId: source.id,
+        });
+        if (!data) throw new Error(`${method} returned null symbol for ${source.constructor.name} ${source.id}`);
+        return this.getOrCreateSymbol(data);
+    }
+
+    async fetchSymbols(source: Symbol | Signature | Type, method: string, handles?: readonly number[], projectId?: string): Promise<readonly Symbol[]> {
+        if (handles) {
+            const result = new Array<Symbol>(handles.length);
+            let allCached = true;
+            for (let i = 0; i < handles.length; i++) {
+                const cached = this.getSymbol(handles[i]);
+                if (!cached) {
+                    allCached = false;
+                    break;
+                }
+                result[i] = cached;
+            }
+            if (allCached) return result;
+        }
+        const symbolData = await this.client.apiRequest<SymbolResponse[] | null>(method, {
+            snapshot: this.snapshotId,
+            project: projectId,
+            objectId: source.id,
+        });
+        if (symbolData == null) return [];
+        else return symbolData.map(data => this.getOrCreateSymbol(data));
+    }
+}
+
+class ProjectObjectRegistry {
     private client: Client;
     private snapshotId: number;
+    private projectId: string;
+    private snapshotRegistry: SnapshotObjectRegistry;
+    private types: Map<number, TypeObject> = new Map();
+    private signatures: Map<number, Signature> = new Map();
 
     constructor(
-        factories: ObjectFactories<Symbol, TypeObject, Signature>,
         client: Client,
         snapshotId: number,
+        projectId: string,
+        snapshotRegistry: SnapshotObjectRegistry,
     ) {
-        super(factories);
         this.client = client;
         this.snapshotId = snapshotId;
+        this.projectId = projectId;
+        this.snapshotRegistry = snapshotRegistry;
+    }
+
+    getOrCreateSymbol(data: SymbolResponse): Symbol {
+        return this.snapshotRegistry.getOrCreateSymbol(data);
+    }
+
+    getSymbol(id: number): Symbol | undefined {
+        return this.snapshotRegistry.getSymbol(id);
+    }
+
+    getOrCreateType(data: TypeResponse): TypeObject {
+        let type = this.types.get(data.id);
+        if (!type) {
+            type = new TypeObject(data, this);
+            this.types.set(data.id, type);
+        }
+        return type;
+    }
+
+    getType(id: number): TypeObject | undefined {
+        return this.types.get(id);
+    }
+
+    getOrCreateSignature(data: SignatureResponse): Signature {
+        let sig = this.signatures.get(data.id);
+        if (!sig) {
+            sig = new Signature(data, this);
+            this.signatures.set(data.id, sig);
+        }
+        return sig;
+    }
+
+    getSignature(id: number): Signature | undefined {
+        return this.signatures.get(id);
+    }
+
+    clear(): void {
+        this.types.clear();
+        this.signatures.clear();
     }
 
     async fetchType<T extends Type>(source: Symbol | Signature | Type, method: string, handle: number | undefined): Promise<T> {
@@ -334,23 +434,15 @@ class SnapshotObjectRegistry extends ObjectRegistry<Symbol, TypeObject, Signatur
 
         const data = await this.client.apiRequest<TypeResponse | null>(method, {
             snapshot: this.snapshotId,
+            project: this.projectId,
             objectId: source.id,
         });
         if (!data) throw new Error(`${method} returned null type for ${source.constructor.name} ${source.id}`);
         return this.getOrCreateType(data) as unknown as T;
     }
 
     async fetchSymbol(source: Symbol | Signature | Type, method: string, handle: number | undefined): Promise<Symbol> {
-        if (!handle) return undefined as unknown as Symbol;
-        const cached = this.getSymbol(handle);
-        if (cached) return cached;
-
-        const data = await this.client.apiRequest<SymbolResponse | null>(method, {
-            snapshot: this.snapshotId,
-            objectId: source.id,
-        });
-        if (!data) throw new Error(`${method} returned null symbol for ${source.constructor.name} ${source.id}`);
-        return this.getOrCreateSymbol(data);
+        return this.snapshotRegistry.fetchSymbol(source, method, handle, this.projectId);
     }
 
     async fetchSignature(source: Symbol | Signature | Type, method: string, handle: number | undefined): Promise<Signature> {
@@ -360,6 +452,7 @@ class SnapshotObjectRegistry extends ObjectRegistry<Symbol, TypeObject, Signatur
 
         const data = await this.client.apiRequest<SignatureResponse | null>(method, {
             snapshot: this.snapshotId,
+            project: this.projectId,
             objectId: source.id,
         });
         if (!data) throw new Error(`${method} returned null signature for ${source.constructor.name} ${source.id}`);
@@ -382,32 +475,15 @@ class SnapshotObjectRegistry extends ObjectRegistry<Symbol, TypeObject, Signatur
         }
         const typesData = await this.client.apiRequest<TypeResponse[] | null>(method, {
             snapshot: this.snapshotId,
+            project: this.projectId,
             objectId: source.id,
         });
         if (typesData == null) return [];
         else return typesData.map(data => this.getOrCreateType(data));
     }
 
     async fetchSymbols(source: Symbol | Signature | Type, method: string, handles?: readonly number[]): Promise<readonly Symbol[]> {
-        if (handles) {
-            const result = new Array<Symbol>(handles.length);
-            let allCached = true;
-            for (let i = 0; i < handles.length; i++) {
-                const cached = this.getSymbol(handles[i]);
-                if (!cached) {
-                    allCached = false;
-                    break;
-                }
-                result[i] = cached;
-            }
-            if (allCached) return result;
-        }
-        const symbolData = await this.client.apiRequest<SymbolResponse[] | null>(method, {
-            snapshot: this.snapshotId,
-            objectId: source.id,
-        });
-        if (symbolData == null) return [];
-        else return symbolData.map(data => this.getOrCreateSymbol(data));
+        return this.snapshotRegistry.fetchSymbols(source, method, handles, this.projectId);
     }
 }
 
@@ -426,9 +502,9 @@ export class Project {
         data: ProjectResponse,
         snapshotId: number,
         client: Client,
-        objectRegistry: SnapshotObjectRegistry,
         sourceFileCache: SourceFileCache,
         toPath: (fileName: string) => Path,
+        snapshotRegistry: SnapshotObjectRegistry,
     ) {
         this.id = data.id;
         this.configFileName = data.configFileName;
@@ -442,6 +518,7 @@ export class Project {
             sourceFileCache,
             toPath,
         );
+        const objectRegistry = new ProjectObjectRegistry(client, snapshotId, this.id, snapshotRegistry);
         this.checker = new Checker(
             snapshotId,
             this.id,
@@ -450,6 +527,10 @@ export class Project {
         );
         this.emitter = new Emitter(client);
     }
+
+    dispose(): void {
+        this.checker.dispose();
+    }
 }
 
 export class Program {
@@ -571,20 +652,24 @@ export class Checker {
     private snapshotId: number;
     private projectId: string;
     private client: Client;
-    private objectRegistry: SnapshotObjectRegistry;
+    private objectRegistry: ProjectObjectRegistry;
 
     constructor(
         snapshotId: number,
         projectId: string,
         client: Client,
-        objectRegistry: SnapshotObjectRegistry,
+        objectRegistry: ProjectObjectRegistry,
     ) {
         this.snapshotId = snapshotId;
         this.projectId = projectId;
         this.client = client;
         this.objectRegistry = objectRegistry;
     }
 
+    dispose(): void {
+        this.objectRegistry.clear();
+    }
+
     getSymbolAtLocation(node: Node): Promise<Symbol | undefined>;
     getSymbolAtLocation(nodes: readonly Node[]): Promise<(Symbol | undefined)[]>;
     async getSymbolAtLocation(nodeOrNodes: Node | readonly Node[]): Promise<Symbol | (Symbol | undefined)[] | undefined> {
@@ -1166,7 +1251,7 @@ export class Symbol {
 }
 
 class TypeObject implements Type {
-    private objectRegistry: SnapshotObjectRegistry;
+    private objectRegistry: ProjectObjectRegistry;
 
     readonly id: number;
     readonly flags: TypeFlags;
@@ -1194,7 +1279,7 @@ class TypeObject implements Type {
     readonly baseType!: number;
     readonly substConstraint!: number;
 
-    constructor(data: TypeResponse, objectRegistry: SnapshotObjectRegistry) {
+    constructor(data: TypeResponse, objectRegistry: ProjectObjectRegistry) {
         this.objectRegistry = objectRegistry;
 
         this.id = data.id;
@@ -1295,7 +1380,7 @@ class TypeObject implements Type {
 
 export class Signature {
     private flags: number;
-    private objectRegistry: SnapshotObjectRegistry;
+    private objectRegistry: ProjectObjectRegistry;
 
     readonly id: number;
     readonly declaration?: NodeHandle | undefined;
@@ -1304,7 +1389,7 @@ export class Signature {
     readonly thisParameter?: number | undefined;
     readonly target?: number | undefined;
 
-    constructor(data: SignatureResponse, objectRegistry: SnapshotObjectRegistry) {
+    constructor(data: SignatureResponse, objectRegistry: ProjectObjectRegistry) {
         this.id = data.id;
         this.flags = data.flags;
         this.objectRegistry = objectRegistry;