Use a more robust highlighter

lib/main/atom-ide/datatipProvider.tsx

@@ -66,7 +66,7 @@ async function highlightCode(code: string) {
     <div
       style={{fontFamily}}
       className="atom-typescript-datatip-tooltip-code"
-      dangerouslySetInnerHTML={{__html: html.join("\n")}}
+      dangerouslySetInnerHTML={{__html: html}}
     />
   )
 }

lib/main/atom/commands/findReferences.tsx

@@ -38,9 +38,8 @@ export async function handleFindReferencesResult(
         ref.contextStart !== undefined && ref.contextEnd !== undefined
           ? fileContents.slice(ref.contextStart.line - 1, ref.contextEnd.line)
           : fileContents
-      const fileHlText = await highlight(context.join("\n"), "source.tsx")
-      // tslint:disable-next-line: strict-boolean-expressions
-      const lineText = fileHlText[ref.start.line - (ref.contextStart?.line || 1)]
+      const fileHlText = (await highlight(context.join("\n"), "source.tsx")).split("\n")
+      const lineText = fileHlText[ref.start.line - (ref.contextStart?.line ?? 1)]
       return {...ref, hlText: lineText}
     }),
   )

lib/main/atom/utils/atom.ts

@@ -3,6 +3,8 @@ import {memoize, throttle} from "lodash"
 import * as path from "path"
 import {FileLocationQuery, Location, pointToLocation} from "./ts"
 
+export {highlight} from "./highlighter"
+
 // Return line/offset position in the editor using 1-indexed coordinates
 function getEditorPosition(editor: Atom.TextEditor): Location {
   const pos = editor.getCursorBufferPosition()
@@ -82,43 +84,3 @@ export function* getOpenEditorsPaths() {
     if (isTypescriptEditorWithPath(ed)) yield ed.getPath()!
   }
 }
-
-export async function highlight(code: string, scopeName: string) {
-  const ed = new Atom.TextEditor({
-    readonly: true,
-    keyboardInputEnabled: false,
-    showInvisibles: false,
-    tabLength: atom.config.get("editor.tabLength"),
-  })
-  const el = atom.views.getView(ed)
-  try {
-    el.setUpdatedSynchronously(true)
-    el.style.pointerEvents = "none"
-    el.style.position = "absolute"
-    el.style.top = "100vh"
-    el.style.width = "100vw"
-    atom.grammars.assignLanguageMode(ed.getBuffer(), scopeName)
-    ed.setText(code)
-    ed.scrollToBufferPosition(ed.getBuffer().getEndPosition())
-    atom.views.getView(atom.workspace).appendChild(el)
-    await editorTokenized(ed)
-    return Array.from(el.querySelectorAll(".line:not(.dummy)")).map((x) => x.innerHTML)
-  } finally {
-    el.remove()
-  }
-}
-
-async function editorTokenized(editor: Atom.TextEditor) {
-  return new Promise((resolve) => {
-    const languageMode = editor.getBuffer().getLanguageMode()
-    const nextUpdatePromise = editor.component.getNextUpdatePromise()
-    if (languageMode.fullyTokenized || languageMode.tree) {
-      resolve(nextUpdatePromise)
-    } else {
-      const disp = editor.onDidTokenize(() => {
-        disp.dispose()
-        resolve(nextUpdatePromise)
-      })
-    }
-  })
-}

lib/main/atom/utils/highlighter-types.d.ts

@@ -0,0 +1,35 @@
+import {Disposable, TextBuffer} from "atom"
+import type {} from "../../../typings/atom" // pull in types
+export {LanguageMode}
+
+interface LanguageMode {
+  readonly fullyTokenized?: boolean
+  readonly tree?: boolean
+  onDidTokenize?(cb: () => void): Disposable
+  buildHighlightIterator(): HighlightIterator
+  classNameForScopeId(id: ScopeId): string
+  startTokenizing?(): void
+}
+
+interface HighlightIterator {
+  seek(pos: {row: number; column: number}): void
+  getPosition(): {row: number; column: number}
+  getOpenScopeIds?(): ScopeId[]
+  getCloseScopeIds?(): ScopeId[]
+  moveToSuccessor(): void
+}
+
+interface ScopeId {}
+
+declare module "atom/dependencies/text-buffer/src/text-buffer" {
+  interface TextBuffer {
+    setLanguageMode(lm: LanguageMode): void
+  }
+}
+
+declare module "atom" {
+  interface GrammarRegistry {
+    grammarForId(id: string): Grammar
+    languageModeForGrammarAndBuffer(g: Grammar, b: TextBuffer): LanguageMode
+  }
+}

