From 40547b47964e03e69ad6fc7371efcab5322198e4 Mon Sep 17 00:00:00 2001
From: Brian Ingles <brianingles@deephaven.io>
Date: Mon, 8 Jul 2024 13:03:51 -0500
Subject: [PATCH] Moved to more generic ExtendedMap class (#3)

---
 src/extension.ts                  |  7 +++----
 src/services/DhService.ts         | 14 +++++++-------
 src/services/DhServiceRegistry.ts |  5 ++---
 src/services/DhcService.ts        |  5 ++---
 src/services/PanelRegistry.ts     | 28 ----------------------------
 src/services/index.ts             |  1 -
 src/util/ExtendedMap.ts           | 17 +++++++++++++++++
 src/util/index.ts                 |  1 +
 8 files changed, 32 insertions(+), 46 deletions(-)
 delete mode 100644 src/services/PanelRegistry.ts
 create mode 100644 src/util/ExtendedMap.ts

diff --git a/src/extension.ts b/src/extension.ts
index 08fc183c..d2ba06be 100644
--- a/src/extension.ts
+++ b/src/extension.ts
@@ -1,20 +1,19 @@
 import * as vscode from 'vscode';
 import {
   ConnectionOption,
+  ExtendedMap,
   createConnectStatusBarItem,
   createConnectTextAndTooltip,
   createConnectionOptions,
   createConnectionQuickPick,
   getTempDir,
 } from './util';
-import { DhcService, PanelRegistry } from './services';
+import { DhcService } from './services';
 import { DhServiceRegistry } from './services';
 import {
   RUN_CODE_COMMAND,
   RUN_SELECTION_COMMAND,
   SELECT_CONNECTION_COMMAND,
-  STATUS_BAR_CONNECTING_TEXT,
-  STATUS_BAR_DISCONNECTED_TEXT,
 } from './common';
 
 export function activate(context: vscode.ExtensionContext) {
@@ -41,7 +40,7 @@ export function activate(context: vscode.ExtensionContext) {
 
   const dhcServiceRegistry = new DhServiceRegistry(
     DhcService,
-    new PanelRegistry(),
+    new ExtendedMap<string, vscode.WebviewPanel>(),
     outputChannel
   );
 
diff --git a/src/services/DhService.ts b/src/services/DhService.ts
index 19318103..2d9b72a6 100644
--- a/src/services/DhService.ts
+++ b/src/services/DhService.ts
@@ -2,10 +2,8 @@ import * as vscode from 'vscode';
 import type { dh as DhcType } from '../dh/dhc-types';
 import { hasErrorCode } from '../util/typeUtils';
 import { ConnectionAndSession } from '../common';
-import { formatTimestamp } from '../util';
-import { PanelFocusManager } from './PanelFocusManager';
+import { ExtendedMap, formatTimestamp } from '../util';
 import { EventDispatcher } from './EventDispatcher';
-import { PanelRegistry } from './PanelRegistry';
 
 /* eslint-disable @typescript-eslint/naming-convention */
 const icons = {
@@ -34,7 +32,7 @@ export abstract class DhService<
 > extends EventDispatcher<'disconnect'> {
   constructor(
     serverUrl: string,
-    panelRegistry: PanelRegistry,
+    panelRegistry: ExtendedMap<string, vscode.WebviewPanel>,
     outputChannel: vscode.OutputChannel
   ) {
     super();
@@ -48,7 +46,7 @@ export abstract class DhService<
   protected readonly subscriptions: (() => void)[] = [];
 
   protected outputChannel: vscode.OutputChannel;
-  private panelRegistry: PanelRegistry;
+  private panelRegistry: ExtendedMap<string, vscode.WebviewPanel>;
   private cachedCreateClient: Promise<TClient> | null = null;
   private cachedCreateSession: Promise<ConnectionAndSession<
     DhcType.IdeConnection,
@@ -289,7 +287,7 @@ export abstract class DhService<
         // );
       }
 
-      const panel = this.panelRegistry.get(title)!;
+      const panel = this.panelRegistry.getOrThrow(title);
       lastPanel = panel;
 
       // See @deprecated comment in PanelFocusManager.onDidChangeViewState
@@ -305,7 +303,9 @@ export abstract class DhService<
           data,
           this.panelRegistry
             .get(title)!
-            .webview.postMessage.bind(this.panelRegistry.get(title)!.webview)
+            .webview.postMessage.bind(
+              this.panelRegistry.getOrThrow(title).webview
+            )
         );
       });
 
diff --git a/src/services/DhServiceRegistry.ts b/src/services/DhServiceRegistry.ts
index 347f46a1..033806bd 100644
--- a/src/services/DhServiceRegistry.ts
+++ b/src/services/DhServiceRegistry.ts
@@ -1,8 +1,7 @@
 import * as vscode from 'vscode';
 import { CacheService } from './CacheService';
 import { DhcService, DhcServiceConstructor } from './DhcService';
-import { ensureHasTrailingSlash } from '../util';
-import { PanelRegistry } from './PanelRegistry';
+import { ensureHasTrailingSlash, ExtendedMap } from '../util';
 
 export class DhServiceRegistry<T extends DhcService> extends CacheService<
   T,
@@ -10,7 +9,7 @@ export class DhServiceRegistry<T extends DhcService> extends CacheService<
 > {
   constructor(
     serviceFactory: DhcServiceConstructor<T>,
-    panelRegistry: PanelRegistry,
+    panelRegistry: ExtendedMap<string, vscode.WebviewPanel>,
     outputChannel: vscode.OutputChannel
   ) {
     super(
diff --git a/src/services/DhcService.ts b/src/services/DhcService.ts
index 665d8c4b..4cd686da 100644
--- a/src/services/DhcService.ts
+++ b/src/services/DhcService.ts
@@ -8,13 +8,12 @@ import {
   initDhcApi,
   initDhcSession,
 } from '../dh/dhc';
-import { getPanelHtml } from '../util';
+import { ExtendedMap, getPanelHtml } from '../util';
 import { ConnectionAndSession } from '../common';
-import { PanelRegistry } from './PanelRegistry';
 
 export type DhcServiceConstructor<T extends DhcService> = new (
   serverUrl: string,
-  panelRegistry: PanelRegistry,
+  panelRegistry: ExtendedMap<string, vscode.WebviewPanel>,
   outputChannel: vscode.OutputChannel
 ) => T;
 
diff --git a/src/services/PanelRegistry.ts b/src/services/PanelRegistry.ts
deleted file mode 100644
index 049374c4..00000000
--- a/src/services/PanelRegistry.ts
+++ /dev/null
@@ -1,28 +0,0 @@
-import * as vscode from 'vscode';
-
-/**
- * Registry of webview panels.
- */
-export class PanelRegistry {
-  private panels = new Map<string, vscode.WebviewPanel>();
-
-  has = (title: string): boolean => {
-    return this.panels.has(title);
-  };
-
-  get = (title: string): vscode.WebviewPanel => {
-    if (!this.has(title)) {
-      throw new Error(`Panel not found: ${title}`);
-    }
-
-    return this.panels.get(title)!;
-  };
-
-  set = (title: string, panel: vscode.WebviewPanel): void => {
-    this.panels.set(title, panel);
-  };
-
-  delete = (title: string): void => {
-    this.panels.delete(title);
-  };
-}
diff --git a/src/services/index.ts b/src/services/index.ts
index 18a28a77..f3d5e4ac 100644
--- a/src/services/index.ts
+++ b/src/services/index.ts
@@ -3,4 +3,3 @@ export * from './Config';
 export * from './DhService';
 export * from './DhcService';
 export * from './DhServiceRegistry';
-export * from './PanelRegistry';
diff --git a/src/util/ExtendedMap.ts b/src/util/ExtendedMap.ts
new file mode 100644
index 00000000..3ec0cbf1
--- /dev/null
+++ b/src/util/ExtendedMap.ts
@@ -0,0 +1,17 @@
+/**
+ * ExtendedMap is a Map with additional utility methods.
+ */
+export class ExtendedMap<K, V> extends Map<K, V> {
+  /**
+   * Return the value for the key if it exists, otherwise throw an error.
+   * @param key
+   * @returns
+   */
+  getOrThrow = (key: K): V => {
+    if (!this.has(key)) {
+      throw new Error(`Key not found: ${key}`);
+    }
+
+    return this.get(key)!;
+  };
+}
diff --git a/src/util/index.ts b/src/util/index.ts
index f902de7e..0eb2c6b3 100644
--- a/src/util/index.ts
+++ b/src/util/index.ts
@@ -1,4 +1,5 @@
 export * from './downloadUtils';
+export * from './ExtendedMap';
 export * from './panelUtils';
 export * from './polyfillUtils';
 export * from './uiUtils';