fix: render inline math while typing

src/components/Editor.tsx

@@ -26,6 +26,7 @@ import {
 } from './editorRawModeSync'
 import { useRawModeWithFlush } from './useRawModeWithFlush'
 import { createArrowLigaturesExtension } from './arrowLigaturesExtension'
+import { createMathInputExtension } from './mathInputExtension'
 import { useFilenameAutolinkGuard } from './useFilenameAutolinkGuard'
 import './Editor.css'
 import './EditorTheme.css'
@@ -185,7 +186,7 @@ function useEditorSetup({
     schema,
     uploadFile: (file: File) => uploadImageFile(file, vaultPathRef.current),
     _tiptapOptions: { injectNonce: RUNTIME_STYLE_NONCE },
-    extensions: [createArrowLigaturesExtension()],
+    extensions: [createArrowLigaturesExtension(), createMathInputExtension()],
   })
   useFilenameAutolinkGuard(editor)
   const activeTab = tabs.find((t) => t.entry.path === activeTabPath) ?? null

src/components/mathInputExtension.test.ts

@@ -0,0 +1,165 @@
+import { describe, expect, it, vi } from 'vitest'
+import { createMathInputExtension } from './mathInputExtension'
+
+function createTransaction() {
+  const transaction = {
+    replaceWith: vi.fn(() => transaction),
+    insertText: vi.fn(() => transaction),
+    scrollIntoView: vi.fn(() => transaction),
+  }
+  return transaction
+}
+
+function createView(beforeText: string, transaction: ReturnType<typeof createTransaction>) {
+  const mathNode = { nodeSize: 1 }
+  const selection = {
+    from: beforeText.length,
+    to: beforeText.length,
+    $from: {
+      parent: {
+        isTextblock: true,
+        textBetween: vi.fn(() => beforeText),
+      },
+      parentOffset: beforeText.length,
+      marks: vi.fn(() => []),
+    },
+  }
+  const mathNodeType = { createChecked: vi.fn(() => mathNode) }
+  const view = {
+    composing: false,
+    dispatch: vi.fn(),
+    state: {
+      schema: { nodes: { mathInline: mathNodeType } },
+      selection,
+      storedMarks: null as Array<{ type: { name: string } }> | null,
+      tr: transaction,
+    },
+  }
+
+  return { mathNode, mathNodeType, view }
+}
+
+function createDom(registerBeforeInput: (listener: (event: InputEvent) => void) => void) {
+  const dom = {
+    addEventListener: vi.fn((type: string, listener: (event: InputEvent) => void) => {
+      if (type === 'beforeinput') {
+        registerBeforeInput(listener)
+      }
+    }),
+  }
+  return dom
+}
+
+function createFixture(beforeText = 'Inline $x^2$') {
+  let beforeInputListener: ((event: InputEvent) => void) | null = null
+  const transaction = createTransaction()
+  const { mathNode, mathNodeType, view } = createView(beforeText, transaction)
+  const dom = createDom((listener) => {
+    beforeInputListener = listener
+  })
+  const editor = {
+    _tiptapEditor: { view },
+    prosemirrorView: view,
+  }
+  const extension = createMathInputExtension()({ editor: editor as never })
+
+  return {
+    dom,
+    extension,
+    fireInput(event: Partial<InputEvent> = {}) {
+      if (!beforeInputListener) {
+        throw new Error('Math input extension did not register a beforeinput listener')
+      }
+
+      const inputEvent = {
+        data: ' ',
+        inputType: 'insertText',
+        isComposing: false,
+        preventDefault: vi.fn(),
+        ...event,
+      }
+
+      beforeInputListener(inputEvent as InputEvent)
+      return inputEvent
+    },
+    mathNode,
+    mathNodeType,
+    mount() {
+      const controller = new AbortController()
+      extension.mount?.({
+        dom: dom as never,
+        root: document,
+        signal: controller.signal,
+      })
+      return controller
+    },
+    transaction,
+    view,
+  }
+}
+
+describe('createMathInputExtension', () => {
+  it('registers a beforeinput listener when the editor mounts', () => {
+    const fixture = createFixture()
+
+    fixture.mount()
+
+    expect(fixture.dom.addEventListener).toHaveBeenCalledWith(
+      'beforeinput',
+      expect.any(Function),
+      expect.objectContaining({
+        capture: true,
+        signal: expect.any(AbortSignal),
+      }),
+    )
+  })
+
+  it('replaces completed inline math before inserting whitespace', () => {
+    const fixture = createFixture()
+    fixture.mount()
+
+    const event = fixture.fireInput()
+
+    expect(fixture.mathNodeType.createChecked).toHaveBeenCalledWith({ latex: 'x^2' })
+    expect(fixture.transaction.replaceWith).toHaveBeenCalledWith(7, 12, fixture.mathNode)
+    expect(fixture.transaction.insertText).toHaveBeenCalledWith(' ', 8)
+    expect(fixture.transaction.scrollIntoView).toHaveBeenCalled()
+    expect(fixture.view.dispatch).toHaveBeenCalledWith(fixture.transaction)
+    expect(event.preventDefault).toHaveBeenCalledTimes(1)
+  })
+
+  it('replaces completed inline math before a new paragraph without swallowing the newline', () => {
+    const fixture = createFixture()
+    fixture.mount()
+
+    const event = fixture.fireInput({ data: null, inputType: 'insertParagraph' })
+
+    expect(fixture.transaction.replaceWith).toHaveBeenCalledWith(7, 12, fixture.mathNode)
+    expect(fixture.transaction.insertText).not.toHaveBeenCalled()
+    expect(fixture.view.dispatch).toHaveBeenCalledWith(fixture.transaction)
+    expect(event.preventDefault).not.toHaveBeenCalled()
+  })
+
+  it('ignores non-whitespace text input', () => {
+    const fixture = createFixture()
+    fixture.mount()
+
+    const event = fixture.fireInput({ data: '.', inputType: 'insertText' })
+
+    expect(fixture.transaction.replaceWith).not.toHaveBeenCalled()
+    expect(fixture.view.dispatch).not.toHaveBeenCalled()
+    expect(event.preventDefault).not.toHaveBeenCalled()
+  })
+
+  it('ignores math-looking input inside inline code', () => {
+    const fixture = createFixture()
+    fixture.view.state.storedMarks = [{ type: { name: 'code' } }]
+    fixture.mount()
+
+    const event = fixture.fireInput()
+
+    expect(fixture.transaction.replaceWith).not.toHaveBeenCalled()
+    expect(fixture.view.dispatch).not.toHaveBeenCalled()
+    expect(event.preventDefault).not.toHaveBeenCalled()
+  })
+})

