diff --git a/polygerrit-ui/app/embed/diff/gr-diff-builder/token-highlight-layer.ts b/polygerrit-ui/app/embed/diff/gr-diff-builder/token-highlight-layer.ts
index b96b156..19e0e22 100644
--- a/polygerrit-ui/app/embed/diff/gr-diff-builder/token-highlight-layer.ts
+++ b/polygerrit-ui/app/embed/diff/gr-diff-builder/token-highlight-layer.ts
@@ -3,7 +3,7 @@
  * Copyright 2021 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
-import {DiffLayer, DiffLayerListener} from '../../../types/types';
+import {DiffLayer} from '../../../types/types';
 import {GrDiffLine, Side, TokenHighlightListener} from '../../../api/diff';
 import {assertIsDefined} from '../../../utils/common-util';
 import {GrAnnotation} from '../gr-diff-highlight/gr-annotation';
@@ -44,12 +44,6 @@
 const TOKEN_OCCURRENCES_LIMIT = 1000;
 
 /**
- * Token highlighting is only useful for code on-screen, so we only highlight
- * the nearest set of tokens up to this limit.
- */
-const TOKEN_HIGHLIGHT_LIMIT = 100;
-
-/**
  * When a user hovers over a token in the diff, then this layer makes sure that
  * all occurrences of this token are annotated with the 'token-highlight' css
  * class. And removes that class when the user moves the mouse away from the
@@ -61,9 +55,6 @@
  * And when that re-rendering happens the appropriate css class is added.
  */
 export class TokenHighlightLayer implements DiffLayer {
-  /** The only listener is typically the renderer of gr-diff. */
-  private listeners: DiffLayerListener[] = [];
-
   /** The currently highlighted token. */
   private currentHighlight?: string;
 
@@ -92,9 +83,7 @@
    * Keeps track of where tokens occur in a file during rendering, so that it is
    * easy to look up when processing mouse events.
    */
-  private tokenToLinesLeft = new Map<string, Set<number>>();
-
-  private tokenToLinesRight = new Map<string, Set<number>>();
+  private tokenToElements = new Map<string, Set<HTMLElement>>();
 
   private hoveredElement?: Element;
 
@@ -110,12 +99,7 @@
     });
   }
 
-  annotate(
-    el: HTMLElement,
-    _: HTMLElement,
-    line: GrDiffLine,
-    side: Side
-  ): void {
+  annotate(el: HTMLElement, _1: HTMLElement, _2: GrDiffLine, _3: Side): void {
     const text = el.textContent;
     if (!text) return;
     // Binary files encoded as text for example can have super long lines
@@ -132,7 +116,7 @@
       if (length > TOKEN_LENGTH_LIMIT) continue;
       atLeastOneTokenMatched = true;
       const highlightTypeClass =
-        token === this.currentHighlight ? CSS_HIGHLIGHT : CSS_TOKEN;
+        token === this.currentHighlight ? CSS_HIGHLIGHT : '';
       const textClass = `${TOKEN_TEXT_PREFIX}${token}`;
       const indexClass = `${TOKEN_INDEX_PREFIX}${index}`;
       // We add the TOKEN_TEXT_PREFIX class so that we can look up the token later easily
@@ -142,12 +126,12 @@
         el,
         index,
         length,
-        `${textClass} ${indexClass} ${highlightTypeClass}`
+        `${textClass} ${indexClass} ${CSS_TOKEN} ${highlightTypeClass}`
       );
       // We could try to detect whether we are re-rendering instead of initially
       // rendering the line. Then we would not have to call storeLineForToken()
       // again. But since the Set swallows the duplicates we don't care.
-      this.storeLineForToken(token, line, side);
+      this.storeElementsForToken(token, el, textClass);
     }
     if (atLeastOneTokenMatched) {
       // These listeners do not have to be cleaned, because listeners are
@@ -162,21 +146,23 @@
     }
   }
 
