Make hat forcing more aggressive (#2602)

## Checklist

- [-] I have added
[tests](https://www.cursorless.org/docs/contributing/test-case-recorder/)
- [-] I have updated the
[docs](https://github.com/cursorless-dev/cursorless/tree/main/docs) and
[cheatsheet](https://github.com/cursorless-dev/cursorless/tree/main/cursorless-talon/src/cheatsheet)
- [-] I have not broken the cheatsheet

packages/common/src/types/HatTokenMap.ts

@@ -6,7 +6,7 @@ import { Token } from "./Token";
  * Maps from (hatStyle, character) pairs to tokens
  */
 export interface HatTokenMap {
-  allocateHats(oldTokenHats?: TokenHat[]): Promise<void>;
+  allocateHats(forceTokenHats?: TokenHat[]): Promise<void>;
   getReadableMap(usePrePhraseSnapshot: boolean): Promise<ReadOnlyHatMap>;
 }
 

packages/cursorless-engine/src/core/HatAllocator.ts

@@ -51,22 +51,30 @@ export class HatAllocator {
   /**
    * Allocate hats to the visible tokens.
    *
-   * @param oldTokenHats If supplied, pretend that this allocation was the
-   * previous allocation when trying to maintain stable hats.  This parameter is
-   * used for testing.
+   * @param forceTokenHats If supplied, force the allocator to use these hats
+   * for the given tokens. This is used for the tutorial, and for testing.
    */
-  async allocateHats(oldTokenHats?: TokenHat[]) {
+  async allocateHats(forceTokenHats?: TokenHat[]) {
     const activeMap = await this.context.getActiveMap();
 
+    // Forced graphemes won't have been normalized
+    forceTokenHats = forceTokenHats?.map((tokenHat) => ({
+      ...tokenHat,
+      grapheme: tokenGraphemeSplitter().normalizeGrapheme(tokenHat.grapheme),
+    }));
+
     const tokenHats = this.hats.isEnabled
-      ? allocateHats(
-          tokenGraphemeSplitter(),
-          this.hats.enabledHatStyles,
-          oldTokenHats ?? activeMap.tokenHats,
-          ide().configuration.getOwnConfiguration("experimental.hatStability"),
-          ide().activeTextEditor,
-          ide().visibleTextEditors,
-        )
+      ? allocateHats({
+          tokenGraphemeSplitter: tokenGraphemeSplitter(),
+          enabledHatStyles: this.hats.enabledHatStyles,
+          forceTokenHats,
+          oldTokenHats: activeMap.tokenHats,
+          hatStability: ide().configuration.getOwnConfiguration(
+            "experimental.hatStability",
+          ),
+          activeTextEditor: ide().activeTextEditor,
+          visibleTextEditors: ide().visibleTextEditors,
+        })
       : [];
 
     activeMap.setTokenHats(tokenHats);

packages/cursorless-engine/src/core/HatTokenMapImpl.ts

@@ -57,12 +57,11 @@ export class HatTokenMapImpl implements HatTokenMap {
   /**
    * Allocate hats to the visible tokens.
    *
-   * @param oldTokenHats If supplied, pretend that this allocation was the
-   * previous allocation when trying to maintain stable hats.  This parameter is
-   * used for testing.
+   * @param forceTokenHats If supplied, force the allocator to use these hats
+   * for the given tokens. This is used for the tutorial, and for testing.
    */
-  allocateHats(oldTokenHats?: TokenHat[]) {
-    return this.hatAllocator.allocateHats(oldTokenHats);
+  allocateHats(forceTokenHats?: TokenHat[]) {
+    return this.hatAllocator.allocateHats(forceTokenHats);
   }
 
   private async getActiveMap() {

packages/cursorless-engine/src/util/allocateHats/allocateHats.ts

@@ -20,6 +20,16 @@ export interface HatCandidate {
   penalty: number;
 }
 
+interface AllocateHatsOptions {
+  tokenGraphemeSplitter: TokenGraphemeSplitter;
+  enabledHatStyles: HatStyleMap;
+  forceTokenHats: readonly TokenHat[] | undefined;
+  oldTokenHats: readonly TokenHat[];
+  hatStability: HatStability;
+  activeTextEditor: TextEditor | undefined;
+  visibleTextEditors: readonly TextEditor[];
+}
+
 /**
  * Allocates hats to all the visible tokens.  Proceeds by ranking tokens
  * according to desirability (how far they are from the cursor), then assigning
@@ -39,14 +49,21 @@ export interface HatCandidate {
  * @returns A hat assignment, which is a list where each entry contains a token
  * and the hat that it will wear
  */
-export function allocateHats(
-  tokenGraphemeSplitter: TokenGraphemeSplitter,
-  enabledHatStyles: HatStyleMap,
-  oldTokenHats: readonly TokenHat[],
-  hatStability: HatStability,
-  activeTextEditor: TextEditor | undefined,
-  visibleTextEditors: readonly TextEditor[],
-): TokenHat[] {
+export function allocateHats({
+  tokenGraphemeSplitter,
+  enabledHatStyles,
+  forceTokenHats,
+  oldTokenHats,
+  hatStability,
+  activeTextEditor,
+  visibleTextEditors,
+}: AllocateHatsOptions): TokenHat[] {
+  /**
+   * Maps from tokens to their forced hat, if any
+   */
+  const forcedHatMap =
+    forceTokenHats == null ? undefined : getTokenOldHatMap(forceTokenHats);
+
   /**
    * Maps from tokens to their assigned hat in previous allocation
    */
@@ -56,7 +73,11 @@ export function allocateHats(
    * A list of tokens in all visible document, ranked by how likely they are to
    * be used.
    */
-  const rankedTokens = getRankedTokens(activeTextEditor, visibleTextEditors);
+  const rankedTokens = getRankedTokens(
+    activeTextEditor,
+    visibleTextEditors,
+    forcedHatMap,
+  );
 
   /**
    * Lookup tables with information about which graphemes / hats appear in which
@@ -102,6 +123,7 @@ export function allocateHats(
         context,
         hatStability,
         tokenRank,
+        forcedHatMap?.get(token),
         tokenOldHatMap.get(token),
         tokenRemainingHatCandidates,
       );

packages/cursorless-engine/src/util/allocateHats/chooseTokenHat.ts

@@ -51,28 +51,32 @@ export function chooseTokenHat(
   { hatOldTokenRanks, graphemeTokenRanks }: RankingContext,
   hatStability: HatStability,
   tokenRank: number,
+  forcedTokenHat: TokenHat | undefined,
   oldTokenHat: TokenHat | undefined,
   candidates: HatCandidate[],
 ): HatCandidate | undefined {
   // We narrow down the candidates by a series of criteria until there is only
   // one left
   return maxByFirstDiffering(candidates, [
-    // 1. Discard any hats that are sufficiently worse than the best hat that we
-    //    wouldn't use them even if they were our old hat
+    // Use forced hat
+    isOldTokenHat(forcedTokenHat),
+
+    // Discard any hats that are sufficiently worse than the best hat that we
+    // wouldn't use them even if they were our old hat
     penaltyEquivalenceClass(hatStability),
 
-    // 2. Use our old hat if it's still in the running
+    // Use our old hat if it's still in the running
     isOldTokenHat(oldTokenHat),
 
-    // 3. Use a free hat if possible; if not, steal the hat of the token with
-    //    lowest rank
+    // Use a free hat if possible; if not, steal the hat of the token with
+    // lowest rank
     hatOldTokenRank(hatOldTokenRanks),
 
-    // 4. Narrow to the hats with the lowest penalty
+    // Narrow to the hats with the lowest penalty
     negativePenalty,
 
-    // 5. Prefer hats that sit on a grapheme that doesn't appear in any highly
-    //    ranked token
+    // Prefer hats that sit on a grapheme that doesn't appear in any highly
+    // ranked token
     minimumTokenRankContainingGrapheme(tokenRank, graphemeTokenRanks),
   ])!;
 }

packages/cursorless-engine/src/util/allocateHats/getRankedTokens.ts

@@ -1,6 +1,10 @@
-import { TextEditor } from "@cursorless/common";
+import {
+  CompositeKeyMap,
+  TextEditor,
+  Token,
+  TokenHat,
+} from "@cursorless/common";
 import { flatten } from "lodash-es";
-import { Token } from "@cursorless/common";
 import { getDisplayLineMap } from "./getDisplayLineMap";
 import { getTokenComparator } from "./getTokenComparator";
 import { getTokensInRange } from "./getTokensInRange";
@@ -13,13 +17,14 @@ import { getTokensInRange } from "./getTokensInRange";
 export function getRankedTokens(
   activeTextEditor: TextEditor | undefined,
   visibleTextEditors: readonly TextEditor[],
+  forcedHatMap: CompositeKeyMap<Token, TokenHat> | undefined,
 ): RankedToken[] {
   const editors: readonly TextEditor[] = getRankedEditors(
     activeTextEditor,
     visibleTextEditors,
   );
 
-  return editors.flatMap((editor) => {
+  const tokens = editors.flatMap((editor) => {
     /**
      * The reference position that will be used to judge how likely a given
      * token is to be used.  Tokens closer to this position will be considered
@@ -44,7 +49,33 @@ export function getRankedTokens(
       ),
     );
 
-    return tokens.map((token, index) => ({ token, rank: -index }));
+    return tokens;
+  });
+
+  return moveForcedHatsToFront(forcedHatMap, tokens).map((token, index) => ({
+    token,
+    rank: -index,
+  }));
+}
+
+function moveForcedHatsToFront(
+  forcedHatMap: CompositeKeyMap<Token, TokenHat> | undefined,
+  tokens: Token[],
+) {
+  if (forcedHatMap == null) {
+    return tokens;
+  }
+
+  return tokens.sort((a, b) => {
+    const aIsForced = forcedHatMap.has(a);
+    const bIsForced = forcedHatMap.has(b);
+    if (aIsForced && !bIsForced) {
+      return -1;
+    }
+    if (!aIsForced && bIsForced) {
+      return 1;
+    }
+    return 0;
   });
 }