latex-render

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Mode, Previewed } from '../types'
5
6import {
7  dvipngColor,
8  figureDocument,
9  hashKey,
10  hasMath,
11  inlineToUnicode,
12  isSafeTex,
13  latexError,
14  nextMode,
15  pngSize,
16  segmentsToUnicode,
17  settledLength,
18  splitMath,
19  styleGuide,
20  texDocument,
21  texToUnicode,
22} from './latex'
23import type { Segment } from './latex'
24
25// Pixels per terminal row the formulas are typeset at; the terminal scales
26// the picture to its cells, so this only sets how crisp they are.
27const ROW_PX = 40
28const VERSION = 1
29// The most rows the preview above the prompt takes.
30const PREVIEW_ROWS = 12
31
32type Display = Segment & { kind: 'display' }
33type Figure = Segment & { kind: 'figure' }
34type Drawn = Display | Figure
35type Palette = { ink: string; paper?: string }
36type Typeset = { path: string; width: number; height: number }
37type Picture = { path: string; columns: number; rows: number }
38
39type Settings = {
40  images: string
41  color: string
42  scale: number
43  cellAspect: number
44  preamble: string
45  background: string
46  preview: boolean
47}
48
49// What the module learns once per load: its options, whether pictures can be
50// drawn, the cache folder, the colors, and every formula and figure typeset so far.
51let settings: Settings = { images: 'auto', color: 'auto', scale: 1.25, cellAspect: 2.1, preamble: '', background: 'auto', preview: true }
52let canDraw: Promise<boolean> | undefined
53let canDrawFigures: Promise<boolean> | undefined
54let cacheDir: Promise<string> | undefined
55let colors: Promise<Palette> | undefined
56// `/latex` flips this; in `source` every reply is drawn as written.
57const mode = atom({ plugin: 'latex-render', key: 'mode' } as const, 'render' as Mode)
58// The formulas and figures of the block being streamed, shown above the prompt until it ends.
59const preview = atom({ plugin: 'latex-render', key: 'preview' } as const, [] as Previewed[])
60// Counts the blocks the preview has been cleared for, so a picture typeset late is not shown under the next.
61let previewed = 0
62const typeset = new Map<string, Promise<Typeset | undefined>>()
63const fitted = new Map<string, Promise<Picture | undefined>>()
64// The width the last reply was drawn at, so a formula is fitted to it while still streaming.
65let lastRoom: number | undefined
66
67export const register: Register = (on, options) => {
68  settings = {
69    images: String(options.images ?? 'auto'),
70    color: String(options.color ?? 'auto'),
71    scale: clamp(Number(options.scale) || 1.25, 0.5, 4),
72    cellAspect: clamp(Number(options.cellAspect) || 2.1, 1, 4),
73    preamble: String(options.preamble ?? ''),
74    background: String(options.background ?? 'auto'),
75    preview: String(options.preview ?? 'on') !== 'off',
76  }
77  // What was found for the previous options no longer holds.
78  canDraw = canDrawFigures = colors = undefined
79
80  // Tells the model how to write math and figures so they are drawn well here.
81  on('prompt.compose', async ($, e, next) => {
82    const composed = await next(e)
83    if (!e.surfaces.includes('terminal')) return composed
84    const [pictures, figures, { paper }] = await Promise.all([imagesOn($), figuresOn($), palette($)])
85    const section = { id: 'latex-render:style', text: styleGuide({ pictures, figures, paper }), scope: 'session' } as const
86    return { sections: [...composed.sections, section] }
87  })
88
89  on('session.start', async ($, e, next) => {
90    await $.command.register({
91      name: 'latex',
92      description: 'Toggle LaTeX in replies between rendered and as written',
93      argumentHint: '[render|source]',
94      immediate: true,
95    })
96    // A reload keeps the mode; the status line has to say so again.
97    showMode($, await read($, mode))
98    return next(e)
99  })
100
101  on('command.run', { command: 'latex' }, async ($, e) => {
102    const chosen = nextMode(await read($, mode), e.args)
103    if (!chosen) return { text: 'Usage: /latex [render|source]' }
104    await update($, mode, () => chosen)
105    showMode($, chosen)
106    return { text: chosen === 'source' ? 'LaTeX is shown as written.' : 'LaTeX is rendered.' }
107  })
108
109  // While a reply streams, each display formula and figure is typeset as
110  // soon as it closes and shown above the prompt, its picture then ready, or
111  // nearly, when the block ends and is drawn. The chunks pass on unchanged.
112  on('turn.step', async function* ($, e, next) {
113    const onTerminal = e.agentId === undefined && (await $.session.surfaces()).includes('terminal')
114    if (!onTerminal || (await read($, mode)) === 'source' || !(await imagesOn($))) return yield* next(e)
115    let block: number | undefined // the index of the text block being written
116    let text = '' // its text so far
117    let settled = 0 // how much of it has been handed to the typesetter
118    let isShowing = false // the preview may hold this block's pictures
119    // The block's row is drawn when it ends, so its preview goes.
120    const endPreview = () => {
121      if (!isShowing) return
122      isShowing = false
123      previewed += 1
124      void update($, preview, () => []).catch(() => undefined)
125    }
126
127    try {
128      for await (const chunk of next(e)) {
129        yield chunk
130        if (chunk.kind !== 'text') {
131          endPreview()
132          continue
133        }
134        if (chunk.index !== block) {
135          endPreview()
136          block = chunk.index
137          text = ''
138          settled = 0
139        }
140        text += chunk.text
141        const upTo = settledLength(text)
142        if (upTo > settled) {
143          const piece = text.slice(settled, upTo)
144          void prepare($, piece)
145          if (settings.preview) {
146            isShowing = true
147            void showPreview($, piece, previewed)
148          }
149          settled = upTo
150        }
151      }
152    } finally {
153      endPreview()
154    }
155  })
156
157  // Above the prompt, while Claude writes: the newest formulas and figures of
158  // the block it is writing, as many as the band has rows for.
159  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
160    // Both read first, so the band is drawn again when either changes.
161    const [items, current] = await Promise.all([read($, preview), read($, mode)])
162    const { hasSurvey, isWorking, maxRows, bodyColumns, view } = e.props
163    // A strip, not a page: outside fullscreen `maxRows` is the terminal's height. A row goes to the heading.
164    const budget = Math.min(maxRows, PREVIEW_ROWS) - 1
165    const isQuiet = items.length === 0 || current === 'source' || hasSurvey || !isWorking || view.agentId !== undefined
166    if (isQuiet || e.surface !== 'terminal' || budget < 1) return next(e)
167
168    const room = Math.max(10, bodyColumns - 2)
169    const pictures = await Promise.all(items.map(item => fit($, item, room, budget)))
170    const shown: { item: Previewed; picture: Picture; index: number }[] = []
171    let left = budget
172    for (let index = items.length - 1; index >= 0; index--) {
173      const item = items[index]
174      const picture = pictures[index]
175      if (!item || !picture) continue
176      if (picture.rows > left) break
177      shown.unshift({ item, picture, index })
178      left -= picture.rows
179    }
180    if (shown.length === 0) return next(e)
181
182    const { Box, Text, Image } = $.ui.resolve(e)
183    return (
184      <Box flexDirection="column">
185        <Text dimColor>LaTeX preview</Text>
186        {shown.map(({ item, picture, index }) => (
187          <Box marginLeft={2}>
188            <Image
189              key={`preview-${index}`}
190              source={{ file: picture.path, format: 'png' }}
191              columns={picture.columns}
192              rows={picture.rows}
193              alt={item.kind === 'figure' ? 'TikZ figure' : texToUnicode(item.tex) || ' '}
194            />
195          </Box>
196        ))}
197      </Box>
198    )
199  })
200
201  on('ui.render', { component: 'AssistantMessage' }, async ($, e, next) => {
202    // Read before anything else, so every reply is drawn again when it changes.
203    if ((await read($, mode)) === 'source') return next(e)
204    const text = e.props.text
205    if (!hasMath(text)) return next(e)
206
207    const segments = splitMath(text)
208    const hasDisplay = segments.some(s => s.kind !== 'markdown')
209    // A summary row stays one line of text: its math in Unicode, never a picture.
210    const drawsPictures = !e.props.isSummary && e.surface === 'terminal' && hasDisplay && (await imagesOn($))
211
212    if (!drawsPictures) {
213      const rewritten = segmentsToUnicode(segments)
214      return rewritten === text ? next(e) : next({ ...e, props: { ...e.props, text: rewritten } })
215    }
216
217    const room = Math.max(10, (e.viewport?.columns ?? 80) - 4)
218    lastRoom = room
219    const pictures = await Promise.all(
220      segments.map(s => (s.kind === 'markdown' ? Promise.resolve(undefined) : fit($, s, room))),
221    )
222
223    const { Box, Text, Markdown, Image } = $.ui.resolve(e)
224    const children: JSX.Element[] = []
225    let pendingText = ''
226    const flushText = () => {
227      const md = pendingText.replace(/^\s*\n/, '').replace(/\n\s*$/, '')
228      if (md.trim() !== '') children.push(<Markdown text={inlineToUnicode(md)} />)
229      pendingText = ''
230    }
231
232    segments.forEach((s, index) => {
233      const picture = pictures[index]
234      if (s.kind === 'markdown') {
235        pendingText += s.text
236      } else if (!picture) {
237        pendingText += segmentsToUnicode([s])
238      } else {
239        flushText()
240        children.push(
241          <Box marginLeft={2}>
242            <Image
243              key={`${s.kind === 'figure' ? 'figure' : 'math'}-${index}`}
244              source={{ file: picture.path, format: 'png' }}
245              columns={picture.columns}
246              rows={picture.rows}
247              alt={s.kind === 'figure' ? 'TikZ figure' : texToUnicode(s.tex) || ' '}
248            />
249          </Box>,
250        )
251      }
252    })
253    flushText()
254
255    return (
256      <Box flexDirection="row">
257        <Box width={2} flexShrink={0}>
258          <Text>{e.props.isFirstOfReply ? '●' : ' '}</Text>
259        </Box>
260        <Box flexDirection="column" flexGrow={1}>
261          {children}
262        </Box>
263      </Box>
264    )
265  })
266}
267
268function drawnIn(text: string): Drawn[] {
269  return splitMath(text).filter((s): s is Drawn => s.kind !== 'markdown')
270}
271
272// Typesets the formulas and figures in a piece of a reply ahead of its drawing.
273function prepare($: EngineInterface, text: string): Promise<unknown> {
274  return Promise.all(drawnIn(text).map(s => (lastRoom === undefined ? typesetSegment($, s) : fit($, s, lastRoom))))
275}
276
277// Adds each formula and figure in a piece of a reply to the preview once it is
278// typeset, in order, unless the block it belongs to has ended by then.
279async function showPreview($: EngineInterface, text: string, block: number): Promise<void> {
280  for (const segment of drawnIn(text)) {
281    const done = await typesetSegment($, segment)
282    if (!done || block !== previewed) continue
283    const item: Previewed =
284      segment.kind === 'figure'
285        ? { kind: 'figure', tex: segment.tex }
286        : segment.env
287          ? { kind: 'display', tex: segment.tex, env: segment.env }
288          : { kind: 'display', tex: segment.tex }
289    await update($, preview, items => [...items, item].slice(-8)).catch(() => undefined)
290  }
291}
292
293function showMode($: EngineInterface, current: Mode): void {
294  $.ui.status(current === 'source' ? 'LaTeX: source (/latex to render)' : undefined)
295}
296
297function imagesOn($: EngineInterface): Promise<boolean> {
298  canDraw ??= detectImages($, settings.images)
299  return canDraw
300}
301
302async function detectImages($: EngineInterface, option: string): Promise<boolean> {
303  if (option === 'off') return false
304  const tools = await Promise.all([
305    $.process.run(['latex', '--version']).then(r => r.exitCode === 0, () => false),
306    $.process.run(['dvipng', '--version']).then(r => r.exitCode === 0, () => false),
307  ])
308  if (!tools.every(Boolean)) return false
309  if (option === 'on') return true
310
311  const [term, program, kitty, tmux, ssh] = await Promise.all([
312    $.env.get('TERM'),
313    $.env.get('TERM_PROGRAM'),
314    $.env.get('KITTY_WINDOW_ID'),
315    $.env.get('TMUX'),
316    $.env.get('SSH_CONNECTION'),
317  ])
318  // The terminal reads the picture from this machine's files over the kitty
319  // graphics protocol, which tmux does not pass and ssh cannot reach.
320  if (tmux || ssh) return false
321  const name = `${term ?? ''} ${program ?? ''}`.toLowerCase()
322  return Boolean(kitty) || /kitty|ghostty|wezterm/.test(name)
323}
324
325// Figures take pdflatex and pdftocairo besides what formulas take.
326function figuresOn($: EngineInterface): Promise<boolean> {
327  canDrawFigures ??= imagesOn($).then(async on => {
328    if (!on) return false
329    const tools = await Promise.all(
330      ['pdflatex', 'pdftocairo'].map(tool => $.process.run([tool, '-v']).then(r => r.exitCode === 0, () => false)),
331    )
332    return tools.every(Boolean)
333  })
334  return canDrawFigures
335}
336
337function cache($: EngineInterface): Promise<string> {
338  cacheDir ??= findCacheDir($)
339  return cacheDir
340}
341
342async function findCacheDir($: EngineInterface): Promise<string> {
343  const [xdg, home] = await Promise.all([$.env.get('XDG_CACHE_HOME'), $.env.get('HOME')])
344  return `${xdg || `${home ?? '/tmp'}/.cache`}/claude-latex-render`
345}
346
347function palette($: EngineInterface): Promise<Palette> {
348  colors ??= pickPalette($)
349  return colors
350}
351
352// The ink, and on a dark theme the paper figures' `white` stands for.
353async function pickPalette($: EngineInterface): Promise<Palette> {
354  const rows = await $.config.list().catch(() => [])
355  const light = String(rows.find(row => row.key === 'theme')?.value ?? 'dark').includes('light')
356  const hex = (option: string, fallback: string) => (dvipngColor(option) ? option.replace('#', '') : fallback)
357  return {
358    ink: hex(settings.color, light ? '1f1f1f' : 'e6e6e6'),
359    paper: light ? undefined : hex(settings.background, '1e1e1e'),
360  }
361}
362
363// Typesets one formula or figure to a tight transparent PNG, cached on disk by content.
364function typesetSegment($: EngineInterface, segment: Drawn): Promise<Typeset | undefined> {
365  const id = segment.kind === 'figure' ? `figure\u0000${segment.tex}` : `${segment.env}\u0000${segment.tex}`
366  let pending = typeset.get(id)
367  if (!pending) {
368    pending = (segment.kind === 'figure' ? runTikz($, segment) : runLatex($, segment)).catch(() => undefined)
369    typeset.set(id, pending)
370  }
371  return pending
372}
373
374async function runLatex($: EngineInterface, segment: Display): Promise<Typeset | undefined> {
375  if (!isSafeTex(segment.tex)) return undefined
376  const dir = await cache($)
377  const { ink, paper } = await palette($)
378  const fg = dvipngColor(ink) ?? 'rgb 0.9 0.9 0.9'
379  const doc = texDocument(segment, settings.preamble, ink, paper)
380  const dpi = dpiFor()
381  const key = hashKey([VERSION, doc, fg, dpi].join('\u0000'))
382  const png = `${dir}/${key}.png`
383  const failed = `${dir}/${key}.err`
384
385  if (await $.fs.exists(failed)) return undefined
386  if (!(await $.fs.exists(png))) {
387    await $.fs.write(`${dir}/${key}.tex`, doc)
388    const latex = await $.process.run(
389      ['latex', '-no-shell-escape', '-interaction=nonstopmode', '-halt-on-error', `${key}.tex`],
390      { cwd: dir, env: { openout_any: 'p', openin_any: 'p' }, timeoutMs: 15000 },
391    )
392    if (latex.exitCode !== 0) {
393      await $.fs.write(failed, latexError(latex.stdout))
394      return undefined
395    }
396    const dvipng = await $.process.run(
397      ['dvipng', '-q', '-D', String(dpi), '-T', 'tight', '-bg', 'Transparent', '-fg', fg, '-o', `${key}.png`, `${key}.dvi`],
398      { cwd: dir, timeoutMs: 15000 },
399    )
400    await $.process.run(['rm', '-f', `${key}.aux`, `${key}.log`, `${key}.dvi`], { cwd: dir }).catch(() => undefined)
401    if (dvipng.exitCode !== 0) {
402      await $.fs.write(failed, dvipng.stderr || 'dvipng failed')
403      return undefined
404    }
405  }
406
407  return measure($, png)
408}
409
410// A figure's text is typeset at the same size as a formula's.
411function dpiFor(): number {
412  return Math.round((settings.scale * ROW_PX * 72.27) / 12)
413}
414
415async function measure($: EngineInterface, png: string): Promise<Typeset | undefined> {
416  const { base64 } = await $.fs.read(png, { as: 'bytes' })
417  const size = pngSize(base64)
418  return size ? { path: png, ...size } : undefined
419}
420
421async function runTikz($: EngineInterface, segment: Figure): Promise<Typeset | undefined> {
422  if (!isSafeTex(segment.tex) || !(await figuresOn($))) return undefined
423  const dir = await cache($)
424  const { ink, paper } = await palette($)
425  const doc = figureDocument(segment.tex, settings.preamble, ink, paper)
426  const dpi = dpiFor()
427  const key = hashKey([VERSION, doc, dpi].join('\u0000'))
428  const png = `${dir}/${key}.png`
429  const failed = `${dir}/${key}.err`
430
431  if (await $.fs.exists(failed)) return undefined
432  if (!(await $.fs.exists(png))) {
433    await $.fs.write(`${dir}/${key}.tex`, doc)
434    const latex = await $.process.run(
435      ['pdflatex', '-no-shell-escape', '-interaction=nonstopmode', '-halt-on-error', `${key}.tex`],
436      { cwd: dir, env: { openout_any: 'p', openin_any: 'p' }, timeoutMs: 30000 },
437    )
438    if (latex.exitCode !== 0) {
439      await $.fs.write(failed, latexError(latex.stdout))
440      return undefined
441    }
442    // pdftocairo names the picture itself, adding `.png` to the stem.
443    const cairo = await $.process.run(
444      ['pdftocairo', '-png', '-transp', '-singlefile', '-r', String(dpi), `${key}.pdf`, key],
445      { cwd: dir, timeoutMs: 15000 },
446    )
447    await $.process.run(['rm', '-f', `${key}.aux`, `${key}.log`, `${key}.pdf`], { cwd: dir }).catch(() => undefined)
448    if (cairo.exitCode !== 0) {
449      await $.fs.write(failed, cairo.stderr || 'pdftocairo failed')
450      return undefined
451    }
452  }
453  return measure($, png)
454}
455
456// Pads (and if needed shrinks) a typeset formula or figure to a whole number
457// of cells, so the terminal's scaling to the box keeps its proportions.
458// Within `room` columns and, when given, `rowLimit` rows.
459function fit($: EngineInterface, segment: Drawn, room: number, rowLimit?: number): Promise<Picture | undefined> {
460  const id = `${segment.kind === 'figure' ? 'figure' : segment.env}\u0000${segment.tex}\u0000${room}\u0000${rowLimit ?? ''}`
461  let pending = fitted.get(id)
462  if (!pending) {
463    pending = padToCells($, segment, room, rowLimit).catch(() => undefined)
464    fitted.set(id, pending)
465  }
466  return pending
467}
468
469async function padToCells($: EngineInterface, segment: Drawn, room: number, rowLimit?: number): Promise<Picture | undefined> {
470  const base = await typesetSegment($, segment)
471  if (!base) return undefined
472  const cellW = ROW_PX / settings.cellAspect
473  const tall = rowLimit === undefined ? 1 : (rowLimit * ROW_PX) / base.height
474  const shrink = Math.min(1, (room * cellW) / base.width, tall)
475  const width = Math.max(1, Math.round(base.width * shrink))
476  const height = Math.max(1, Math.round(base.height * shrink))
477  const columns = clamp(Math.ceil(width / cellW), 1, 255)
478  const rows = clamp(Math.ceil(height / ROW_PX), 1, 255)
479  const out = base.path.replace(/\.png$/, `-${columns}x${rows}.png`)
480
481  if (!(await $.fs.exists(out))) {
482    const args = ['magick', base.path]
483    if (shrink < 1) args.push('-resize', `${width}x${height}!`)
484    args.push('-background', 'none', '-gravity', 'west', '-extent', `${Math.round(columns * cellW)}x${rows * ROW_PX}`, out)
485    const padded = await $.process.run(args, { timeoutMs: 15000 }).catch(() => undefined)
486    // Without ImageMagick the picture is drawn unpadded, a little stretched.
487    if (padded?.exitCode !== 0) return { path: base.path, columns, rows }
488  }
489  return { path: out, columns, rows }
490}
491
492function clamp(n: number, low: number, high: number): number {
493  return Math.min(high, Math.max(low, n))
494}
495

