Repository: h2oai/wave
File: ui/src/markdown.tsx (TypeScript)

Code before:
  return highlightedCode
}

export const Markdown = ({ source }: { source: S }) => {
  const
    prevHighlights = React.useRef<S[]>([]), // Prevent flicker during streaming.
    codeBlockIdx = React.useRef(0), // MarkdownIt parses code blocks sequentially, which is a problem for streaming.
    markdown = React.useMemo(() => MarkdownIt({
      html: true, linkify: true, typographer: true, highlight: (str, lang) => {
        const codeBlockId = codeBlockIdx.current.toString()
        if (prevHighlights.current.length === codeBlockIdx.current) prevHighlights.current.push('')

        // HACK: MarkdownIt does not support async rules.
        // https://github.com/markdown-it/markdown-it/blob/master/docs/development.md#i-need-async-rule-how-to-do-it
        setTimeout(async () => prevHighlights.current[+codeBlockId] = await highlightSyntax(str, lang, codeBlockId).finally(() => {
          // Add copy button once code block is higlighted.
          const codeBlockContainer = document.getElementById(codeBlockId)?.parentElement
          if (codeBlockContainer) {
            const buttonContainer = document.createElement('span')
            ReactDOM.render(<ClipboardCopyButton value={str} anchorElement={codeBlockContainer} showOnHoverOnly />, buttonContainer,
              () => {
                codeBlockContainer.style.position = 'relative'
                codeBlockContainer.insertAdjacentElement('afterbegin', buttonContainer)
              })
          }
        }), 0)

        // TODO: Sanitize the HTML.
        const ret = `<code id='${codeBlockId}' class="hljs">${prevHighlights.current[codeBlockIdx.current] || str}</code>`
        codeBlockIdx.current++
        return ret
      }
    }), []),
    html = React.useMemo(() => {
      const html = markdown.render(source)
      codeBlockIdx.current = 0
      return html
    }, [markdown, source]),
    onClick = (e: React.MouseEvent<HTMLDivElement>) => {
      const hrefAttr = (e.target as HTMLAnchorElement).getAttribute('href')
      if (e.target instanceof HTMLAnchorElement && hrefAttr?.startsWith('?')) {
        // Handled in app.tsx. 
        window.dispatchEvent(new CustomEvent('md-link-click', { detail: hrefAttr.substring(1) }))
        e.preventDefault()
        e.stopPropagation()
        // Prevent navigation behavior.
        return false
      }
    }
  return <div onClick={onClick} className={clas(css.markdown, 'wave-markdown')} dangerouslySetInnerHTML={{ __html: html }} />
}

Review comment (suggestion):
* Do not mix async/await with promise callbacks.
* Why is finally needed?
* Why is ReactDOM.render callback needed?
* Can it be part of `highlightSyntax` function instead? That way, the codeBlockContainer would be available.
* Why is `insertAdjacentElement('afterbegin'` needed as opposed to an `appendChild`? I thought button is absolutely positioned so its insertion order does not matter.

Code after:
      // TS cannot do dynamic JSON imports properly. Use any as a workaround.
      // https://stackoverflow.com/questions/70601733/dynamic-import-with-json-file-doesnt-work-typescript.
      const langAliases: any = await import('./markdownCodeSyntaxHighlighting.json')
      language = langAliases[language]
      if (language && !hljs.getLanguage(language)) {
        // Need to use relative path due to https://www.npmjs.com/package/@rollup/plugin-dynamic-import-vars.
        const languageModule = await import(`../node_modules/highlight.js/es/languages/${language}.js`)
        hljs.registerLanguage(language, languageModule.default)
      }
    } catch (e) {
      language = '' // Fallback to auto-detection.
    }
  }
  const highlightedCode = language
    ? hljs.highlight(str, { language, ignoreIllegals: true }).value
    : hljs.highlightAuto(str).value

  codeBlock.innerHTML = highlightedCode
  const btnContainer = document.createElement('div')
  btnContainer.classList.add(css.copyBtnWrapper)
  ReactDOM.render(<ClipboardCopyButton value={str} />, codeBlock.appendChild(btnContainer))

  return highlightedCode
}

export const Markdown = ({ source }: { source: S }) => {
  const
    prevHighlights = React.useRef<S[]>([]), // Prevent flicker during streaming.
    codeBlockIdx = React.useRef(0), // MarkdownIt parses code blocks sequentially, which is a problem for streaming.
    markdown = React.useMemo(() => MarkdownIt({
      html: true, linkify: true, typographer: true, highlight: (str, lang) => {
        const codeBlockId = codeBlockIdx.current.toString()
        if (prevHighlights.current.length === codeBlockIdx.current) prevHighlights.current.push('')

        // HACK: MarkdownIt does not support async rules.
        // https://github.com/markdown-it/markdown-it/blob/master/docs/development.md#i-need-async-rule-how-to-do-it
        setTimeout(async () => prevHighlights.current[+codeBlockId] = await highlightSyntax(str, lang, codeBlockId), 0)

        // TODO: Sanitize the HTML.
        const ret = `<code id='${codeBlockId}' class="hljs ${css.codeblock}">${prevHighlights.current[codeBlockIdx.current] || str}</code>`
        codeBlockIdx.current++
        return ret
      }
    }), []),
    html = React.useMemo(() => {
      const html = markdown.render(source)
      codeBlockIdx.current = 0
      return html
    }, [markdown, source]),
    onClick = (e: React.MouseEvent<HTMLDivElement>) => {
      const hrefAttr = (e.target as HTMLAnchorElement).getAttribute('href')