Fixes

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

@@ -59,15 +59,17 @@ export class HatAllocator {
     const activeMap = await this.context.getActiveMap();
 
     const tokenHats = this.hats.isEnabled
-      ? allocateHats(
-          tokenGraphemeSplitter(),
-          this.hats.enabledHatStyles,
+      ? allocateHats({
+          tokenGraphemeSplitter: tokenGraphemeSplitter(),
+          enabledHatStyles: this.hats.enabledHatStyles,
           forceTokenHats,
-          activeMap.tokenHats,
-          ide().configuration.getOwnConfiguration("experimental.hatStability"),
-          ide().activeTextEditor,
-          ide().visibleTextEditors,
-        )
+          oldTokenHats: activeMap.tokenHats,
+          hatStability: ide().configuration.getOwnConfiguration(
+            "experimental.hatStability",
+          ),
+          activeTextEditor: ide().activeTextEditor,
+          visibleTextEditors: ide().visibleTextEditors,
+        })
       : [];
 
     activeMap.setTokenHats(tokenHats);

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,19 +49,20 @@ 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,
-  forceTokenHats: readonly TokenHat[] | undefined,
-  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 assigned hat in previous allocation
+   * Maps from tokens to their forced hat, if any
    */
-  const forcedHatMap = getTokenOldHatMap(forceTokenHats ?? []);
+  const forcedHatMap =
+    forceTokenHats == null ? undefined : getTokenOldHatMap(forceTokenHats);
 
   /**
    * Maps from tokens to their assigned hat in previous allocation
@@ -62,15 +73,18 @@ 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
    * tokens
    */
   const context = getHatRankingContext(
     rankedTokens,
-    forcedHatMap,
     tokenOldHatMap,
     tokenGraphemeSplitter,
   );
@@ -109,7 +123,7 @@ export function allocateHats(
         context,
         hatStability,
         tokenRank,
-        forcedHatMap.get(token),
+        forcedHatMap?.get(token),
         tokenOldHatMap.get(token),
         tokenRemainingHatCandidates,
       );

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

@@ -8,18 +8,6 @@ import { TokenGraphemeSplitter } from "../../tokenGraphemeSplitter";
 import { RankedToken } from "./getRankedTokens";
 
 export interface RankingContext {
-  /**
-   * Maps from a hat candidate (grapheme+style combination) to the score of the
-   * token that used the given hat in the previous hat allocation.
-   */
-  hatForcedTokenRanks: CompositeKeyMap<
-    {
-      grapheme: string;
-      hatStyle: HatStyleName;
-    },
-    number
-  >;
-
   /**
    * Maps from a hat candidate (grapheme+style combination) to the score of the
    * token that used the given hat in the previous hat allocation.
@@ -43,7 +31,6 @@ export interface RankingContext {
 
 export function getHatRankingContext(
   tokens: RankedToken[],
-  forcedTokenHatMap: CompositeKeyMap<Token, TokenHat>,
   oldTokenHatMap: CompositeKeyMap<Token, TokenHat>,
   tokenGraphemeSplitter: TokenGraphemeSplitter,
 ): RankingContext {
@@ -55,20 +42,12 @@ export function getHatRankingContext(
     { grapheme: string; hatStyle: HatStyleName },
     number
   >(({ grapheme, hatStyle }) => [grapheme, hatStyle]);
-  const hatForcedTokenRanks = new CompositeKeyMap<
-    { grapheme: string; hatStyle: HatStyleName },
-    number
-  >(({ grapheme, hatStyle }) => [grapheme, hatStyle]);
 
   tokens.forEach(({ token, rank }) => {
     const existingTokenHat = oldTokenHatMap.get(token);
     if (existingTokenHat != null) {
       hatOldTokenRanks.set(existingTokenHat, rank);
     }
-    const forcedTokenHat = forcedTokenHatMap.get(token);
-    if (forcedTokenHat != null) {
-      hatForcedTokenRanks.set(forcedTokenHat, rank);
-    }
     tokenGraphemeSplitter
       .getTokenGraphemes(token.text)
       .forEach(({ text: graphemeText }) => {
@@ -86,7 +65,6 @@ export function getHatRankingContext(
   });
 
   return {
-    hatForcedTokenRanks,
     hatOldTokenRanks,
     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;
   });
 }