Migrate gr-diff-highlight from PolymerElement to plain class

Release-Notes: skip
Google-Bug-Id: n/a
Change-Id: I23e4f19fe630f5a0308c5ab01e75392301295cde
(cherry picked from commit f661777cbfe7f6292ba0bb795ff6b7535b38e21d)
diff --git a/polygerrit-ui/app/embed/diff/gr-diff-cursor/gr-diff-cursor_test.js b/polygerrit-ui/app/embed/diff/gr-diff-cursor/gr-diff-cursor_test.js
index 9b093af..c64f484 100644
--- a/polygerrit-ui/app/embed/diff/gr-diff-cursor/gr-diff-cursor_test.js
+++ b/polygerrit-ui/app/embed/diff/gr-diff-cursor/gr-diff-cursor_test.js
@@ -548,7 +548,7 @@
         end_line: 6,
         end_character: 1,
       };
-      diffElement.$.highlights.selectedRange = {
+      diffElement.highlights.selectedRange = {
         side: 'right',
         range: someRange,
       };
diff --git a/polygerrit-ui/app/embed/diff/gr-diff-highlight/gr-diff-highlight.ts b/polygerrit-ui/app/embed/diff/gr-diff-highlight/gr-diff-highlight.ts
index b64839b..0714645 100644
--- a/polygerrit-ui/app/embed/diff/gr-diff-highlight/gr-diff-highlight.ts
+++ b/polygerrit-ui/app/embed/diff/gr-diff-highlight/gr-diff-highlight.ts
@@ -1,28 +1,13 @@
 /**
  * @license
- * Copyright (C) 2016 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
+ * Copyright 2016 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
  */
 import '../../../styles/shared-styles';
 import '../gr-selection-action-box/gr-selection-action-box';
-import {dom, EventApi} from '@polymer/polymer/lib/legacy/polymer.dom';
-import {PolymerElement} from '@polymer/polymer/polymer-element';
-import {htmlTemplate} from './gr-diff-highlight_html';
 import {GrAnnotation} from './gr-annotation';
 import {normalize} from './gr-range-normalizer';
 import {strToClassName} from '../../../utils/dom-util';
-import {customElement, property} from '@polymer/decorators';
 import {Side} from '../../../constants/constants';
 import {CommentRange} from '../../../types/common';
 import {GrSelectionActionBox} from '../gr-selection-action-box/gr-selection-action-box';
@@ -34,7 +19,7 @@
   GrDiffThreadElement,
 } from '../gr-diff/gr-diff-utils';
 import {debounce, DelayedTask} from '../../../utils/async-util';
-import {queryAndAssert} from '../../../utils/common-util';
+import {assertIsDefined, queryAndAssert} from '../../../utils/common-util';
 
 interface SidedRange {
   side: Side;
@@ -61,45 +46,57 @@
   getContentTdByLineEl(lineEl?: Element): Element | null;
 }
 
