From a4e9a2e94ac7034511de7b97349f60654fb79e47 Mon Sep 17 00:00:00 2001
From: Josh Bleecher Snyder <josharian@gmail.com>
Date: Tue, 22 Aug 2023 17:10:10 -0700
Subject: [PATCH 1/2] memoize minimumTokenRankContainingGrapheme

minimumTokenRankContainingGrapheme was accidentally
quadratic in the number of tokens sharing a grapheme.

It was executed for every token, and for each token,
it considered every single other token sharing a candidate grapheme.

It dominated hat allocation performance for larger numbers of tokens.

Memoizing the result by grapheme text makes it linear again.

No functional changes.
---
 .../src/util/allocateHats/HatMetrics.ts       | 36 +++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/packages/cursorless-engine/src/util/allocateHats/HatMetrics.ts b/packages/cursorless-engine/src/util/allocateHats/HatMetrics.ts
index 3aae820a47..5254ae254a 100644
--- a/packages/cursorless-engine/src/util/allocateHats/HatMetrics.ts
+++ b/packages/cursorless-engine/src/util/allocateHats/HatMetrics.ts
@@ -50,8 +50,14 @@ export function minimumTokenRankContainingGrapheme(
   tokenRank: number,
   graphemeTokenRanks: { [key: string]: number[] },
 ): HatMetric {
-  return ({ grapheme: { text } }) =>
-    min(graphemeTokenRanks[text].filter((r) => r > tokenRank)) ?? Infinity;
+  return memoizedHatMetric(
+    ({ grapheme: { text } }): number => {
+      return (
+        min(graphemeTokenRanks[text].filter((r) => r > tokenRank)) ?? Infinity
+      );
+    },
+    ({ grapheme }) => grapheme.text,
+  );
 }
 
 /**
@@ -85,3 +91,29 @@ export function penaltyEquivalenceClass(hatStability: HatStability): HatMetric {
       return (_) => 0;
   }
 }
+
+/**
+ * Memoizes a hat metric based on a key function.
+ * Hat allocation can be highly repetitive across any given dimension
+ * (grapheme, hat style, etc).
+ * This helps us avoid accidentally quadratic behavior in the number of tokens
+ * in minimumTokenRankContainingGrapheme.
+ * @param fn The hat metric to memoize
+ * @param key A function that returns a key for a given hat candidate
+ * @returns A memoized version of the hat metric
+ */
+function memoizedHatMetric(
+  fn: HatMetric,
+  key: (hat: HatCandidate) => any,
+): HatMetric {
+  const cache = new Map<any, number>();
+  return (hat: HatCandidate): number => {
+    const k = key(hat);
+    if (cache.has(k)) {
+      return cache.get(k) as number;
+    }
+    const result = fn(hat);
+    cache.set(k, result);
+    return result;
+  };
+}

From 37ba8bc468cdcc4d1a216a13f130cc8a99d232d5 Mon Sep 17 00:00:00 2001
From: Pokey Rule <755842+pokey@users.noreply.github.com>
Date: Fri, 1 Sep 2023 15:04:54 +0100
Subject: [PATCH 2/2] Attempt to memoize just core function

---
 .../src/util/allocateHats/HatMetrics.ts       | 43 ++++---------------
 1 file changed, 8 insertions(+), 35 deletions(-)

diff --git a/packages/cursorless-engine/src/util/allocateHats/HatMetrics.ts b/packages/cursorless-engine/src/util/allocateHats/HatMetrics.ts
index 5254ae254a..cd747f57a4 100644
--- a/packages/cursorless-engine/src/util/allocateHats/HatMetrics.ts
+++ b/packages/cursorless-engine/src/util/allocateHats/HatMetrics.ts
@@ -1,5 +1,5 @@
 import { CompositeKeyMap, HatStability, TokenHat } from "@cursorless/common";
-import { min } from "lodash";
+import { memoize, min } from "lodash";
 import { HatCandidate } from "./allocateHats";
 
 /**
@@ -50,14 +50,13 @@ export function minimumTokenRankContainingGrapheme(
   tokenRank: number,
   graphemeTokenRanks: { [key: string]: number[] },
 ): HatMetric {
-  return memoizedHatMetric(
-    ({ grapheme: { text } }): number => {
-      return (
-        min(graphemeTokenRanks[text].filter((r) => r > tokenRank)) ?? Infinity
-      );
-    },
-    ({ grapheme }) => grapheme.text,
-  );
+  const coreMetric = memoize((graphemeText: string): number => {
+    return (
+      min(graphemeTokenRanks[graphemeText].filter((r) => r > tokenRank)) ??
+      Infinity
+    );
+  });
+  return ({ grapheme: { text } }) => coreMetric(text);
 }
 
 /**
@@ -91,29 +90,3 @@ export function penaltyEquivalenceClass(hatStability: HatStability): HatMetric {
       return (_) => 0;
   }
 }
-
-/**
- * Memoizes a hat metric based on a key function.
- * Hat allocation can be highly repetitive across any given dimension
- * (grapheme, hat style, etc).
- * This helps us avoid accidentally quadratic behavior in the number of tokens
- * in minimumTokenRankContainingGrapheme.
- * @param fn The hat metric to memoize
- * @param key A function that returns a key for a given hat candidate
- * @returns A memoized version of the hat metric
- */
-function memoizedHatMetric(
-  fn: HatMetric,
-  key: (hat: HatCandidate) => any,
-): HatMetric {
-  const cache = new Map<any, number>();
-  return (hat: HatCandidate): number => {
-    const k = key(hat);
-    if (cache.has(k)) {
-      return cache.get(k) as number;
-    }
-    const result = fn(hat);
-    cache.set(k, result);
-    return result;
-  };
-}