-  private storeLineForToken(token: string, line: GrDiffLine, side: Side) {
-    const tokenToLines =
-      side === Side.LEFT ? this.tokenToLinesLeft : this.tokenToLinesRight;
-    // Just to make sure that we don't break down on large files.
-    if (tokenToLines.size > TOKEN_COUNT_LIMIT) return;
-    let numbers = tokenToLines.get(token);
-    if (!numbers) {
-      numbers = new Set<number>();
-      tokenToLines.set(token, numbers);
+  private storeElementsForToken(
+    token: string,
+    lineEl: HTMLElement,
+    cssClass: string
+  ) {
+    for (const el of lineEl.querySelectorAll(`.${cssClass}`)) {
+      let tokenEls = this.tokenToElements.get(token);
+      if (!tokenEls) {
+        // Just to make sure that we don't break down on large files.
+        if (this.tokenToElements.size > TOKEN_COUNT_LIMIT) return;
+        tokenEls = new Set<HTMLElement>();
+        this.tokenToElements.set(token, tokenEls);
+      }
+      // Just to make sure that we don't break down on large files.
+      if (tokenEls.size > TOKEN_OCCURRENCES_LIMIT) return;
+      tokenEls.add(el as HTMLElement);
     }
-    // Just to make sure that we don't break down on large files.
-    if (numbers.size > TOKEN_OCCURRENCES_LIMIT) return;
-    const lineNumber =
-      side === Side.LEFT ? line.beforeNumber : line.afterNumber;
-    numbers.add(Number(lineNumber));
   }
 
   private handleTokenMouseOut(e: MouseEvent) {
@@ -234,10 +220,7 @@
   } {
     if (!(el instanceof Element))
       return {line: 0, token: undefined, element: undefined};
-    if (
-      el.classList.contains(CSS_TOKEN) ||
-      el.classList.contains(CSS_HIGHLIGHT)
-    ) {
+    if (el.classList.contains(CSS_TOKEN)) {
       const tkTextClass = [...el.classList].find(c =>
         c.startsWith(TOKEN_TEXT_PREFIX)
       );
@@ -265,8 +248,8 @@
       this.currentHighlightLineNumber === newLineNumber
     )
       return;
+
     const oldHighlight = this.currentHighlight;
-    const oldLineNumber = this.currentHighlightLineNumber;
     this.currentHighlight = newHighlight;
     this.currentHighlightLineNumber = newLineNumber;
     this.triggerTokenHighlightEvent(
@@ -274,8 +257,21 @@
       newLineNumber,
       newHoveredElement
     );
-    this.notifyForToken(oldHighlight, oldLineNumber);
-    this.notifyForToken(newHighlight, newLineNumber);
+    this.toggleTokenHighlight(oldHighlight, CSS_HIGHLIGHT);
+    this.toggleTokenHighlight(newHighlight, CSS_HIGHLIGHT);
+  }
+
+  private toggleTokenHighlight(token: string | undefined, cssClass: string) {
+    if (!token) {
+      return;
+    }
+    const tokenEls = this.tokenToElements.get(token);
+    if (!tokenEls) {
+      return;
+    }
+    for (const el of tokenEls) {
+      el.classList.toggle(cssClass);
+    }
   }
 
   triggerTokenHighlightEvent(
@@ -306,58 +302,4 @@
     };
     this.tokenHighlightListener({token, element, side, range});
   }
-
-  getSortedLinesForSide(
-    lineMapping: Map<string, Set<number>>,
-    token: string | undefined,
-    lineNumber: number
-  ): Array<number> {
-    if (!token) return [];
-    const lineSet = lineMapping.get(token);
-    if (!lineSet) return [];
-    const lines = [...lineSet];
-    lines.sort((a, b) => {
-      const da = Math.abs(a - lineNumber);
-      const db = Math.abs(b - lineNumber);
-      // For equal distance, prefer lines later in the file over earlier in the
-      // file. This ensures total ordering.
-      if (da === db) return b - a;
-      // Compare the distance to lineNumber.
-      return da - db;
-    });
-    return lines.slice(0, TOKEN_HIGHLIGHT_LIMIT);
-  }
-
-  notifyForToken(token: string | undefined, lineNumber: number) {
-    const leftLines = this.getSortedLinesForSide(
-      this.tokenToLinesLeft,
-      token,
-      lineNumber
-    );
-    for (const line of leftLines) {
-      this.notifyListeners(line, Side.LEFT);
-    }
-    const rightLines = this.getSortedLinesForSide(
-      this.tokenToLinesRight,
-      token,
-      lineNumber
-    );
-    for (const line of rightLines) {
-      this.notifyListeners(line, Side.RIGHT);
-    }
-  }
-
-  addListener(listener: DiffLayerListener) {
-    this.listeners.push(listener);
-  }
-
-  removeListener(listener: DiffLayerListener) {
-    this.listeners = this.listeners.filter(f => f !== listener);
-  }
-
-  notifyListeners(line: number, side: Side) {
-    for (const listener of this.listeners) {
-      listener(line, line, side);
-    }
-  }
 }
diff --git a/polygerrit-ui/app/embed/diff/gr-diff-builder/token-highlight-layer_test.ts b/polygerrit-ui/app/embed/diff/gr-diff-builder/token-highlight-layer_test.ts
index ef2e3d9..c0812b7 100644
--- a/polygerrit-ui/app/embed/diff/gr-diff-builder/token-highlight-layer_test.ts
+++ b/polygerrit-ui/app/embed/diff/gr-diff-builder/token-highlight-layer_test.ts
@@ -55,6 +55,7 @@
     highlightDetails?: TokenHighlightEventDetails
   ) {
     tokenHighlightingCalls.push({details: highlightDetails});
+    listener.notify({details: highlightDetails});
   }
 
   setup(async () => {
@@ -63,7 +64,6 @@
     container = document.createElement('div');
     document.body.appendChild(container);
     highlighter = new TokenHighlightLayer(container, tokenHighlightListener);
-    highlighter.addListener((...args) => listener.notify(...args));
   });
 
   teardown(() => {
@@ -126,21 +126,21 @@
         el,
         0,
         5,
-        'tk-text-these tk-index-0 token'
+        'tk-text-these tk-index-0 token '
       );
       assertAnnotation(
         annotateElementStub.args[1],
         el,
         6,
         3,
-        'tk-text-are tk-index-6 token'
+        'tk-text-are tk-index-6 token '
       );
       assertAnnotation(
         annotateElementStub.args[2],
         el,
         10,
         5,
-        'tk-text-words tk-index-10 token'
+        'tk-text-words tk-index-10 token '
       );
     });
 