-@customElement('gr-diff-highlight')
-export class GrDiffHighlight extends PolymerElement {
-  static get template() {
-    return htmlTemplate;
-  }
-
-  @property({type: Object})
-  _cachedDiffBuilder?: DiffBuilderInterface;
-
-  @property({type: Object, notify: true})
+/**
+ * Handles showing, positioning and interacting with <gr-selection-action-box>.
+ *
+ * Toggles a css class for highlighting comment ranges when the mouse leaves or
+ * enters a comment thread element.
+ */
+export class GrDiffHighlight {
   selectedRange?: SidedRange;
 
+  private diffBuilder?: DiffBuilderInterface;
+
+  private diffTable?: HTMLElement;
+
   private selectionChangeTask?: DelayedTask;
 
-  constructor() {
-    super();
-    this.addEventListener('comment-thread-mouseleave', e =>
-      this._handleCommentThreadMouseleave(e)
+  init(diffTable: HTMLElement, diffBuilder: DiffBuilderInterface) {
+    this.cleanup();
+
+    this.diffTable = diffTable;
+    this.diffBuilder = diffBuilder;
+
+    diffTable.addEventListener(
+      'comment-thread-mouseleave',
+      this.handleCommentThreadMouseleave
     );
-    this.addEventListener('comment-thread-mouseenter', e =>
-      this._handleCommentThreadMouseenter(e)
+    diffTable.addEventListener(
+      'comment-thread-mouseenter',
+      this.handleCommentThreadMouseenter
     );
-    this.addEventListener('create-comment-requested', e =>
-      this._handleRangeCommentRequest(e)
+    diffTable.addEventListener(
+      'create-comment-requested',
+      this.handleRangeCommentRequest
     );
   }
 
-  override disconnectedCallback() {
+  cleanup() {
     this.selectionChangeTask?.cancel();
-    super.disconnectedCallback();
-  }
-
-  get diffBuilder(): DiffBuilderInterface {
-    if (!this._cachedDiffBuilder) {
-      this._cachedDiffBuilder = this.querySelector(
-        'gr-diff-builder'
-      ) as DiffBuilderInterface;
+    if (this.diffTable) {
+      this.diffTable.removeEventListener(
+        'comment-thread-mouseleave',
+        this.handleCommentThreadMouseleave
+      );
+      this.diffTable.removeEventListener(
+        'comment-thread-mouseenter',
+        this.handleCommentThreadMouseenter
+      );
+      this.diffTable.removeEventListener(
+        'create-comment-requested',
+        this.handleRangeCommentRequest
+      );
     }
-    return this._cachedDiffBuilder;
   }
 
   /**
@@ -128,18 +125,17 @@
     // removed.
     // If you wait longer than 50 ms, then you don't properly catch a very
     // quick 'c' press after the selection change. If you wait less than 10
-    // ms, then you will have about 50 _handleSelection calls when doing a
+    // ms, then you will have about 50 handleSelection() calls when doing a
     // simple drag for select.
     this.selectionChangeTask = debounce(
       this.selectionChangeTask,
-      () => this._handleSelection(selection, isMouseUp),
+      () => this.handleSelection(selection, isMouseUp),
       10
     );
   }
 
-  _getThreadEl(e: Event): GrDiffThreadElement | null {
-    const path = (dom(e) as EventApi).path || [];
-    for (const pathEl of path) {
+  private getThreadEl(e: Event): GrDiffThreadElement | null {
+    for (const pathEl of e.composedPath()) {
       if (
         pathEl instanceof HTMLElement &&
         pathEl.classList.contains('comment-thread')
@@ -150,69 +146,53 @@
     return null;
   }
 
-  _toggleRangeElHighlight(
+  private toggleRangeElHighlight(
     threadEl: GrDiffThreadElement | null,
     highlightRange = false
   ) {
-    if (!threadEl) return;
-    // We don't want to re-create the line just for highlighting the range which
-    // is creating annoying bugs: @see Issue 12934
-    // As gr-ranged-comment-layer now does not notify the layer re-render and
-    // lack of access to the thread or the lineEl from the ranged-comment-layer,
-    // need to update range class for styles here.
-    let curNode: HTMLElement | null = threadEl.assignedSlot;
-    while (curNode) {
-      if (curNode.nodeName === 'TABLE') break;
-      curNode = curNode.parentElement;
-    }
-    if (curNode?.querySelectorAll) {
-      if (highlightRange) {
-        const rangeNodes = curNode.querySelectorAll(
-          `.range.${strToClassName(threadEl.rootId)}`
+    const rootId = threadEl?.rootId;
+    if (!rootId) return;
+    if (!this.diffTable) return;
+    if (highlightRange) {
+      const selector = `.range.${strToClassName(rootId)}`;
+      const rangeNodes = this.diffTable.querySelectorAll(selector);
+      rangeNodes.forEach(rangeNode => {
+        rangeNode.classList.add('rangeHoverHighlight');
+      });
+      const hintNode = this.diffTable.querySelector(
+        `gr-ranged-comment-hint[threadElRootId="${rootId}"]`
+      );
+      hintNode?.shadowRoot
+        ?.querySelectorAll('.rangeHighlight')
+        .forEach(highlightNode =>
+          highlightNode.classList.add('rangeHoverHighlight')
         );
-        rangeNodes.forEach(rangeNode => {
-          rangeNode.classList.add('rangeHoverHighlight');
-        });
-        const hintNode = threadEl.parentElement?.querySelector(
-          `gr-ranged-comment-hint[threadElRootId="${threadEl.rootId}"]`
+    } else {
+      const selector = `.rangeHoverHighlight.${strToClassName(rootId)}`;
+      const rangeNodes = this.diffTable.querySelectorAll(selector);
+      rangeNodes.forEach(rangeNode => {
+        rangeNode.classList.remove('rangeHoverHighlight');
+      });
+      const hintNode = this.diffTable.querySelector(
+        `gr-ranged-comment-hint[threadElRootId="${rootId}"]`
+      );
+      hintNode?.shadowRoot
+        ?.querySelectorAll('.rangeHoverHighlight')
+        .forEach(highlightNode =>
+          highlightNode.classList.remove('rangeHoverHighlight')
         );
-        if (hintNode) {
-          hintNode.shadowRoot
-            ?.querySelectorAll('.rangeHighlight')
-            .forEach(highlightNode =>
-              highlightNode.classList.add('rangeHoverHighlight')
-            );
-        }
-      } else {
-        const rangeNodes = curNode.querySelectorAll(
-          `.rangeHoverHighlight.${strToClassName(threadEl.rootId)}`
-        );
-        rangeNodes.forEach(rangeNode => {
-          rangeNode.classList.remove('rangeHoverHighlight');
-        });
-        const hintNode = threadEl.parentElement?.querySelector(
-          `gr-ranged-comment-hint[threadElRootId="${threadEl.rootId}"]`
-        );
-        if (hintNode) {
-          hintNode.shadowRoot
-            ?.querySelectorAll('.rangeHoverHighlight')
-            .forEach(highlightNode =>
-              highlightNode.classList.remove('rangeHoverHighlight')
-            );
-        }
-      }
     }
   }
 
-  _handleCommentThreadMouseenter(e: Event) {
-    const threadEl = this._getThreadEl(e);
-    this._toggleRangeElHighlight(threadEl, /* highlightRange= */ true);
-  }
+  private handleCommentThreadMouseenter = (e: Event) => {
+    const threadEl = this.getThreadEl(e);
+    this.toggleRangeElHighlight(threadEl, /* highlightRange= */ true);
+  };
 
-  _handleCommentThreadMouseleave(e: Event) {
-    const threadEl = this._getThreadEl(e);
-    this._toggleRangeElHighlight(threadEl, /* highlightRange= */ false);
-  }
+  private handleCommentThreadMouseleave = (e: Event) => {
+    const threadEl = this.getThreadEl(e);
+    this.toggleRangeElHighlight(threadEl, /* highlightRange= */ false);
+  };
 
   /**
    * Get current normalized selection.
@@ -220,20 +200,20 @@
    * syntax highligh, convert native DOM Range objects to Gerrit concepts
    * (line, side, etc).
    */
-  _getNormalizedRange(selection: Selection | Range) {
+  private getNormalizedRange(selection: Selection | Range) {
     /* On Safari the ShadowRoot.getSelection() isn't there and the only thing
        we can get is a single Range */
     if (selection instanceof Range) {
-      return this._normalizeRange(selection);
+      return this.normalizeRange(selection);
     }
     const rangeCount = selection.rangeCount;
     if (rangeCount === 0) {
       return null;
     } else if (rangeCount === 1) {
-      return this._normalizeRange(selection.getRangeAt(0));
+      return this.normalizeRange(selection.getRangeAt(0));
     } else {
-      const startRange = this._normalizeRange(selection.getRangeAt(0));
-      const endRange = this._normalizeRange(
+      const startRange = this.normalizeRange(selection.getRangeAt(0));
+      const endRange = this.normalizeRange(
         selection.getRangeAt(rangeCount - 1)
       );
       return {
@@ -248,15 +228,15 @@
    *
    * @return fixed normalized range
    */
-  _normalizeRange(domRange: Range): NormalizedRange {
+  private normalizeRange(domRange: Range): NormalizedRange {
     const range = normalize(domRange);
-    return this._fixTripleClickSelection(
+    return this.fixTripleClickSelection(
       {
-        start: this._normalizeSelectionSide(
+        start: this.normalizeSelectionSide(
           range.startContainer,
           range.startOffset
         ),
-        end: this._normalizeSelectionSide(range.endContainer, range.endOffset),
+        end: this.normalizeSelectionSide(range.endContainer, range.endOffset),
       },
       domRange
     );
@@ -272,7 +252,7 @@
    * @param domRange DOM Range object
    * @return fixed normalized range
    */
-  _fixTripleClickSelection(range: NormalizedRange, domRange: Range) {
+  private fixTripleClickSelection(range: NormalizedRange, domRange: Range) {
     if (!range.start) {
       // Selection outside of current diff.
       return range;
@@ -293,7 +273,7 @@
       end.column === 0 &&
       end.line === start.line + 1;
     const content = domRange.cloneContents().querySelector('.contentText');
-    const lineLength = (content && this._getLength(content)) || 0;
+    const lineLength = (content && this.getLength(content)) || 0;
     if (lineLength && (endsAtBeginningOfNextLine || endsAtOtherEmptySide)) {
       // Move the selection to the end of the previous line.
       range.end = {
@@ -314,12 +294,14 @@
    * @param node td.content child
    * @param offset offset within node
    */
-  _normalizeSelectionSide(
+  private normalizeSelectionSide(
     node: Node | null,
     offset: number
   ): NormalizedPosition | null {
     let column;
-    if (!node || !this.contains(node)) return null;
+    if (!this.diffTable) return null;
+    if (!this.diffBuilder) return null;
+    if (!node || !this.diffTable.contains(node)) return null;
     const lineEl = getLineElByChild(node);
     if (!lineEl) return null;
     const side = getSideByLineEl(lineEl);
@@ -335,10 +317,10 @@
     } else {
       const thread = contentTd.querySelector('.comment-thread');
       if (thread?.contains(node)) {
-        column = this._getLength(contentText);
+        column = this.getLength(contentText);
         node = contentText;
       } else {
-        column = this._convertOffsetToColumn(node, offset);
+        column = this.convertOffsetToColumn(node, offset);
       }
     }
 
@@ -357,7 +339,8 @@
    * collapsed section, so don't need to worry about this case for
    * positioning the tooltip.
    */
-  _positionActionBox(
+  // visible for testing
+  positionActionBox(
     actionBox: GrSelectionActionBox,
     startLine: number,
     range: Text | Element | Range
@@ -371,7 +354,7 @@
     actionBox.placeBelow(range);
   }
 
-  _isRangeValid(range: NormalizedRange | null) {
+  private isRangeValid(range: NormalizedRange | null) {
     if (!range || !range.start || !range.start.node || !range.end) {
       return false;
     }
@@ -384,15 +367,16 @@
     );
   }
 
-  _handleSelection(selection: Selection | Range, isMouseUp: boolean) {
+  // visible for testing
+  handleSelection(selection: Selection | Range, isMouseUp: boolean) {
     /* On Safari, the selection events may return a null range that should
        be ignored */
-    if (!selection) {
-      return;
-    }
-    const normalizedRange = this._getNormalizedRange(selection);
-    if (!this._isRangeValid(normalizedRange)) {
-      this._removeActionBox();
+    if (!selection) return;
+    if (!this.diffTable) return;
+
+    const normalizedRange = this.getNormalizedRange(selection);
+    if (!this.isRangeValid(normalizedRange)) {
+      this.removeActionBox();
       return;
     }
     /* On Safari the ShadowRoot.getSelection() isn't there and the only thing
@@ -422,8 +406,8 @@
       // start.column with the content length), we just check if the selection
       // is empty to see that it's at the end of a line.
       const content = domRange.cloneContents().querySelector('.contentText');
-      if (isMouseUp && this._getLength(content) === 0) {
-        this._fireCreateRangeComment(start.side, {
+      if (isMouseUp && this.getLength(content) === 0) {
+        this.fireCreateRangeComment(start.side, {
           start_line: start.line,
           start_character: 0,
           end_line: start.line,
@@ -433,10 +417,10 @@
       return;
     }
 
-    let actionBox = this.shadowRoot!.querySelector('gr-selection-action-box');
+    let actionBox = this.diffTable.querySelector('gr-selection-action-box');
     if (!actionBox) {
       actionBox = document.createElement('gr-selection-action-box');
-      this.root!.insertBefore(actionBox, this.root!.firstElementChild);
+      this.diffTable.appendChild(actionBox);
     }
     this.selectedRange = {
       range: {
@@ -448,10 +432,10 @@
       side: start.side,
     };
     if (start.line === end.line) {
-      this._positionActionBox(actionBox, start.line, domRange);
+      this.positionActionBox(actionBox, start.line, domRange);
     } else if (start.node instanceof Text) {
       if (start.column) {
-        this._positionActionBox(
+        this.positionActionBox(
           actionBox,
           start.line,
           start.node.splitText(start.column)
@@ -464,44 +448,41 @@
       (start.node.firstChild instanceof Element ||
         start.node.firstChild instanceof Text)
     ) {
-      this._positionActionBox(actionBox, start.line, start.node.firstChild);
+      this.positionActionBox(actionBox, start.line, start.node.firstChild);
     } else if (start.node instanceof Element || start.node instanceof Text) {
-      this._positionActionBox(actionBox, start.line, start.node);
+      this.positionActionBox(actionBox, start.line, start.node);
     } else {
       console.warn('Failed to position comment action box.');
-      this._removeActionBox();
+      this.removeActionBox();
     }
   }
 
-  _fireCreateRangeComment(side: Side, range: CommentRange) {
-    this.dispatchEvent(
+  private fireCreateRangeComment(side: Side, range: CommentRange) {
+    this.diffTable?.dispatchEvent(
       new CustomEvent('create-range-comment', {
         detail: {side, range},
         composed: true,
         bubbles: true,
       })
     );
-    this._removeActionBox();
+    this.removeActionBox();
   }
 
-  _handleRangeCommentRequest(e: Event) {
+  private handleRangeCommentRequest = (e: Event) => {
     e.stopPropagation();
-    if (!this.selectedRange) {
-      throw Error('Selected Range is needed for new range comment!');
-    }
+    assertIsDefined(this.selectedRange, 'selectedRange');
     const {side, range} = this.selectedRange;
-    this._fireCreateRangeComment(side, range);
-  }
+    this.fireCreateRangeComment(side, range);
+  };
 
-  _removeActionBox() {
+  // visible for testing
+  removeActionBox() {
     this.selectedRange = undefined;
-    const actionBox = this.shadowRoot!.querySelector('gr-selection-action-box');
-    if (actionBox) {
-      this.root!.removeChild(actionBox);
-    }
+    const actionBox = this.diffTable?.querySelector('gr-selection-action-box');
+    if (actionBox) actionBox.remove();
   }
 
-  _convertOffsetToColumn(el: Node, offset: number) {
+  private convertOffsetToColumn(el: Node, offset: number) {
     if (el instanceof Element && el.classList.contains('content')) {
       return offset;
     }
@@ -511,7 +492,7 @@
     ) {
       if (el.previousSibling) {
         el = el.previousSibling;
-        offset += this._getLength(el);
+        offset += this.getLength(el);
       } else {
         el = el.parentElement!;
       }
@@ -525,10 +506,11 @@
    *
    * @param node this is sometimes passed as null.
    */
-  _getLength(node: Node | null): number {
+  // visible for testing
+  getLength(node: Node | null): number {
     if (node === null) return 0;
     if (node instanceof Element && node.classList.contains('content')) {
-      return this._getLength(queryAndAssert(node, '.contentText'));
+      return this.getLength(queryAndAssert(node, '.contentText'));
     } else {
       return GrAnnotation.getLength(node);
     }
@@ -541,9 +523,6 @@
 }
 
 declare global {
-  interface HTMLElementTagNameMap {
-    'gr-diff-highlight': GrDiffHighlight;
-  }
   interface HTMLElementEventMap {
     'create-range-comment': CustomEvent<CreateRangeCommentEventDetail>;
   }
diff --git a/polygerrit-ui/app/embed/diff/gr-diff-highlight/gr-diff-highlight_html.ts b/polygerrit-ui/app/embed/diff/gr-diff-highlight/gr-diff-highlight_html.ts
deleted file mode 100644
index 5a6cb1c..0000000
--- a/polygerrit-ui/app/embed/diff/gr-diff-highlight/gr-diff-highlight_html.ts
+++ /dev/null
@@ -1,35 +0,0 @@
-/**
- * @license
- * Copyright (C) 2020 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-import {html} from '@polymer/polymer/lib/utils/html-tag';
-
-export const htmlTemplate = html`
-  <style include="shared-styles">
-    :host {
-      position: relative;
-    }
-    gr-selection-action-box {
-      /**
-         * Needs z-index to appear above wrapped content, since it's inserted
-         * into DOM before it.
-         */
-      z-index: 10;
-    }
-  </style>
-  <div class="contentWrapper">
-    <slot></slot>
-  </div>
-`;
diff --git a/polygerrit-ui/app/embed/diff/gr-diff-highlight/gr-diff-highlight_test.ts b/polygerrit-ui/app/embed/diff/gr-diff-highlight/gr-diff-highlight_test.ts
index 1adfa4a..b819754 100644
--- a/polygerrit-ui/app/embed/diff/gr-diff-highlight/gr-diff-highlight_test.ts
+++ b/polygerrit-ui/app/embed/diff/gr-diff-highlight/gr-diff-highlight_test.ts
@@ -6,7 +6,7 @@
 import '../../../test/common-test-setup-karma';
 import './gr-diff-highlight';
 import {_getTextOffset} from './gr-range-normalizer';
-import {html} from '@polymer/polymer/lib/utils/html-tag';
+import {fixture, fixtureCleanup, html} from '@open-wc/testing-helpers';
 import {
   GrDiffHighlight,
   DiffBuilderInterface,
@@ -16,155 +16,147 @@
 import {SinonStubbedMember} from 'sinon';
 import {queryAndAssert} from '../../../utils/common-util';
 import {GrDiffThreadElement} from '../gr-diff/gr-diff-utils';
+import {waitQueryAndAssert, waitUntil} from '../../../test/test-utils';
+import {GrSelectionActionBox} from '../gr-selection-action-box/gr-selection-action-box';
 
 // Splitting long lines in html into shorter rows breaks tests:
 // zero-length text nodes and new lines are not expected in some places
 /* eslint-disable max-len, lit/prefer-static-styles */
 /* prettier-ignore */
-const basicFixture = fixtureFromTemplate(html`
-  <style>
-    .tab-indicator:before {
-      color: #c62828;
-      /* >> character */
-      content: '\\00BB';
-    }
-  </style>
-  <gr-diff-highlight>
-    <table id="diffTable">
-      <tbody class="section both">
-        <tr class="diff-row side-by-side" left-type="both" right-type="both">
-          <td class="left lineNum" data-value="1"></td>
-          <td class="content both"><div class="contentText">[1] Nam cum ad me in Cumanum salutandi causa uterque venisset,</div></td>
-          <td class="right lineNum" data-value="1"></td>
-          <td class="content both"><div class="contentText">[1] Nam cum ad me in Cumanum salutandi causa uterque</div></td>
-        </tr>
-      </tbody>
+const diffTable = html`
+  <table id="diffTable">
+    <tbody class="section both">
+      <tr class="diff-row side-by-side" left-type="both" right-type="both">
+        <td class="left lineNum" data-value="1"></td>
+        <td class="content both"><div class="contentText">[1] Nam cum ad me in Cumanum salutandi causa uterque venisset,</div></td>
+        <td class="right lineNum" data-value="1"></td>
+        <td class="content both"><div class="contentText">[1] Nam cum ad me in Cumanum salutandi causa uterque</div></td>
+      </tr>
+    </tbody>
 
-      <tbody class="section delta">
-        <tr class="diff-row side-by-side" left-type="remove" right-type="add">
-          <td class="left lineNum" data-value="2"></td>
-          <!-- Next tag is formatted to eliminate zero-length text nodes. -->
-          <td class="content remove"><div class="contentText">na💢ti <hl class="foo range id314">te, inquit</hl>, sumus<hl class="bar">aliquando</hl> otiosum, <hl>certe</hl> a<hl><span class="tab-indicator" style="tab-size:8;"> </span></hl>udiam, <hl>quid</hl> sit,<span class="tab-indicator" style="tab-size:8;"> </span>quod<hl>Epicurum</hl></div></td>
-          <td class="right lineNum" data-value="2"></td>
-          <!-- Next tag is formatted to eliminate zero-length text nodes. -->
-          <td class="content add"><div class="contentText">nacti , <hl>,</hl> sumus<hl><span class="tab-indicator" style="tab-size:8;"> </span></hl>otiosum,<span class="tab-indicator" style="tab-size:8;"> </span> audiam,sit, quod</div></td>
-        </tr>
-      </tbody>
+    <tbody class="section delta">
+      <tr class="diff-row side-by-side" left-type="remove" right-type="add">
+        <td class="left lineNum" data-value="2"></td>
+        <!-- Next tag is formatted to eliminate zero-length text nodes. -->
+        <td class="content remove"><div class="contentText">na💢ti <hl class="foo range generated_id314">te, inquit</hl>, sumus<hl class="bar">aliquando</hl> otiosum, <hl>certe</hl> a<hl><span class="tab-indicator" style="tab-size:8;"> </span></hl>udiam, <hl>quid</hl> sit,<span class="tab-indicator" style="tab-size:8;"> </span>quod<hl>Epicurum</hl></div></td>
+        <td class="right lineNum" data-value="2"></td>
+        <!-- Next tag is formatted to eliminate zero-length text nodes. -->
+        <td class="content add"><div class="contentText">nacti , <hl>,</hl> sumus<hl><span class="tab-indicator" style="tab-size:8;"> </span></hl>otiosum,<span class="tab-indicator" style="tab-size:8;"> </span> audiam,sit, quod</div></td>
+      </tr>
+    </tbody>
 
-      <tbody class="section both">
-        <tr class="diff-row side-by-side" left-type="both" right-type="both">
-          <td class="left lineNum" data-value="138"></td>
-          <td class="content both"><div class="contentText">[14] Nam cum ad me in Cumanum salutandi causa uterque venisset,</div></td>
-          <td class="right lineNum" data-value="119"></td>
-          <td class="content both"><div class="contentText">[14] Nam cum ad me in Cumanum salutandi causa uterque venisset,</div></td>
-        </tr>
-      </tbody>
+    <tbody class="section both">
+      <tr class="diff-row side-by-side" left-type="both" right-type="both">
+        <td class="left lineNum" data-value="138"></td>
+        <td class="content both"><div class="contentText">[14] Nam cum ad me in Cumanum salutandi causa uterque venisset,</div></td>
+        <td class="right lineNum" data-value="119"></td>
+        <td class="content both"><div class="contentText">[14] Nam cum ad me in Cumanum salutandi causa uterque venisset,</div></td>
+      </tr>
+    </tbody>
 
-      <tbody class="section delta">
-        <tr class="diff-row side-by-side" left-type="remove" right-type="add">
-          <td class="left lineNum" data-value="140"></td>
-          <!-- Next tag is formatted to eliminate zero-length text nodes. -->
-          <td class="content remove"><div class="contentText">na💢ti <hl class="foo">te, inquit</hl>, sumus <hl class="bar">aliquando</hl> otiosum, <hl>certe</hl> a <hl><span class="tab-indicator" style="tab-size:8;">\u0009</span></hl>udiam, <hl>quid</hl> sit, <span class="tab-indicator" style="tab-size:8;">\u0009</span>quod <hl>Epicurum</hl></div><div class="comment-thread">
-            [Yet another random diff thread content here]
-          </div></td>
-          <td class="right lineNum" data-value="120"></td>
-          <!-- Next tag is formatted to eliminate zero-length text nodes. -->
-          <td class="content add"><div class="contentText">nacti , <hl>,</hl> sumus <hl><span class="tab-indicator" style="tab-size:8;">\u0009</span></hl> otiosum,  <span class="tab-indicator" style="tab-size:8;">\u0009</span> audiam,  sit, quod</div></td>
-        </tr>
-      </tbody>
+    <tbody class="section delta">
+      <tr class="diff-row side-by-side" left-type="remove" right-type="add">
+        <td class="left lineNum" data-value="140"></td>
+        <!-- Next tag is formatted to eliminate zero-length text nodes. -->
+        <td class="content remove"><div class="contentText">na💢ti <hl class="foo">te, inquit</hl>, sumus <hl class="bar">aliquando</hl> otiosum, <hl>certe</hl> a <hl><span class="tab-indicator" style="tab-size:8;">\u0009</span></hl>udiam, <hl>quid</hl> sit, <span class="tab-indicator" style="tab-size:8;">\u0009</span>quod <hl>Epicurum</hl></div><div class="comment-thread">
+          [Yet another random diff thread content here]
+        </div></td>
+        <td class="right lineNum" data-value="120"></td>
+        <!-- Next tag is formatted to eliminate zero-length text nodes. -->
+        <td class="content add"><div class="contentText">nacti , <hl>,</hl> sumus <hl><span class="tab-indicator" style="tab-size:8;">\u0009</span></hl> otiosum,  <span class="tab-indicator" style="tab-size:8;">\u0009</span> audiam,  sit, quod</div></td>
+      </tr>
+    </tbody>
 
-      <tbody class="section both">
-        <tr class="diff-row side-by-side" left-type="both" right-type="both">
-          <td class="left lineNum" data-value="141"></td>
-          <td class="content both"><div class="contentText">nam et<hl><span class="tab-indicator" style="tab-size:8;"> </span></hl>complectitur<span class="tab-indicator" style="tab-size:8;"></span>verbis, quod vult, et dicit plane, quod intellegam;</div></td>
-          <td class="right lineNum" data-value="130"></td>
-          <td class="content both"><div class="contentText">nam et complectitur verbis, quod vult, et dicit plane, quodintellegam;</div></td>
-        </tr>
-      </tbody>
+    <tbody class="section both">
+      <tr class="diff-row side-by-side" left-type="both" right-type="both">
+        <td class="left lineNum" data-value="141"></td>
+        <td class="content both"><div class="contentText">nam et<hl><span class="tab-indicator" style="tab-size:8;"> </span></hl>complectitur<span class="tab-indicator" style="tab-size:8;"></span>verbis, quod vult, et dicit plane, quod intellegam;</div></td>
+        <td class="right lineNum" data-value="130"></td>
+        <td class="content both"><div class="contentText">nam et complectitur verbis, quod vult, et dicit plane, quodintellegam;</div></td>
+      </tr>
+    </tbody>
 
-      <tbody class="section contextControl">
-        <tr
-          class="diff-row side-by-side"
-          left-type="contextControl"
-          right-type="contextControl"
-        >
-          <td class="left contextLineNum"></td>
-          <td>
-            <gr-button>+10↑</gr-button>
-            -
-            <gr-button>Show 21 common lines</gr-button>
-            -
-            <gr-button>+10↓</gr-button>
-          </td>
-          <td class="right contextLineNum"></td>
-          <td>
-            <gr-button>+10↑</gr-button>
-            -
-            <gr-button>Show 21 common lines</gr-button>
-            -
-            <gr-button>+10↓</gr-button>
-          </td>
-        </tr>
-      </tbody>
+    <tbody class="section contextControl">
+      <tr
+        class="diff-row side-by-side"
+        left-type="contextControl"
+        right-type="contextControl"
+      >
+        <td class="left contextLineNum"></td>
+        <td>
+          <gr-button>+10↑</gr-button>
+          -
+          <gr-button>Show 21 common lines</gr-button>
+          -
+          <gr-button>+10↓</gr-button>
+        </td>
+        <td class="right contextLineNum"></td>
+        <td>
+          <gr-button>+10↑</gr-button>
+          -
+          <gr-button>Show 21 common lines</gr-button>
+          -
+          <gr-button>+10↓</gr-button>
+        </td>
+      </tr>
+    </tbody>
 
-      <tbody class="section delta total">
-        <tr class="diff-row side-by-side" left-type="blank" right-type="add">
-          <td class="left"></td>
-          <td class="blank"></td>
-          <td class="right lineNum" data-value="146"></td>
-          <td class="content add"><div class="contentText">[17] Quid igitur est? inquit; audire enim cupio, quid non probes. Principio, inquam,</div></td>
-        </tr>
-      </tbody>
+    <tbody class="section delta total">
+      <tr class="diff-row side-by-side" left-type="blank" right-type="add">
+        <td class="left"></td>
+        <td class="blank"></td>
+        <td class="right lineNum" data-value="146"></td>
+        <td class="content add"><div class="contentText">[17] Quid igitur est? inquit; audire enim cupio, quid non probes. Principio, inquam,</div></td>
+      </tr>
+    </tbody>
 
-      <tbody class="section both">
-        <tr class="diff-row side-by-side" left-type="both" right-type="both">
-          <td class="left lineNum" data-value="165"></td>
-          <td class="content both"><div class="contentText"></div></td>
-          <td class="right lineNum" data-value="147"></td>
-          <td class="content both"><div class="contentText">in physicis, <hl><span class="tab-indicator" style="tab-size:8;"> </span></hl>quibus maxime gloriatur, primum totus est alienus. Democritea dicit</div></td>
-        </tr>
-      </tbody>
-    </table>
-  </gr-diff-highlight>
-`);
+    <tbody class="section both">
+      <tr class="diff-row side-by-side" left-type="both" right-type="both">
+        <td class="left lineNum" data-value="165"></td>
+        <td class="content both"><div class="contentText"></div></td>
+        <td class="right lineNum" data-value="147"></td>
+        <td class="content both"><div class="contentText">in physicis, <hl><span class="tab-indicator" style="tab-size:8;"> </span></hl>quibus maxime gloriatur, primum totus est alienus. Democritea dicit</div></td>
+      </tr>
+    </tbody>
+  </table>
+`;
 /* eslint-enable max-len */
 
 suite('gr-diff-highlight', () => {
-  let element: GrDiffHighlight;
-
-  setup(() => {
-    element = (
-      basicFixture.instantiate() as HTMLElement[]
-    )[1] as GrDiffHighlight;
-  });
-
   suite('comment events', () => {
-    let builder;
     let threadEl: GrDiffThreadElement;
     let hlRange: HTMLElement;
+    let element: GrDiffHighlight;
+    let diff: HTMLElement;
+    let builder: {
+      getContentTdByLineEl: SinonStubbedMember<
+        DiffBuilderInterface['getContentTdByLineEl']
+      >;
+    };
 
-    setup(() => {
+    setup(async () => {
+      diff = await fixture<HTMLTableElement>(diffTable);
       builder = {
-        getContentTdByLineEl: () => null,
+        getContentTdByLineEl: sinon.stub(),
       };
-      element._cachedDiffBuilder = builder;
-
-      const diff = queryAndAssert(element, '#diffTable');
-      hlRange = queryAndAssert(diff, 'hl.range.id314');
+      element = new GrDiffHighlight();
+      element.init(diff, builder);
+      hlRange = queryAndAssert(diff, 'hl.range.generated_id314');
 
       threadEl = document.createElement(
         'div'
       ) as unknown as GrDiffThreadElement;
       threadEl.className = 'comment-thread';
       threadEl.rootId = 'id314';
-      document.body.appendChild(threadEl);
+      diff.appendChild(threadEl);
     });
 
     teardown(() => {
+      element.cleanup();
       threadEl.remove();
     });
 
-    test('comment-thread-mouseenter toggles rangeHoverHighlight class', () => {
+    test('comment-thread-mouseenter toggles rangeHoverHighlight class', async () => {
       assert.isFalse(hlRange.classList.contains('rangeHoverHighlight'));
       threadEl.dispatchEvent(
         new CustomEvent('comment-thread-mouseenter', {
@@ -172,12 +164,11 @@
           composed: true,
         })
       );
-      // TODO: Enable this test when we don't need `assignedSlot` anymore.
-      // This will be done in a direct follow-up change.
-      // assert.isTrue(hlRange.classList.contains('rangeHoverHighlight'));
+      await waitUntil(() => hlRange.classList.contains('rangeHoverHighlight'));
+      assert.isTrue(hlRange.classList.contains('rangeHoverHighlight'));
     });
 
-    test('comment-thread-mouseleave toggles rangeHoverHighlight class', () => {
+    test('comment-thread-mouseleave toggles rangeHoverHighlight class', async () => {
       hlRange.classList.add('rangeHoverHighlight');
       threadEl.dispatchEvent(
         new CustomEvent('comment-thread-mouseleave', {
@@ -185,14 +176,13 @@
           composed: true,
         })
       );
-      // TODO: Enable this test when we don't need `assignedSlot` anymore.
-      // This will be done in a direct follow-up change.
-      // assert.isFalse(hlRange.classList.contains('rangeHoverHighlight'));
+      await waitUntil(() => !hlRange.classList.contains('rangeHoverHighlight'));
+      assert.isFalse(hlRange.classList.contains('rangeHoverHighlight'));
     });
 
     test(`create-range-comment for range when create-comment-requested
           is fired`, () => {
-      const removeActionBoxStub = sinon.stub(element, '_removeActionBox');
+      const removeActionBoxStub = sinon.stub(element, 'removeActionBox');
       element.selectedRange = {
         side: Side.LEFT,
         range: {
@@ -204,10 +194,10 @@
       };
       const requestEvent = new CustomEvent('create-comment-requested');
       let createRangeEvent: CustomEvent<CreateRangeCommentEventDetail>;
-      element.addEventListener('create-range-comment', e => {
+      diff.addEventListener('create-range-comment', e => {
         createRangeEvent = e;
       });
-      element.dispatchEvent(requestEvent);
+      diff.dispatchEvent(requestEvent);
       if (!createRangeEvent!) assert.fail('event not set');
       assert.deepEqual(element.selectedRange, createRangeEvent.detail);
       assert.isTrue(removeActionBoxStub.called);
@@ -215,6 +205,7 @@
   });
 
   suite('selection', () => {
+    let element: GrDiffHighlight;
     let diff: HTMLElement;
     let builder: {
       getContentTdByLineEl: SinonStubbedMember<
@@ -223,6 +214,25 @@
     };
     let contentStubs;
 
+    setup(async () => {
+      diff = await fixture<HTMLTableElement>(diffTable);
+      builder = {
+        getContentTdByLineEl: sinon.stub(),
+      };
+      element = new GrDiffHighlight();
+      element.init(diff, builder);
+      contentStubs = [];
+      stub('gr-selection-action-box', 'placeAbove');
+      stub('gr-selection-action-box', 'placeBelow');
+    });
+
+    teardown(() => {
+      fixtureCleanup();
+      element.cleanup();
+      contentStubs = null;
+      document.getSelection()!.removeAllRanges();
+    });
+
     const stubContent = (line: number, side: Side) => {
       const contentTd = diff.querySelector(
         `.${side}.lineNum[data-value="${line}"] ~ .content`
@@ -247,38 +257,22 @@
       endNode: Node,
       endOffset: number
     ) => {
-      const selection = document.getSelection()!;
+      const selection = document.getSelection();
+      if (!selection) assert.fail('no selection');
+      selection.removeAllRanges();
       const range = document.createRange();
       range.setStart(startNode, startOffset);
       range.setEnd(endNode, endOffset);
       selection.addRange(range);
-      element._handleSelection(selection, false);
+      element.handleSelection(selection, false);
     };
 
-    setup(() => {
-      contentStubs = [];
-      stub('gr-selection-action-box', 'placeAbove');
-      stub('gr-selection-action-box', 'placeBelow');
-      diff = queryAndAssert(element, '#diffTable');
-      builder = {
-        getContentTdByLineEl: sinon.stub(),
-      };
-      element._cachedDiffBuilder = builder;
-    });
-
-    teardown(() => {
-      contentStubs = null;
-      document.getSelection()!.removeAllRanges();
-    });
-
     test('single first line', () => {
       const content = stubContent(1, Side.RIGHT);
-      sinon.spy(element, '_positionActionBox');
+      sinon.spy(element, 'positionActionBox');
       if (!content?.firstChild) assert.fail('content first child not found');
       emulateSelection(content.firstChild, 5, content.firstChild, 12);
-      const actionBox = element.shadowRoot?.querySelector(
-        'gr-selection-action-box'
-      );
+      const actionBox = diff.querySelector('gr-selection-action-box');
       if (!actionBox) assert.fail('action box not found');
       assert.isTrue(actionBox.positionBelow);
     });
@@ -286,7 +280,7 @@
     test('multiline starting on first line', () => {
       const startContent = stubContent(1, Side.RIGHT);
       const endContent = stubContent(2, Side.RIGHT);
-      sinon.spy(element, '_positionActionBox');
+      sinon.spy(element, 'positionActionBox');
       if (!startContent?.firstChild) {
         assert.fail('first child of start content not found');
       }
@@ -294,22 +288,20 @@
         assert.fail('last child of end content not found');
       }
       emulateSelection(startContent.firstChild, 10, endContent.lastChild, 7);
-      const actionBox = element.shadowRoot?.querySelector(
-        'gr-selection-action-box'
-      );
+      const actionBox = diff.querySelector('gr-selection-action-box');
       if (!actionBox) assert.fail('action box not found');
       assert.isTrue(actionBox.positionBelow);
     });
 
-    test('single line', () => {
+    test('single line', async () => {
       const content = stubContent(138, Side.LEFT);
-      sinon.spy(element, '_positionActionBox');
+      sinon.spy(element, 'positionActionBox');
       if (!content?.firstChild) assert.fail('content first child not found');
       emulateSelection(content.firstChild, 5, content.firstChild, 12);
-      const actionBox = element.shadowRoot?.querySelector(
+      const actionBox = await waitQueryAndAssert<GrSelectionActionBox>(
+        diff,
         'gr-selection-action-box'
       );
-      if (!actionBox) assert.fail('action box not found');
       if (!element.selectedRange) assert.fail('no range selected');
       const {range, side} = element.selectedRange;
       assert.deepEqual(range, {
@@ -325,7 +317,7 @@
     test('multiline', () => {
       const startContent = stubContent(119, Side.RIGHT);
       const endContent = stubContent(120, Side.RIGHT);
-      sinon.spy(element, '_positionActionBox');
+      sinon.spy(element, 'positionActionBox');
       if (!startContent?.firstChild) {
         assert.fail('first child of start content not found');
       }
@@ -333,9 +325,7 @@
         assert.fail('last child of end content');
       }
       emulateSelection(startContent.firstChild, 10, endContent.lastChild, 7);
-      const actionBox = element.shadowRoot?.querySelector(
-        'gr-selection-action-box'
-      );
+      const actionBox = diff.querySelector('gr-selection-action-box');
       if (!actionBox) assert.fail('action box not found');
       if (!element.selectedRange) assert.fail('no range selected');
       const {range, side} = element.selectedRange;
@@ -378,7 +368,7 @@
         getRangeAt: getRangeAtStub,
         removeAllRanges: sinon.stub(),
       } as unknown as Selection;
-      element._handleSelection(selection, false);
+      element.handleSelection(selection, false);
       if (!element.selectedRange) assert.fail('no range selected');
       const {range} = element.selectedRange;
       assert.deepEqual(range, {
@@ -684,7 +674,7 @@
       assert.equal(side, Side.LEFT);
     });
 
-    test('GrRangeNormalizer._getTextOffset computes text offset', () => {
+    test('GrRangeNormalizer.getTextOffset computes text offset', () => {
       let content = stubContent(140, Side.LEFT);
       if (!content) assert.fail('content element not found');
       if (!content.lastChild) assert.fail('last child of content not found');
@@ -700,7 +690,7 @@
       assert.equal(result, 0);
     });
 
-    test('_fixTripleClickSelection', () => {
+    test('fixTripleClickSelection', () => {
       const startContent = stubContent(119, Side.RIGHT);
       const endContent = stubContent(120, Side.RIGHT);
       if (!startContent?.firstChild) {
@@ -715,7 +705,7 @@
         start_line: 119,
         start_character: 0,
         end_line: 119,
-        end_character: element._getLength(startContent),
+        end_character: element.getLength(startContent),
       });
       assert.equal(side, Side.RIGHT);
     });
diff --git a/polygerrit-ui/app/embed/diff/gr-diff/gr-diff.ts b/polygerrit-ui/app/embed/diff/gr-diff/gr-diff.ts
index 22da871..34c2a33 100644
--- a/polygerrit-ui/app/embed/diff/gr-diff/gr-diff.ts
+++ b/polygerrit-ui/app/embed/diff/gr-diff/gr-diff.ts
@@ -100,7 +100,6 @@
 
 export interface GrDiff {
   $: {
-    highlights: GrDiffHighlight;
     diffBuilder: GrDiffBuilderElement;
     diffTable: HTMLTableElement;
   };
@@ -297,6 +296,8 @@
 
   private diffSelection = new GrDiffSelection();
 
+  private highlights = new GrDiffHighlight();
+
   constructor() {
     super();
     this._setLoading(true);
@@ -319,6 +320,7 @@
     this._unobserveIncrementalNodes();
     this._unobserveNodes();
     this.diffSelection.cleanup();
+    this.highlights.cleanup();
     super.disconnectedCallback();
   }
 
@@ -361,7 +363,7 @@
     // and pass the shadow DOM selection into gr-diff-highlight, where the
     // corresponding range is determined and normalized.
     const selection = this._getShadowOrDocumentSelection();
-    this.$.highlights.handleSelectionChange(selection, false);
+    this.highlights.handleSelectionChange(selection, false);
   };
 
   private readonly handleMouseUp = () => {
@@ -369,7 +371,7 @@
     // mouse-up if there's a selection that just covers a line change. We
     // can't do that on selection change since the user may still be dragging.
     const selection = this._getShadowOrDocumentSelection();
-    this.$.highlights.handleSelectionChange(selection, true);
+    this.highlights.handleSelectionChange(selection, true);
   };
 
   /** Gets the current selection, preferring the shadow DOM selection. */
@@ -501,7 +503,7 @@
   }
 
   isRangeSelected() {
-    return !!this.$.highlights.selectedRange;
+    return !!this.highlights.selectedRange;
   }
 
   toggleLeftDiff() {
@@ -593,7 +595,7 @@
     if (!this.isRangeSelected()) {
       throw Error('Selection is needed for new range comment');
     }
-    const selectedRange = this.$.highlights.selectedRange;
+    const selectedRange = this.highlights.selectedRange;
     if (!selectedRange) throw Error('selected range not set');
     const {side, range} = selectedRange;
     this._createCommentForSelection(side, range);
@@ -818,6 +820,7 @@
     }
     if (this.diff) {
       this.diffSelection.init(this.diff, this.$.diffTable);
+      this.highlights.init(this.$.diffTable, this.$.diffBuilder);
     }
   }
 
diff --git a/polygerrit-ui/app/embed/diff/gr-diff/gr-diff_html.ts b/polygerrit-ui/app/embed/diff/gr-diff/gr-diff_html.ts
index 5d9e5c3..6d36b89 100644
--- a/polygerrit-ui/app/embed/diff/gr-diff/gr-diff_html.ts
+++ b/polygerrit-ui/app/embed/diff/gr-diff/gr-diff_html.ts
@@ -487,6 +487,10 @@
       color: var(--link-color);
       padding: var(--spacing-m) 0 var(--spacing-m) 48px;
     }
+    #diffTable {
+      /* for gr-selection-action-box positioning */
+      position: relative;
+    }
     #diffTable:focus {
       outline: none;
     }
@@ -670,6 +674,14 @@
     .token-highlight {
       background-color: var(--token-highlighting-color, #fffd54);
     }
+
+    gr-selection-action-box {
+      /**
+       * Needs z-index to appear above wrapped content, since it's inserted
+       * into DOM before it.
+       */
+      z-index: 10;
+    }
   </style>
   <style include="gr-syntax-theme">
     /* Workaround for empty style block - see https://github.com/Polymer/tools/issues/408 */
@@ -686,38 +698,36 @@
     class$="[[_computeContainerClass(loggedIn, viewMode, displayLine)]]"
     on-click="_handleTap"
   >
-    <gr-diff-highlight id="highlights">
-      <gr-diff-builder
-        id="diffBuilder"
-        comment-ranges="[[_commentRanges]]"
-        coverage-ranges="[[coverageRanges]]"
-        diff="[[diff]]"
-        path="[[path]]"
-        view-mode="[[viewMode]]"
-        is-image-diff="[[isImageDiff]]"
-        base-image="[[baseImage]]"
-        layers="[[layers]]"
-        revision-image="[[revisionImage]]"
-        use-new-image-diff-ui="[[useNewImageDiffUi]]"
-      >
-        <table
-          id="diffTable"
-          class$="[[_diffTableClass]]"
-          role="presentation"
-          contenteditable$="[[isContentEditable]]"
-        ></table>
+    <gr-diff-builder
+      id="diffBuilder"
+      comment-ranges="[[_commentRanges]]"
+      coverage-ranges="[[coverageRanges]]"
+      diff="[[diff]]"
+      path="[[path]]"
+      view-mode="[[viewMode]]"
+      is-image-diff="[[isImageDiff]]"
+      base-image="[[baseImage]]"
+      layers="[[layers]]"
+      revision-image="[[revisionImage]]"
+      use-new-image-diff-ui="[[useNewImageDiffUi]]"
+    >
+      <table
+        id="diffTable"
+        class$="[[_diffTableClass]]"
+        role="presentation"
+        contenteditable$="[[isContentEditable]]"
+      ></table>
 
-        <template
-          is="dom-if"
-          if="[[showNoChangeMessage(_loading, prefs, _diffLength, diff)]]"
-        >
-          <div class="whitespace-change-only-message">
-            This file only contains whitespace changes. Modify the whitespace
-            setting to see the changes.
-          </div>
-        </template>
-      </gr-diff-builder>
-    </gr-diff-highlight>
+      <template
+        is="dom-if"
+        if="[[showNoChangeMessage(_loading, prefs, _diffLength, diff)]]"
+      >
+        <div class="whitespace-change-only-message">
+          This file only contains whitespace changes. Modify the whitespace
+          setting to see the changes.
+        </div>
+      </template>
+    </gr-diff-builder>
   </div>
   <div class$="[[_computeNewlineWarningClass(_newlineWarning, _loading)]]">
     [[_newlineWarning]]
diff --git a/polygerrit-ui/app/embed/diff/gr-diff/gr-diff_test.js b/polygerrit-ui/app/embed/diff/gr-diff/gr-diff_test.js
index ba4fa9f..c8d8a2f 100644
--- a/polygerrit-ui/app/embed/diff/gr-diff/gr-diff_test.js
+++ b/polygerrit-ui/app/embed/diff/gr-diff/gr-diff_test.js
@@ -50,21 +50,21 @@
 
     setup(() => {
       element = basicFixture.instantiate();
-      sinon.stub(element.$.highlights, 'handleSelectionChange');
+      sinon.stub(element.highlights, 'handleSelectionChange');
     });
 
     test('enabled if logged in', async () => {
       element.loggedIn = true;
       emulateSelection();
       await flush();
-      assert.isTrue(element.$.highlights.handleSelectionChange.called);
+      assert.isTrue(element.highlights.handleSelectionChange.called);
     });
 
     test('ignored if logged out', async () => {
       element.loggedIn = false;
       emulateSelection();
       await flush();
-      assert.isFalse(element.$.highlights.handleSelectionChange.called);
+      assert.isFalse(element.highlights.handleSelectionChange.called);
     });
   });