lib/main/atom/utils/highlighter.ts

@@ -0,0 +1,91 @@
+import {TextBuffer} from "atom"
+import type {LanguageMode} from "./highlighter-types"
+
+function eventLoopYielder(delayMs: number, maxTimeMs: number) {
+  const started = performance.now()
+  let lastYield = started
+  // tslint:disable-next-line: only-arrow-functions
+  return async function (): Promise<boolean> {
+    const now = performance.now()
+    if (now - lastYield > delayMs) {
+      await new Promise(setImmediate)
+      lastYield = now
+    }
+    return now - started <= maxTimeMs
+  }
+}
+
+/** Throws maximum time reached error */
+function maxTimeError(name: string, timeS: number) {
+  const err = new Error("Max time reached")
+  atom.notifications.addError(`${name} took more than ${timeS} seconds to complete`, {
+    dismissable: true,
+    description: `${name} took too long to complete and was terminated.`,
+    stack: err.stack,
+  })
+  return err
+}
+
+export async function highlight(sourceCode: string, scopeName: string) {
+  const yielder = eventLoopYielder(100, 5000)
+  const buf = new TextBuffer()
+  try {
+    const grammar = atom.grammars.grammarForId(scopeName)
+    const lm = atom.grammars.languageModeForGrammarAndBuffer(grammar, buf)
+    buf.setLanguageMode(lm)
+    buf.setText(sourceCode)
+    const end = buf.getEndPosition()
+    if (lm.startTokenizing) lm.startTokenizing()
+    await tokenized(lm)
+    const iter = lm.buildHighlightIterator()
+    if (iter.getOpenScopeIds && iter.getCloseScopeIds) {
+      let pos = {row: 0, column: 0}
+      iter.seek(pos)
+      const res = []
+      while (pos.row < end.row || (pos.row === end.row && pos.column <= end.column)) {
+        res.push(
+          ...iter.getCloseScopeIds().map(() => "</span>"),
+          ...iter.getOpenScopeIds().map((x) => `<span class="${lm.classNameForScopeId(x)}">`),
+        )
+        iter.moveToSuccessor()
+        const nextPos = iter.getPosition()
+        res.push(escapeHTML(buf.getTextInRange([pos, nextPos])))
+
+        if (!(await yielder())) {
+          console.error(maxTimeError("Atom-TypeScript: Highlighter", 5))
+          break
+        }
+        pos = nextPos
+      }
+      return res.join("")
+    } else {
+      return sourceCode
+    }
+  } finally {
+    buf.destroy()
+  }
+}
+
+async function tokenized(lm: LanguageMode) {
+  return new Promise((resolve) => {
+    if (lm.fullyTokenized || lm.tree) {
+      resolve(undefined)
+    } else if (lm.onDidTokenize) {
+      const disp = lm.onDidTokenize(() => {
+        disp.dispose()
+        resolve(undefined)
+      })
+    } else {
+      resolve(undefined) // null language mode
+    }
+  })
+}
+
+function escapeHTML(str: string) {
+  return str
+    .replace(/&/g, "&amp;")
+    .replace(/</g, "&lt;")
+    .replace(/>/g, "&gt;")
+    .replace(/"/g, "&quot;")
+    .replace(/'/g, "&#039;")
+}

lib/typings/atom.d.ts

@@ -1,13 +1,15 @@
 export {}
 import {Disposable} from "atom"
-declare module "atom" {
+declare module "atom/dependencies/text-buffer/src/text-buffer" {
   interface TextBuffer {
     emitDidStopChangingEvent(): void
     getLanguageMode(): {
       readonly fullyTokenized?: boolean
       readonly tree?: boolean
     }
   }
+}
+declare module "atom" {
   interface TextEditor {
     onDidTokenize(callback: () => void): Disposable
     isDestroyed(): boolean