@@ -196,10 +196,8 @@
 
       // After a total of HOVER_DELAY_MS ms the hover behavior should trigger.
       clock.tick(HOVER_DELAY_MS - 100);
-      assert.equal(listener.pending, 2);
+      assert.equal(listener.pending, 1);
       assert.equal(_testOnly_allTasks.size, 0);
-      assert.deepEqual(listener.shift(), [1, 1, Side.LEFT]);
-      assert.deepEqual(listener.shift(), [2, 2, Side.RIGHT]);
     });
 
     test('highlighting spans many lines', async () => {
@@ -216,10 +214,14 @@
 
       // After a total of HOVER_DELAY_MS ms the hover behavior should trigger.
       clock.tick(HOVER_DELAY_MS);
-      assert.equal(listener.pending, 2);
+      assert.equal(listener.pending, 1);
       assert.equal(_testOnly_allTasks.size, 0);
-      assert.deepEqual(listener.shift(), [1, 1, Side.LEFT]);
-      assert.deepEqual(listener.shift(), [1000, 1000, Side.RIGHT]);
+      assert.deepEqual(tokenHighlightingCalls[0].details, {
+        token: 'words',
+        side: Side.RIGHT,
+        element: words1,
+        range: {start_line: 1, start_column: 5, end_line: 1, end_column: 9},
+      });
     });
 
     test('highlighting mouse out before delay', async () => {
@@ -259,10 +261,12 @@
         element: words1,
         range: {start_line: 1, start_column: 5, end_line: 1, end_column: 9},
       });
+      assert.isTrue(words1.classList.contains('token-highlight'));
 
       container.click();
       assert.equal(tokenHighlightingCalls.length, 2);
       assert.deepEqual(tokenHighlightingCalls[1].details, undefined);
+      assert.isFalse(words1.classList.contains('token-highlight'));
     });
 
     test('triggers listener on token with single occurrence', async () => {
@@ -303,13 +307,13 @@
       dispatchMouseEvent('mouseover', words1);
       assert.equal(listener.pending, 0);
       clock.tick(HOVER_DELAY_MS);
-      assert.equal(listener.pending, 2);
+      assert.equal(listener.pending, 1);
       listener.flush();
       assert.equal(listener.pending, 0);
+      assert.isTrue(words1.classList.contains('token-highlight'));
       container.click();
-      assert.equal(listener.pending, 2);
-      assert.deepEqual(listener.shift(), [1, 1, Side.LEFT]);
-      assert.deepEqual(listener.shift(), [2, 2, Side.RIGHT]);
+      assert.equal(listener.pending, 1);
+      assert.isFalse(words1.classList.contains('token-highlight'));
     });
 
     test('clicking on word does not clear highlight', async () => {
@@ -323,11 +327,13 @@
       dispatchMouseEvent('mouseover', words1);
       assert.equal(listener.pending, 0);
       clock.tick(HOVER_DELAY_MS);
-      assert.equal(listener.pending, 2);
+      assert.equal(listener.pending, 1);
       listener.flush();
       assert.equal(listener.pending, 0);
+      assert.isTrue(words1.classList.contains('token-highlight'));
       words1.click();
       assert.equal(listener.pending, 0);
+      assert.isTrue(words1.classList.contains('token-highlight'));
     });
   });
 });