1// Pure LaTeX helpers: finding math and TikZ figures in markdown, Unicode
2// conversion, the documents a display formula or a figure is typeset in. No `$` here, so tests reach it all.
3
4import type { Mode } from '../types'
5
6export type Segment =
7  | { kind: 'markdown'; text: string }
8  | { kind: 'display'; tex: string; env?: string }
9  | { kind: 'figure'; tex: string }
10
11const DISPLAY_ENVS = 'equation|align|gather|multline|eqnarray|flalign|alignat'
12
13const FIGURE_ENVS = 'tikzpicture|tikzcd|axis'
14
15const DISPLAY = new RegExp(
16  String.raw`\$\$([\s\S]+?)\$\$|\\\[([\s\S]+?)\\\]|\\begin\{(${DISPLAY_ENVS})(\*?)\}([\s\S]*?)\\end\{\3\4\}` +
17    String.raw`|\\begin\{(${FIGURE_ENVS})\}[\s\S]*?\\end\{\6\}`,
18  'g',
19)
20
21// A `$` glued to a digit or a letter before it, or a digit after it, is a
22// price (`US$5`, `5$`, `$5-$10`), so `$n$th` and `_$x$_` are math and `$5-$10` is not.
23const INLINE = /\\\((.+?)\\\)|(?<![\\$0-9A-Za-z])\$(?![\s$])([^$\n]+?)(?<![\s\\])\$(?![$\d])/g
24
25// A number written as math, `$0$`, `$-3.5$` or `$1,000$`.
26const NUMBER = /^-?(?:\d{1,3}(?:,\d{3})+|\d+)(?:\.\d+)?$/
27
28const FENCE = /^ {0,3}(`{3,}|~{3,})/
29
30// A fenced block whose info string is `tikz` holds a figure.
31const TIKZ_FENCE = /^ {0,3}(`{3,}|~{3,})[ \t]*tikz[ \t]*$/m
32
33// Private-use characters stand in for inline code spans while math is found.
34const HOLE = ''
35
36/**
37 * The mode `/latex <args>` switches to from `current`: no argument flips it,
38 * `render` or `source` (or `on` / `off`) sets it, anything else is undefined.
39 */
40export function nextMode(current: Mode, args: string): Mode | undefined {
41  const word = args.trim().toLowerCase()
42  if (word === '') return current === 'render' ? 'source' : 'render'
43  if (word === 'render' || word === 'on') return 'render'
44  if (word === 'source' || word === 'off') return 'source'
45  return undefined
46}
47
48export function hasMath(text: string): boolean {
49  return (
50    text.includes('$') || text.includes('\\(') || text.includes('\\[') || text.includes('\\begin{') || TIKZ_FENCE.test(text)
51  )
52}
53
54/**
55 * Splits markdown into text, display-math and figure segments, leaving other
56 * fenced code and inline code spans alone. Inline math stays in the text
57 * segments. A figure is a closed ```tikz fence or a bare `tikzpicture`,
58 * `tikzcd` or pgfplots `axis` environment.
59 */
60export function splitMath(text: string): Segment[] {
61  const out: Segment[] = []
62  const pushText = (chunk: string) => {
63    if (chunk === '') return
64    const last = out[out.length - 1]
65    if (last?.kind === 'markdown') last.text += chunk
66    else out.push({ kind: 'markdown', text: chunk })
67  }
68
69  for (const block of splitFences(text)) {
70    if (block.isCode) {
71      const figure = tikzFence(block.text)
72      if (figure) out.push(figure)
73      else pushText(block.text)
74      continue
75    }
76    const { masked, spans } = maskCodeSpans(block.text)
77    let at = 0
78    for (const m of masked.matchAll(DISPLAY)) {
79      const start = m.index ?? 0
80      pushText(unmask(masked.slice(at, start), spans))
81      if (m[6]) {
82        out.push({ kind: 'figure', tex: unmask(m[0], spans).trim() })
83        at = start + m[0].length
84        continue
85      }
86      const env = m[3]
87      const tex = unmask(m[1] ?? m[2] ?? m[5] ?? '', spans).trim()
88      if (tex === '') pushText(unmask(m[0], spans))
89      else out.push(env ? { kind: 'display', tex, env: env + (m[4] ?? '') } : { kind: 'display', tex })
90      at = start + m[0].length
91    }
92    pushText(unmask(masked.slice(at), spans))
93  }
94
95  return out
96}
97
98/**
99 * How much of a reply still being written is settled: its length up to the
100 * end of the line on which its last finished display formula or figure
101 * closes, when nothing but spaces follows the close on that line; 0 when no
102 * formula or figure has closed so. Math inside code is never a close.
103 */
104export function settledLength(text: string): number {
105  let settled = 0
106  let offset = 0
107  for (const block of splitFences(text)) {
108    if (block.isCode) {
109      if (block.text.endsWith('\n') && tikzFence(block.text)) settled = offset + block.text.length
110    } else {
111      // Code spans keep their length, so offsets in the masked text hold in the reply.
112      const masked = block.text.replace(/(`+)[\s\S]*?\1/g, span => HOLE.repeat(span.length))
113      for (const m of masked.matchAll(DISPLAY)) {
114        const end = (m.index ?? 0) + m[0].length
115        const line = /^[ \t]*\n/.exec(masked.slice(end))
116        if (line) settled = offset + end + line[0].length
117      }
118    }
119    offset += block.text.length
120  }
121  return settled
122}
123
124/** Replaces inline math in markdown with Unicode, escaped for markdown. */
125export function inlineToUnicode(markdown: string): string {
126  return splitFences(markdown)
127    .map(block => {
128      if (block.isCode) return block.text
129      const { masked, spans } = maskCodeSpans(block.text)
130      const replaced = masked.replace(INLINE, (whole, paren?: string, dollar?: string) => {
131        const tex = paren ?? dollar ?? ''
132        if (NUMBER.test(tex)) return tex
133        // Digits that are no number (`$5,$`) are left as written.
134        if (dollar !== undefined && /^[\d.,\s]+$/.test(dollar)) return whole
135        return escapeMarkdown(texToUnicode(tex))
136      })
137      return unmask(replaced, spans)
138    })
139    .join('')
140}
141
142/** The whole reply as markdown with every formula in Unicode, every figure as its code. */
143export function segmentsToUnicode(segments: Segment[]): string {
144  return segments
145    .map(s =>
146      s.kind === 'markdown' ? inlineToUnicode(s.text) : s.kind === 'figure' ? figureToCode(s.tex) : displayToUnicode(s.tex),
147    )
148    .join('')
149}
150
151/** A figure that cannot be drawn, shown as the TikZ it was written in. */
152export function figureToCode(tex: string): string {
153  return '\n\n```tex\n' + tex + '\n```\n\n'
154}
155
156// The figure in a closed ```tikz fence, undefined for any other block.
157function tikzFence(block: string): Segment | undefined {
158  const lines = block.split(/(?<=\n)/)
159  const marks = TIKZ_FENCE.exec(lines[0] ?? '')?.[1]
160  const last = lines[lines.length - 1]?.trim() ?? ''
161  if (!marks || lines.length < 2 || last[0] !== marks[0] || !/^(`+|~+)$/.test(last) || last.length < marks.length) {
162    return undefined
163  }
164  const tex = lines.slice(1, -1).join('').trim()
165  return tex === '' ? undefined : { kind: 'figure', tex }
166}
167
168/** A display formula as indented Unicode lines, one per `\\` row. */
169export function displayToUnicode(tex: string): string {
170  const body = tex.replace(/\\(?:begin|end)\{(?:aligned|gathered|split|align\*?|gather\*?)\}/g, '')
171  const lines = body
172    .split(/\\\\(?:\[[^\]]*\])?/)
173    .map(line => texToUnicode(line).trim())
174    .filter(line => line !== '')
175  return '\n\n' + lines.map(line => '    ' + escapeMarkdown(line)).join('  \n') + '\n\n'
176}
177
178function splitFences(text: string): { text: string; isCode: boolean }[] {
179  const blocks: { text: string; isCode: boolean }[] = []
180  let current = ''
181  let fence: string | undefined
182  for (const line of text.split(/(?<=\n)/)) {
183    const marks = FENCE.exec(line)?.[1]
184    if (fence === undefined && marks) {
185      if (current) blocks.push({ text: current, isCode: false })
186      current = line
187      fence = marks
188    } else if (fence !== undefined && marks && marks[0] === fence[0] && marks.length >= fence.length && line.trim() === marks) {
189      blocks.push({ text: current + line, isCode: true })
190      current = ''
191      fence = undefined
192    } else {
193      current += line
194    }
195  }
196  if (current) blocks.push({ text: current, isCode: fence !== undefined })
197  return blocks
198}
199
200function maskCodeSpans(text: string): { masked: string; spans: string[] } {
201  const spans: string[] = []
202  const masked = text.replace(/(`+)[\s\S]*?\1/g, span => {
203    spans.push(span)
204    return HOLE + (spans.length - 1) + HOLE
205  })
206  return { masked, spans }
207}
208
209function unmask(text: string, spans: string[]): string {
210  return text.replace(/(\d+)/g, (_, i: string) => spans[Number(i)] ?? '')
211}
212
213export function escapeMarkdown(text: string): string {
214  return text.replace(/[\\`*_[\]<>|~]/g, c => '\\' + c)
215}
216
217// ---------------------------------------------------------------------------
218// LaTeX to Unicode
219
220const SYMBOLS: Record<string, string> = {
221  alpha: 'α', beta: 'β', gamma: 'γ', delta: 'δ', epsilon: 'ϵ', varepsilon: 'ε', zeta: 'ζ', eta: 'η',
222  theta: 'θ', vartheta: 'ϑ', iota: 'ι', kappa: 'κ', lambda: 'λ', mu: 'μ', nu: 'ν', xi: 'ξ', pi: 'π',
223  varpi: 'ϖ', rho: 'ρ', varrho: 'ϱ', sigma: 'σ', varsigma: 'ς', tau: 'τ', upsilon: 'υ', phi: 'ϕ',
224  varphi: 'φ', chi: 'χ', psi: 'ψ', omega: 'ω',
225  Gamma: 'Γ', Delta: 'Δ', Theta: 'Θ', Lambda: 'Λ', Xi: 'Ξ', Pi: 'Π', Sigma: 'Σ', Upsilon: 'Υ',
226  Phi: 'Φ', Psi: 'Ψ', Omega: 'Ω',
227  cdot: '·', cdots: '⋯', ldots: '…', dots: '…', vdots: '⋮', ddots: '⋱', times: '×', div: '÷',
228  pm: '±', mp: '∓', ast: '∗', star: '⋆', circ: '∘', bullet: '•', oplus: '⊕', otimes: '⊗',
229  leq: '≤', le: '≤', geq: '≥', ge: '≥', neq: '≠', ne: '≠', approx: '≈', equiv: '≡', sim: '∼',
230  simeq: '≃', cong: '≅', propto: '∝', ll: '≪', gg: '≫', prec: '≺', succ: '≻', perp: '⊥',
231  parallel: '∥', mid: '∣', nmid: '∤', vert: '|', Vert: '‖', lvert: '|', rvert: '|', lVert: '‖',
232  rVert: '‖', langle: '⟨', rangle: '⟩', lfloor: '⌊', rfloor: '⌋', lceil: '⌈', rceil: '⌉',
233  lbrace: '{', rbrace: '}', '{': '{', '}': '}', '|': '‖', '%': '%', '$': '$', '&': '&', '#': '#', _: '_',
234  in: '∈', notin: '∉', ni: '∋', subset: '⊂', supset: '⊃', subseteq: '⊆', supseteq: '⊇',
235  subsetneq: '⊊', cup: '∪', cap: '∩', bigcup: '⋃', bigcap: '⋂', setminus: '∖', emptyset: '∅',
236  varnothing: '∅', forall: '∀', exists: '∃', nexists: '∄', neg: '¬', lnot: '¬', land: '∧',
237  wedge: '∧', lor: '∨', vee: '∨', top: '⊤', bot: '⊥', vdash: '⊢', models: '⊨',
238  to: '→', rightarrow: '→', leftarrow: '←', gets: '←', leftrightarrow: '↔', Rightarrow: '⇒',
239  Leftarrow: '⇐', Leftrightarrow: '⇔', implies: '⟹', impliedby: '⟸', iff: '⟺', mapsto: '↦',
240  longrightarrow: '⟶', longleftarrow: '⟵', longleftrightarrow: '⟷', Longrightarrow: '⟹',
241  Longleftarrow: '⟸', Longleftrightarrow: '⟺', longmapsto: '⟼', uparrow: '↑', downarrow: '↓',
242  hookrightarrow: '↪', rightharpoonup: '⇀', nearrow: '↗', searrow: '↘',
243  infty: '∞', partial: '∂', nabla: '∇', sum: '∑', prod: '∏', coprod: '∐', int: '∫', iint: '∬',
244  iiint: '∭', oint: '∮', angle: '∠', triangle: '△', square: '□', Box: '□',
245  ell: 'ℓ', hbar: 'ℏ', Re: 'ℜ', Im: 'ℑ', aleph: 'ℵ', wp: '℘', prime: '′', degree: '°',
246  qed: '∎', blacksquare: '■', checkmark: '✓', dagger: '†', ddagger: '‡',
247  colon: ':', ',': ' ', ';': ' ', ':': ' ', '!': '', ' ': ' ', quad: '  ', qquad: '    ',
248  enspace: ' ', thinspace: ' ', medspace: ' ', thickspace: ' ', cr: '; ', newline: '; ',
249}
250
251const FUNCTIONS = new Set([
252  'sin', 'cos', 'tan', 'cot', 'sec', 'csc', 'arcsin', 'arccos', 'arctan', 'sinh', 'cosh', 'tanh',
253  'coth', 'log', 'ln', 'lg', 'exp', 'lim', 'liminf', 'limsup', 'max', 'min', 'sup', 'inf', 'det',
254  'dim', 'ker', 'deg', 'gcd', 'lcm', 'arg', 'Pr', 'hom', 'mod', 'bmod', 'tr', 'rank', 'span',
255])
256
257const IGNORED = new Set([
258  'left', 'right', 'big', 'Big', 'bigg', 'Bigg', 'bigl', 'bigr', 'Bigl', 'Bigr', 'biggl', 'biggr',
259  'Biggl', 'Biggr', 'middle', 'displaystyle', 'textstyle', 'scriptstyle', 'limits', 'nolimits',
260  'label', 'nonumber', 'notag', 'tag', 'mathstrut', 'strut', 'phantom', 'hphantom', 'vphantom',
261])
262
263const TEXT_COMMANDS = new Set([
264  'text', 'textrm', 'textit', 'textbf', 'textsf', 'texttt', 'mathrm', 'mathit', 'mathbf', 'mathsf',
265  'mathtt', 'boldsymbol', 'bm', 'mbox', 'hbox', 'emph', 'pmb', 'mathnormal',
266])
267
268const ACCENTS: Record<string, string> = {
269  hat: '̂', widehat: '̂', bar: '̅', overline: '̅', tilde: '̃',
270  widetilde: '̃', dot: '̇', ddot: '̈', vec: '⃗', overrightarrow: '⃗',
271  check: '̌', acute: '́', grave: '̀', breve: '̆', underline: '̲',
272}
273
274const DOUBLE_STRUCK: Record<string, string> = {
275  C: 'ℂ', H: 'ℍ', N: 'ℕ', P: 'ℙ', Q: 'ℚ', R: 'ℝ', Z: 'ℤ', '1': '𝟙',
276}
277
278const SCRIPT: Record<string, string> = {
279  B: 'ℬ', E: 'ℰ', F: 'ℱ', H: 'ℋ', I: 'ℐ', L: 'ℒ', M: 'ℳ', R: 'ℛ', e: 'ℯ', g: 'ℊ', o: 'ℴ',
280}
281
282const SUPERSCRIPT: Record<string, string> = {
283  '0': '⁰', '1': '¹', '2': '²', '3': '³', '4': '⁴', '5': '⁵', '6': '⁶', '7': '⁷', '8': '⁸', '9': '⁹',
284  '+': '⁺', '-': '⁻', '−': '⁻', '=': '⁼', '(': '⁽', ')': '⁾', a: 'ᵃ', b: 'ᵇ', c: 'ᶜ', d: 'ᵈ',
285  e: 'ᵉ', f: 'ᶠ', g: 'ᵍ', h: 'ʰ', i: 'ⁱ', j: 'ʲ', k: 'ᵏ', l: 'ˡ', m: 'ᵐ', n: 'ⁿ', o: 'ᵒ', p: 'ᵖ',
286  r: 'ʳ', s: 'ˢ', t: 'ᵗ', u: 'ᵘ', v: 'ᵛ', w: 'ʷ', x: 'ˣ', y: 'ʸ', z: 'ᶻ', A: 'ᴬ', B: 'ᴮ', D: 'ᴰ',
287  E: 'ᴱ', G: 'ᴳ', H: 'ᴴ', I: 'ᴵ', J: 'ᴶ', K: 'ᴷ', L: 'ᴸ', M: 'ᴹ', N: 'ᴺ', O: 'ᴼ', P: 'ᴾ', R: 'ᴿ',
288  T: 'ᵀ', U: 'ᵁ', V: 'ⱽ', W: 'ᵂ', 'β': 'ᵝ', 'γ': 'ᵞ', 'δ': 'ᵟ', 'θ': 'ᶿ', 'ϕ': 'ᵠ', 'φ': 'ᵠ',
289  'χ': 'ᵡ', '∗': '*', '*': '*', '′': '′', ',': '˒', ' ': '', '∘': '°',
290}
291
292const SUBSCRIPT: Record<string, string> = {
293  '0': '₀', '1': '₁', '2': '₂', '3': '₃', '4': '₄', '5': '₅', '6': '₆', '7': '₇', '8': '₈', '9': '₉',
294  '+': '₊', '-': '₋', '−': '₋', '=': '₌', '(': '₍', ')': '₎', a: 'ₐ', e: 'ₑ', h: 'ₕ', i: 'ᵢ', j: 'ⱼ',
295  k: 'ₖ', l: 'ₗ', m: 'ₘ', n: 'ₙ', o: 'ₒ', p: 'ₚ', r: 'ᵣ', s: 'ₛ', t: 'ₜ', u: 'ᵤ', v: 'ᵥ', x: 'ₓ',
296  'β': 'ᵦ', 'γ': 'ᵧ', 'ρ': 'ᵨ', 'ϕ': 'ᵩ', 'φ': 'ᵩ', 'χ': 'ᵪ', ',': ',', ' ': '',
297}
298
299// The letters Unicode has a subscript for, as the style guide lists them.
300const SUBSCRIPT_LETTERS = Object.keys(SUBSCRIPT).filter(k => /^[a-z]$/.test(k)).join(' ')
301
302const MATRIX_ENVS: Record<string, [string, string]> = {
303  matrix: ['', ''], pmatrix: ['(', ')'], bmatrix: ['[', ']'], Bmatrix: ['{', '}'],
304  vmatrix: ['|', '|'], Vmatrix: ['‖', '‖'], smallmatrix: ['', ''], cases: ['{ ', ''],
305  array: ['', ''],
306}
307
308// Mark a function name's ends: a space where a letter or digit touches it.
309const FUNCTION_START = '\uE002'
310const FUNCTION_END = '\uE001'
311
312// Relations and binary operators drawn with a space either side.
313const SPACED = /[+=<>≤≥≠≈≡∼≃≅∝∈∉∋⊂⊃⊆⊇⊊→←↔⇒⇐⇔⟹⟸⟺↦⟶⟵⟷⊢⊨≪≫·×±∓÷]/g
314
315// A sign with nothing to add to: first, or after a bracket, a comma or another operator.
316const SIGNS = '+±∓'
317const AFTER_SIGN = /(?:^|[([{,;=<>≤≥≠≈≡∼≃≅∝∈∉∋⊂⊃⊆⊇⊊→←↔⇒⇐⇔⟹⟸⟺↦⟶⊢⊨≪≫·×±∓÷+])$/
318
319function spaceOperator(op: string, at: number, all: string): string {
320  return SIGNS.includes(op) && AFTER_SIGN.test(all.slice(0, at).trimEnd()) ? op : ` ${op} `
321}
322
323type Token = { type: 'command'; name: string } | { type: 'char'; value: string }
324
325const isChar = (t: Token | undefined, value: string) => t?.type === 'char' && t.value === value
326
327function tokenize(tex: string): Token[] {
328  const tokens: Token[] = []
329  for (let i = 0; i < tex.length; ) {
330    const c = tex[i]
331    if (c === '\\') {
332      const m = /^[a-zA-Z]+/.exec(tex.slice(i + 1))
333      if (m) {
334        tokens.push({ type: 'command', name: m[0] })
335        i += 1 + m[0].length
336        // TeX drops the spaces after a control word.
337        while (/[ \t\n]/.test(tex[i] ?? '')) i++
338      } else {
339        tokens.push({ type: 'command', name: tex[i + 1] ?? '' })
340        i += 2
341      }
342    } else if (c === '%') {
343      const end = tex.indexOf('\n', i)
344      i = end === -1 ? tex.length : end + 1
345    } else {
346      const ch = String.fromCodePoint(tex.codePointAt(i) ?? 0)
347      tokens.push({ type: 'char', value: ch })
348      i += ch.length
349    }
350  }
351  return tokens
352}
353
354/** Converts a LaTeX math expression to readable Unicode text. */
355export function texToUnicode(tex: string): string {
356  const tokens = tokenize(tex)
357  let i = 0
358
359  const skipSpace = () => {
360    for (let t = tokens[i]; t?.type === 'char' && /\s/.test(t.value); t = tokens[i]) i++
361  }
362
363  // One argument: a braced group or a single token.
364  const arg = (): string => {
365    skipSpace()
366    const t = tokens[i]
367    if (t === undefined) return ''
368    if (t.type === 'char' && t.value === '{') {
369      i++
370      return expr('}')
371    }
372    i++
373    return atom(t)
374  }
375
376  const optionalArg = (): string | undefined => {
377    skipSpace()
378    const t = tokens[i]
379    if (t?.type !== 'char' || t.value !== '[') return undefined
380    i++
381    return expr(']')
382  }
383
384  const rawGroup = (): string => {
385    skipSpace()
386    if (!isChar(tokens[i], '{')) return ''
387    i++
388    let depth = 1
389    let s = ''
390    for (let t = tokens[i]; t !== undefined; t = tokens[i]) {
391      i++
392      if (t.type === 'char' && t.value === '{') depth++
393      if (t.type === 'char' && t.value === '}' && --depth === 0) break
394      s += t.type === 'char' ? t.value : '\\' + t.name
395    }
396    return s
397  }
398
399  const environment = (name: string): string => {
400    if (name === 'array') rawGroup()
401    let row: string[] = []
402    const rows: string[][] = [row]
403    let cell = ''
404    for (let t = tokens[i]; t !== undefined; t = tokens[i]) {
405      if (t.type === 'command' && t.name === 'end') {
406        i++
407        rawGroup()
408        break
409      }
410      if (t.type === 'char' && t.value === '&') {
411        i++
412        row.push(cell.trim())
413        cell = ''
414      } else if (t.type === 'command' && t.name === '\\') {
415        i++
416        optionalArg()
417        row.push(cell.trim())
418        row = []
419        rows.push(row)
420        cell = ''
421      } else {
422        i++
423        cell += atom(t)
424      }
425    }
426    row.push(cell.trim())
427    const kept = rows.filter(r => r.some(c => c !== ''))
428    const [open, close] = (Object.hasOwn(MATRIX_ENVS, name) ? MATRIX_ENVS[name] : undefined) ?? ['', '']
429    if (name === 'cases') return open + kept.map(r => r.filter(c => c !== '').join(', ')).join('; ')
430    const sep = /matrix|array/.test(name) ? ' ' : ''
431    return open + kept.map(r => r.join(sep)).join('; ') + close
432  }
433
434  const command = (name: string): string => {
435    const symbol = Object.hasOwn(SYMBOLS, name) ? SYMBOLS[name] : undefined
436    if (symbol !== undefined && !IGNORED.has(name)) return symbol
437    if (FUNCTIONS.has(name)) return FUNCTION_START + name + FUNCTION_END
438    if (name === 'operatorname') return FUNCTION_START + arg() + FUNCTION_END
439    if (IGNORED.has(name)) {
440      if (['label', 'tag', 'phantom', 'hphantom', 'vphantom'].includes(name)) rawGroup()
441      if (name === 'left' || name === 'right') {
442        skipSpace()
443        const t = tokens[i]
444        if (t?.type === 'char' && t.value === '.') i++
445      }
446      return ''
447    }
448    const mark = Object.hasOwn(ACCENTS, name) ? ACCENTS[name] : undefined
449    if (mark !== undefined) {
450      const body = arg()
451      if (name === 'vec' || name === 'overrightarrow' || name.startsWith('wide')) return body + mark
452      return [...body].map(ch => ch + mark).join('')
453    }
454    if (TEXT_COMMANDS.has(name)) return arg()
455    switch (name) {
456      case 'frac':
457      case 'dfrac':
458      case 'tfrac':
459      case 'cfrac': {
460        const a = arg()
461        const b = arg()
462        return wrap(a) + '/' + wrap(b)
463      }
464      case 'binom':
465      case 'dbinom':
466      case 'tbinom':
467        return 'C(' + arg() + ', ' + arg() + ')'
468      case 'sqrt': {
469        const n = optionalArg()
470        const body = arg()
471        const root = n === '3' ? '∛' : n === '4' ? '∜' : n ? superscript(n) + '√' : '√'
472        return root + wrap(body)
473      }
474      case 'mathbb':
475        return [...arg()].map(ch => DOUBLE_STRUCK[ch] ?? ch).join('')
476      case 'mathcal':
477      case 'mathscr':
478        return [...arg()].map(ch => SCRIPT[ch] ?? ch).join('')
479      case 'mathfrak':
480        return arg()
481      case 'not': {
482        const next = arg()
483        return next === '=' ? '≠' : next === '∈' ? '∉' : next + '̸'
484      }
485      case 'begin': {
486        const env = rawGroup().replace(/\*$/, '')
487        return environment(env)
488      }
489      case 'end':
490        rawGroup()
491        return ''
492      case 'overset':
493      case 'stackrel': {
494        arg()
495        return arg()
496      }
497      case 'underbrace':
498      case 'overbrace':
499        return arg()
500      // Color has no Unicode form: drop it, keep what it colors.
501      case 'color':
502        optionalArg()
503        rawGroup()
504        return ''
505      case 'textcolor':
506      case 'colorbox':
507        optionalArg()
508        rawGroup()
509        return arg()
510      case 'fcolorbox':
511        optionalArg()
512        rawGroup()
513        rawGroup()
514        return arg()
515      case 'boxed':
516        return arg()
517      case 'pmod':
518        return ' (mod ' + arg() + ')'
519      case '\\':
520        return '; '
521      default:
522        return name
523    }
524  }
525
526  const atom = (t: Token): string => {
527    if (t.type === 'command') return command(t.name)
528    switch (t.value) {
529      case '^':
530        return script(arg(), superscript, '^')
531      case '_':
532        return script(arg(), subscript, '_')
533      case "'":
534        return '′'
535      case '~':
536        return ' '
537      case '&':
538        return ' '
539      case '{':
540        return expr('}')
541      default:
542        return /\s/.test(t.value) ? ' ' : t.value
543    }
544  }
545
546  const expr = (until?: string): string => {
547    let s = ''
548    for (let t = tokens[i]; t !== undefined; t = tokens[i]) {
549      if (until && isChar(t, until)) {
550        i++
551        break
552      }
553      i++
554      s += atom(t)
555    }
556    return s
557  }
558
559  return expr()
560    .replace(new RegExp(FUNCTION_END + '(?=[A-Za-z0-9\u0370-\u03ff])', 'g'), ' ')
561    .replaceAll(FUNCTION_END, '')
562    .replace(new RegExp('(?<=[A-Za-z0-9)\u0370-\u03ff′])' + FUNCTION_START, 'g'), ' ')
563    .replaceAll(FUNCTION_START, '')
564    .replace(SPACED, spaceOperator)
565    .replace(/[ \t\n]+/g, ' ')
566    .replace(/ ([,;.)\]])/g, '$1')
567    .replace(/([([]) /g, '$1')
568    .trim()
569}
570
571function script(body: string, convert: (s: string) => string | undefined, mark: string): string {
572  const trimmed = body.trim()
573  if (/^′+$/.test(trimmed)) return trimmed
574  if (trimmed === '') return ''
575  const plain = [...trimmed].length === 1 || /^\(.*\)$/.test(trimmed) ? trimmed : '(' + trimmed + ')'
576  return convert(trimmed) ?? mark + plain
577}
578
579function superscript(s: string): string | undefined {
580  return mapAll(s, SUPERSCRIPT)
581}
582
583function subscript(s: string): string | undefined {
584  return mapAll(s, SUBSCRIPT)
585}
586
587function mapAll(s: string, table: Record<string, string>): string | undefined {
588  let out = ''
589  for (const ch of s) {
590    const mapped = table[ch]
591    if (mapped === undefined) return undefined
592    out += mapped
593  }
594  return out
595}
596
597function wrap(s: string): string {
598  const t = s.trim()
599  if ([...t].length <= 1 || /^[\p{L}\p{N}.′]+$/u.test(t) || /^\(.*\)$/.test(t)) return t
600  return '(' + t + ')'
601}
602
603// ---------------------------------------------------------------------------
604// Typesetting a display formula
605
606const STARRED: Record<string, string> = {
607  equation: 'gather*', 'equation*': 'gather*', multline: 'gather*', 'multline*': 'gather*',
608  flalign: 'align*', 'flalign*': 'align*', align: 'align*', gather: 'gather*', alignat: 'alignat*',
609  eqnarray: 'eqnarray*',
610}
611
612/**
613 * The LaTeX document a display formula is typeset alone in, drawn in `ink`:
614 * xcolor sets the page in explicit `black`, which would override dvipng's
615 * `-fg`, so `black` is the ink. With `paper`, a dark theme's background, named
616 * colors too dim to read on it are lightened. Both are 6-digit hex.
617 */
618export function texDocument(segment: { tex: string; env?: string }, preamble: string, ink?: string, paper?: string): string {
619  const body = segment.tex.replace(/\\label\{[^}]*\}/g, '')
620  let env = segment.env ? (STARRED[segment.env] ?? segment.env) : undefined
621  if (env === undefined) env = body.includes('&') && !body.includes('\\begin{') ? 'align*' : 'gather*'
622  if (env === 'alignat*' && !body.trimStart().startsWith('{')) env = 'align*'
623  return [
624    '\\documentclass[10pt]{article}',
625    '\\usepackage{amsmath,amssymb,amsfonts,mathtools,bm,xcolor}',
626    ...figureColors(Boolean(paper)),
627    preamble,
628    ...(ink ? [`\\definecolor{black}{HTML}{${ink.replace('#', '')}}`] : []),
629    '\\pagestyle{empty}',
630    '\\setlength{\\textwidth}{60cm}',
631    '\\begin{document}',
632    ...(ink ? ['\\color{black}'] : []),
633    `\\begin{${env}}`,
634    body,
635    `\\end{${env}}`,
636    '\\end{document}',
637    '',
638  ].join('\n')
639}
640
641/**
642 * The standalone document a TikZ figure is typeset alone in, drawn in `ink`.
643 * With `paper`, the background of a dark theme, `white` (and what colors are
644 * mixed with, as in `blue!20`) becomes that background, so fills, masks and
645 * tints stay dark; `black` is always the ink. Both are 6-digit hex.
646 */
647export function figureDocument(tex: string, preamble: string, ink: string, paper?: string): string {
648  const body = /\\begin\{(?:tikzpicture|tikzcd)\}/.test(tex)
649    ? tex
650    : `\\begin{tikzpicture}\n${tex}\n\\end{tikzpicture}`
651  const dark = paper ? hexRgb(paper) : undefined
652  return [
653    '\\documentclass[tikz,border=4pt]{standalone}',
654    '\\usepackage{amsmath,amssymb,amsfonts,mathtools,bm,pgfplots}',
655    '\\pgfplotsset{compat=1.18,every axis/.append style={axis on top}}',
656    '\\usetikzlibrary{arrows.meta,calc,cd,decorations.pathreplacing,patterns,positioning,shapes.geometric}',
657    ...figureColors(Boolean(paper)),
658    preamble,
659    `\\definecolor{black}{HTML}{${ink.replace('#', '')}}`,
660    ...(dark
661      ? [
662          `\\definecolor{white}{HTML}{${paper!.replace('#', '')}}`,
663          '\\makeatletter',
664          `\\def\\XC@clr@rgb@white{${dark.join(',')}}`,
665          `\\def\\XC@clr@gray@white{${((dark[0] + dark[1] + dark[2]) / 3).toFixed(3)}}`,
666          `\\def\\XC@clr@cmy@white{${dark.map(c => (1 - c).toFixed(3)).join(',')}}`,
667          `\\def\\XC@clr@cmyk@white{${dark.map(c => (1 - c).toFixed(3)).join(',')},0}`,
668          '\\makeatother',
669        ]
670      : []),
671    '\\begin{document}',
672    '\\color{black}',
673    body,
674    '\\end{document}',
675    '',
676  ].join('\n')
677}
678
679// Lowercase SVG names xcolor lacks, defined so a figure naming one compiles.
680const EXTRA_COLORS: Record<string, string> = {
681  navy: '000080',
682  darkblue: '00008B',
683  midnightblue: '191970',
684  indigo: '4B0082',
685  maroon: '800000',
686  darkred: '8B0000',
687  darkgreen: '006400',
688}
689
690// On a dark theme, the named colors too dim to read on the background, as
691// lighter shades of themselves.
692const DARK_COLORS: Record<string, string> = {
693  blue: '6EA8FF',
694  navy: '7AA2F7',
695  darkblue: '6C8CFF',
696  midnightblue: '8C9EFF',
697  violet: 'B48CFF',
698  indigo: 'A58CFF',
699  purple: 'D38BD8',
700  brown: 'D09A66',
701  maroon: 'E07A7A',
702  darkred: 'FF6B6B',
703  darkgreen: '5CC96B',
704  teal: '3CC8C8',
705  darkgray: 'A8A8A8',
706}
707
708// Ahead of the extra preamble, so a color the person defines there wins.
709function figureColors(dark: boolean): string[] {
710  const colors = dark ? { ...EXTRA_COLORS, ...DARK_COLORS } : EXTRA_COLORS
711  return Object.entries(colors).map(([name, hex]) => `\\definecolor{${name}}{HTML}{${hex}}`)
712}
713
714/** What a reply drawn here can show: pictures at all, figures, and a dark theme's background. */
715export type StyleFacts = { pictures: boolean; figures: boolean; paper?: string }
716
717/**
718 * The system prompt section that tells the model how to write math and
719 * figures so this mod draws them well.
720 */
721export function styleGuide({ pictures, figures, paper }: StyleFacts): string {
722  const lines = [
723    '# Math and figures',
724    'Your replies are shown in a terminal that renders LaTeX (the latex-render plugin). Write math and figures in the forms it renders:',
725    '- Write math in LaTeX rather than plain-text approximations like x^2 or sqrt(x): inline `$…$`, display `$$…$$` on lines of its own (or an `align`, `gather` or `equation` environment).',
726    '- Inline math is turned into Unicode text, so keep it to symbols, Greek letters, sub- and superscripts, operators and short fractions like `$a/b$`. Anything taller (stacked fractions, matrices, `cases`, sums or integrals with limits) goes in display math.',
727    `- Inline subscripts become Unicode only for digits, signs and the letters ${SUBSCRIPT_LETTERS}. Any other, like the N in \`$S_N$\`, is left as \`_N\`. When the notation is yours to choose, pick letters that have one (\`$S_n$\`, \`$x_k$\`); when it is fixed, put a formula that should look typeset in display math.`,
728    '- A bare number is plain text: write 0 or 3.5 rather than `$0$` or `$3.5$`.',
729    '- Never put math meant to be read in a code block or inline code, which show raw LaTeX. Use a ```latex block only when asked for LaTeX source.',
730    '- Figures, graphs, plots and commutative diagrams (including pgfplots `axis` plots) always go in a ```tikz fenced block, never loose in the text and never in a ```latex block: bare TikZ commands, or a whole `tikzpicture`, `axis` or `tikzcd` environment. pgfplots, tikz-cd and the TikZ libraries arrows.meta, calc, cd, decorations.pathreplacing, patterns, positioning and shapes.geometric are loaded; write no `\\documentclass`, `\\usepackage` or `\\begin{document}`.',
731    '- A formula or figure using `\\def`, `\\newcommand`, `\\let`, `\\input`, `\\write`, `\\csname` or other macro and file primitives is never compiled; use `\\pgfmathsetmacro`, `\\tikzset` or plain repetition instead.',
732  ]
733  if (!pictures) {
734    lines.push(
735      '- This terminal shows no images: display math is shown as Unicode text and figures as their TikZ code. Keep display math simple, and draw a figure only when asked.',
736    )
737  } else if (!figures) {
738    lines.push('- Figures cannot be drawn here and are shown as their TikZ code; draw one only when asked.')
739  }
740  if (pictures && figures && paper) {
741    lines.push(
742      `- The background is dark (#${paper.replace('#', '')}). In figures, leave lines and text in the default color, the terminal's light text color (\`black\` means the same). \`white\` is the background: a white fill hides what is under it, and tints like \`blue!20\` are dim fills, good for shading regions.`,
743      '- Color figures with named colors: red, orange, yellow, green, cyan, blue, magenta, violet, purple, teal, brown, gray. Dark ones are lightened to read on the background. Do not define colors by RGB or hex, and do not use dvipsnames or svgnames (`NavyBlue`, `Navy`): they are not loaded.',
744    )
745  } else if (pictures && figures) {
746    lines.push('- Color figures with named xcolor colors (red, blue, green, orange, …); dvipsnames and svgnames (`NavyBlue`, `Navy`) are not loaded.')
747  }
748  return lines.join('\n')
749}
750
751function hexRgb(hex: string): [number, number, number] | undefined {
752  const m = /^#?([0-9a-f]{2})([0-9a-f]{2})([0-9a-f]{2})$/i.exec(hex.trim())
753  if (!m) return undefined
754  return [1, 2, 3].map(k => Number((parseInt(m[k]!, 16) / 255).toFixed(3))) as [number, number, number]
755}
756
757/** A short stable hex key for a string (two FNV-1a passes). */
758export function hashKey(s: string): string {
759  let a = 0x811c9dc5
760  let b = 0x01000193
761  for (let i = 0; i < s.length; i++) {
762    const c = s.charCodeAt(i)
763    a = Math.imul(a ^ c, 0x01000193)
764    b = Math.imul(b ^ c, 0x5bd1e995) ^ (b >>> 13)
765  }
766  return (a >>> 0).toString(16).padStart(8, '0') + (b >>> 0).toString(16).padStart(8, '0')
767}
768
769/** Width and height of a PNG from the start of its base64. */
770export function pngSize(base64: string): { width: number; height: number } | undefined {
771  const head = atob(base64.slice(0, 32))
772  if (head.length < 24 || head.slice(1, 4) !== 'PNG') return undefined
773  const u32 = (at: number) =>
774    ((head.charCodeAt(at) << 24) | (head.charCodeAt(at + 1) << 16) | (head.charCodeAt(at + 2) << 8) | head.charCodeAt(at + 3)) >>> 0
775  return { width: u32(16), height: u32(20) }
776}
777
778/** `#rrggbb` as dvipng's `rgb r g b`. */
779export function dvipngColor(hex: string): string | undefined {
780  const m = /^#?([0-9a-f]{2})([0-9a-f]{2})([0-9a-f]{2})$/i.exec(hex.trim())
781  if (!m) return undefined
782  return 'rgb ' + m.slice(1, 4).map(h => (parseInt(h, 16) / 255).toFixed(3)).join(' ')
783}
784
785/** The first error line of a LaTeX log. */
786export function latexError(log: string): string {
787  const line = log.split('\n').find(l => l.startsWith('! '))
788  return line ? line.slice(2).trim() : 'LaTeX failed'
789}
790
791const UNSAFE = /\\(?:input|include|includeonly|InputIfFileExists|openin|openout|closein|closeout|read|readline|write|message|immediate|special|catcode|uccode|lccode|csname|endcsname|string|def|edef|gdef|xdef|let|futurelet|expandafter|noexpand|newcommand|renewcommand|providecommand|DeclareRobustCommand|usepackage|RequirePackage|documentclass|begin\s*\{\s*document|end\s*\{\s*document|directlua|latelua|pdffilesize|filesize|verbatiminput|lstinputlisting|includegraphics|newread|newwrite|jobname|loop|repeat)(?![a-zA-Z])|\^\^/
792
793/** Whether a formula is plain math: no file access, no macro or catcode tricks. */
794export function isSafeTex(tex: string): boolean {
795  return !UNSAFE.test(tex)
796}
797

1/** How replies are drawn: LaTeX typeset as pictures, or left as written. */
2export type Mode = 'render' | 'source'
3
4/** A formula or figure typeset while its reply still streams, shown above the prompt. */
5export type Previewed = { kind: 'display'; tex: string; env?: string } | { kind: 'figure'; tex: string }
6
7declare module 'claude-code' {
8  interface PluginState {
9    'latex-render': { mode: Mode; preview: Previewed[] }
10  }
11}
12