src/components/mathInputExtension.ts

@@ -0,0 +1,112 @@
+import { createExtension } from '@blocknote/core'
+import type { useCreateBlockNote } from '@blocknote/react'
+import { MATH_INLINE_TYPE, readCompletedInlineMathAtEnd } from '../utils/mathMarkdown'
+
+const INLINE_WHITESPACE_RE = /^[^\S\r\n]$/
+const NEWLINE_INPUT_TYPES = new Set(['insertParagraph', 'insertLineBreak'])
+type EditorViewLike = NonNullable<ReturnType<typeof useCreateBlockNote>['prosemirrorView']>
+
+interface CursorText {
+  beforeText: string
+  parentStart: number
+}
+
+interface InlineMathReplacement {
+  from: number
+  latex: string
+  to: number
+}
+
+function isInsertedInlineWhitespace(event: InputEvent): event is InputEvent & { data: string } {
+  return event.inputType === 'insertText'
+    && typeof event.data === 'string'
+    && INLINE_WHITESPACE_RE.test(event.data)
+}
+
+function shouldHandleInput(event: InputEvent): boolean {
+  return isInsertedInlineWhitespace(event) || NEWLINE_INPUT_TYPES.has(event.inputType)
+}
+
+function shouldSkipInput(event: InputEvent, view: EditorViewLike): boolean {
+  if (event.isComposing) return true
+  if (view.composing) return true
+  return !shouldHandleInput(event)
+}
+
+function selectionHasCodeMark(view: EditorViewLike): boolean {
+  const marks = view.state.storedMarks ?? view.state.selection.$from.marks()
+  return marks.some((mark: { type: { name: string } }) => mark.type.name === 'code')
+}
+
+function readCursorText(view: EditorViewLike): CursorText | null {
+  const { from, to, $from } = view.state.selection
+  if (from !== to) return null
+  if (!$from.parent.isTextblock) return null
+
+  return {
+    beforeText: $from.parent.textBetween(0, $from.parentOffset, '', ''),
+    parentStart: from - $from.parentOffset,
+  }
+}
+
+function readInlineMathReplacement(view: EditorViewLike): InlineMathReplacement | null {
+  if (selectionHasCodeMark(view)) return null
+
+  const cursorText = readCursorText(view)
+  if (!cursorText) return null
+
+  const math = readCompletedInlineMathAtEnd({ text: cursorText.beforeText })
+  if (!math) return null
+
+  return {
+    from: cursorText.parentStart + math.start,
+    latex: math.latex,
+    to: cursorText.parentStart + math.end + 1,
+  }
+}
+
+function replaceCompletedInlineMath(
+  view: EditorViewLike,
+  trailingText?: string,
+): EditorViewLike['state']['tr'] | null {
+  const replacement = readInlineMathReplacement(view)
+  const mathNodeType = view.state.schema.nodes[MATH_INLINE_TYPE]
+  if (!replacement || !mathNodeType) return null
+
+  const mathNode = mathNodeType.createChecked({ latex: replacement.latex })
+  const transaction = view.state.tr.replaceWith(replacement.from, replacement.to, mathNode)
+
+  if (trailingText !== undefined) {
+    transaction.insertText(trailingText, replacement.from + mathNode.nodeSize)
+  }
+
+  return transaction.scrollIntoView()
+}
+
+export const createMathInputExtension = createExtension(({ editor }) => {
+  const readView = () => editor._tiptapEditor?.view ?? editor.prosemirrorView
+
+  return {
+    key: 'mathInput',
+    mount: ({ dom, signal }) => {
+      const handleBeforeInput = (event: InputEvent) => {
+        const view = readView()
+        if (!view || shouldSkipInput(event, view)) return
+
+        const trailingText = isInsertedInlineWhitespace(event) ? event.data : undefined
+        const transaction = replaceCompletedInlineMath(view, trailingText)
+        if (!transaction) return
+
+        view.dispatch(transaction)
+        if (trailingText !== undefined) {
+          event.preventDefault()
+        }
+      }
+
+      dom.addEventListener('beforeinput', handleBeforeInput as EventListener, {
+        capture: true,
+        signal,
+      })
+    },
+  } as const
+})