Start adding client-side word count (#24892)

GitOrigin-RevId: 6c17d7bf7095794c003e17939a8302fc6b059262

services/web/app/src/Features/Project/ProjectController.js

@@ -355,6 +355,7 @@ const _ProjectController = {
       'paywall-change-compile-timeout',
       'overleaf-assist-bundle',
       'wf-feature-rebrand',
+      'word-count-client',
     ].filter(Boolean)
 
     const getUserValues = async userId =>

services/web/frontend/js/features/editor-left-menu/components/actions-word-count.tsx

@@ -1,15 +1,10 @@
 import { useState, useCallback } from 'react'
-import { useTranslation } from 'react-i18next'
-import { useDetachCompileContext as useCompileContext } from '../../../shared/context/detach-compile-context'
 import WordCountModal from '../../word-count-modal/components/word-count-modal'
-import LeftMenuButton from './left-menu-button'
 import * as eventTracking from '../../../infrastructure/event-tracking'
-import OLTooltip from '@/features/ui/components/ol/ol-tooltip'
+import { WordCountButton } from '@/features/word-count-modal/components/word-count-button'
 
 export default function ActionsWordCount() {
   const [showModal, setShowModal] = useState(false)
-  const { pdfUrl } = useCompileContext()
-  const { t } = useTranslation()
 
   const handleShowModal = useCallback(() => {
     eventTracking.sendMB('left-menu-count')
@@ -18,32 +13,7 @@ export default function ActionsWordCount() {
 
   return (
     <>
-      {pdfUrl ? (
-        <LeftMenuButton onClick={handleShowModal} icon="match_case">
-          {t('word_count')}
-        </LeftMenuButton>
-      ) : (
-        <OLTooltip
-          id="disabled-word-count"
-          description={t('please_compile_pdf_before_word_count')}
-          overlayProps={{
-            placement: 'top',
-          }}
-        >
-          {/* OverlayTrigger won't fire unless the child is a non-react html element (e.g div, span) */}
-          <div>
-            <LeftMenuButton
-              icon="match_case"
-              disabled
-              disabledAccesibilityText={t(
-                'please_compile_pdf_before_word_count'
-              )}
-            >
-              {t('word_count')}
-            </LeftMenuButton>
-          </div>
-        </OLTooltip>
-      )}
+      <WordCountButton handleShowModal={handleShowModal} />
       <WordCountModal show={showModal} handleHide={() => setShowModal(false)} />
     </>
   )

services/web/frontend/js/features/word-count-modal/components/word-count-button.tsx

@@ -0,0 +1,44 @@
+import { useDetachCompileContext as useCompileContext } from '@/shared/context/detach-compile-context'
+import { useTranslation } from 'react-i18next'
+import { isSplitTestEnabled } from '@/utils/splitTestUtils'
+import LeftMenuButton from '@/features/editor-left-menu/components/left-menu-button'
+import OLTooltip from '@/features/ui/components/ol/ol-tooltip'
+import { memo } from 'react'
+
+export const WordCountButton = memo<{
+  handleShowModal: () => void
+}>(function WordCountButton({ handleShowModal }) {
+  const { pdfUrl } = useCompileContext()
+  const { t } = useTranslation()
+
+  const enabled = pdfUrl || isSplitTestEnabled('word-count-client')
+
+  if (!enabled) {
+    return (
+      <OLTooltip
+        id="disabled-word-count"
+        description={t('please_compile_pdf_before_word_count')}
+        overlayProps={{
+          placement: 'top',
+        }}
+      >
+        {/* OverlayTrigger won't fire unless the child is a non-react html element (e.g div, span) */}
+        <div>
+          <LeftMenuButton
+            icon="match_case"
+            disabled
+            disabledAccesibilityText={t('please_compile_pdf_before_word_count')}
+          >
+            {t('word_count')}
+          </LeftMenuButton>
+        </div>
+      </OLTooltip>
+    )
+  }
+
+  return (
+    <LeftMenuButton onClick={handleShowModal} icon="match_case">
+      {t('word_count')}
+    </LeftMenuButton>
+  )
+})

services/web/frontend/js/features/word-count-modal/components/word-count-client.tsx

@@ -0,0 +1,105 @@
+import { FC, useEffect, useMemo, useState } from 'react'
+import { WordCountData } from '@/features/word-count-modal/components/word-count-data'
+import { WordCountLoading } from '@/features/word-count-modal/components/word-count-loading'
+import { WordCountError } from '@/features/word-count-modal/components/word-count-error'
+import { useProjectContext } from '@/shared/context/project-context'
+import useAbortController from '@/shared/hooks/use-abort-controller'
+import { useProjectSettingsContext } from '@/features/editor-left-menu/context/project-settings-context'
+import { useEditorManagerContext } from '@/features/ide-react/context/editor-manager-context'
+import { useFileTreePathContext } from '@/features/file-tree/contexts/file-tree-path'
+import { debugConsole } from '@/utils/debugging'
+import { signalWithTimeout } from '@/utils/abort-signal'
+import { isMainFile } from '@/features/pdf-preview/util/editor-files'
+import { countWordsInFile } from '@/features/word-count-modal/utils/count-words-in-file'
+import { WordCounts } from '@/features/word-count-modal/components/word-counts'
+import { createSegmenters } from '@/features/word-count-modal/utils/segmenters'
+
+export const WordCountClient: FC = () => {
+  const [loading, setLoading] = useState(true)
+  const [error, setError] = useState(false)
+  const [data, setData] = useState<WordCountData | null>(null)
+  const { projectSnapshot, rootDocId } = useProjectContext()
+  const { spellCheckLanguage } = useProjectSettingsContext()
+  const { openDocs, currentDocument } = useEditorManagerContext()
+  const { pathInFolder } = useFileTreePathContext()
+
+  const { signal } = useAbortController()
+
+  const segmenters = useMemo(() => {
+    return createSegmenters(spellCheckLanguage?.replace(/_/, '-'))
+  }, [spellCheckLanguage])
+
+  useEffect(() => {
+    if (currentDocument && segmenters) {
+      const countWords = async () => {
+        await openDocs.awaitBufferedOps(signalWithTimeout(signal, 5000))
+        await projectSnapshot.refresh()
+
+        if (signal.aborted) return null
+
+        const currentDocSnapshot = currentDocument.getSnapshot()
+        const currentRootDocId = isMainFile(currentDocSnapshot)
+          ? currentDocument.doc_id
+          : rootDocId
+        if (!currentRootDocId) return null
+
+        const currentRootDocPath = pathInFolder(currentRootDocId)
+        if (!currentRootDocPath) return null
+
+        const data: WordCountData = {
+          encode: 'ascii',
+          textWords: 0,
+          textCharacters: 0,
+          headWords: 0,
+          headCharacters: 0,
+          abstractWords: 0,
+          abstractCharacters: 0,
+          captionWords: 0,
+          captionCharacters: 0,
+          footnoteWords: 0,
+          footnoteCharacters: 0,
+          outside: 0,
+          outsideCharacters: 0,
+          headers: 0,
+          elements: 0,
+          mathInline: 0,
+          mathDisplay: 0,
+          errors: 0,
+          messages: '',
+        }
+
+        countWordsInFile(data, projectSnapshot, currentRootDocPath, segmenters)
+
+        return data
+      }
+
+      countWords()
+        .then(data => {
+          setData(data)
+        })
+        .catch(error => {
+          debugConsole.error(error)
+          setError(true)
+        })
+        .finally(() => {
+          setLoading(false)
+        })
+    }
+  }, [
+    signal,
+    openDocs,
+    projectSnapshot,
+    segmenters,
+    currentDocument,
+    rootDocId,
+    pathInFolder,
+  ])
+
+  return (
+    <>
+      {loading && !error && <WordCountLoading />}
+      {error && <WordCountError />}
+      {data && <WordCounts data={data} source="client" />}
+    </>
+  )
+}

services/web/frontend/js/features/word-count-modal/components/word-count-data.ts

@@ -0,0 +1,25 @@
+export type ServerWordCountData = {
+  encode: string
+  textWords: number
+  headWords: number
+  outside: number
+  headers: number
+  elements: number
+  mathInline: number
+  mathDisplay: number
+  errors: number
+  messages: string
+}
+
+export type WordCountData = ServerWordCountData & {
+  textCharacters: number
+  headCharacters: number
+  captionWords: number
+  captionCharacters: number
+  footnoteWords: number
+  footnoteCharacters: number
+  abstractWords: number
+  abstractCharacters: number
+  // outsideWords: number
+  outsideCharacters: number
+}

services/web/frontend/js/features/word-count-modal/components/word-count-error.tsx

@@ -0,0 +1,10 @@
+import OLNotification from '@/features/ui/components/ol/ol-notification'
+import { useTranslation } from 'react-i18next'
+
+export const WordCountError = () => {
+  const { t } = useTranslation()
+
+  return (
+    <OLNotification type="error" content={t('generic_something_went_wrong')} />
+  )
+}

services/web/frontend/js/features/word-count-modal/components/word-count-loading.tsx

@@ -0,0 +1,14 @@
+import { Spinner } from 'react-bootstrap-5'
+import { useTranslation } from 'react-i18next'
+
+export const WordCountLoading = () => {
+  const { t } = useTranslation()
+
+  return (
+    <div className="loading">
+      <Spinner animation="border" aria-hidden="true" size="sm" role="status" />
+      &nbsp;
+      {t('loading')}…
+    </div>
+  )
+}

services/web/frontend/js/features/word-count-modal/components/word-count-modal-content.jsx

@@ -1,109 +0,0 @@
-import PropTypes from 'prop-types'
-import { useTranslation } from 'react-i18next'
-import { useProjectContext } from '@/shared/context/project-context'
-import { useLocalCompileContext } from '@/shared/context/local-compile-context'
-import { useWordCount } from '../hooks/use-word-count'
-import {
-  OLModalBody,
-  OLModalFooter,
-  OLModalHeader,
-  OLModalTitle,
-} from '@/features/ui/components/ol/ol-modal'
-import OLNotification from '@/features/ui/components/ol/ol-notification'
-import OLRow from '@/features/ui/components/ol/ol-row'
-import OLCol from '@/features/ui/components/ol/ol-col'
-import OLButton from '@/features/ui/components/ol/ol-button'
-import { Spinner } from 'react-bootstrap-5'
-
-// NOTE: this component is only mounted when the modal is open
-export default function WordCountModalContent({ handleHide }) {
-  const { _id: projectId } = useProjectContext()
-  const { clsiServerId } = useLocalCompileContext()
-  const { t } = useTranslation()
-  const { data, error, loading } = useWordCount(projectId, clsiServerId)
-
-  return (
-    <>
-      <OLModalHeader closeButton>
-        <OLModalTitle>{t('word_count')}</OLModalTitle>
-      </OLModalHeader>
-
-      <OLModalBody>
-        {loading && !error && (
-          <div className="loading">
-            <Spinner
-              animation="border"
-              aria-hidden="true"
-              size="sm"
-              role="status"
-            />
-            &nbsp;
-            {t('loading')}…
-          </div>
-        )}
-
-        {error && (
-          <OLNotification
-            type="error"
-            content={t('generic_something_went_wrong')}
-          />
-        )}
-
-        {data && (
-          <div className="container-fluid">
-            {data.messages && (
-              <OLRow>
-                <OLCol xs={12}>
-                  <OLNotification
-                    type="error"
-                    content={
-                      <p style={{ whiteSpace: 'pre-wrap' }}>{data.messages}</p>
-                    }
-                  />
-                </OLCol>
-              </OLRow>
-            )}
-
-            <OLRow>
-              <OLCol xs={4}>
-                <div className="pull-right">{t('total_words')}:</div>
-              </OLCol>
-              <OLCol xs={6}>{data.textWords}</OLCol>
-            </OLRow>
-
-            <OLRow>
-              <OLCol xs={4}>
-                <div className="pull-right">{t('headers')}:</div>
-              </OLCol>
-              <OLCol xs={6}>{data.headers}</OLCol>
-            </OLRow>
-
-            <OLRow>
-              <OLCol xs={4}>
-                <div className="pull-right">{t('math_inline')}:</div>
-              </OLCol>
-              <OLCol xs={6}>{data.mathInline}</OLCol>
-            </OLRow>
-
-            <OLRow>
-              <OLCol xs={4}>
-                <div className="pull-right">{t('math_display')}:</div>
-              </OLCol>
-              <OLCol xs={6}>{data.mathDisplay}</OLCol>
-            </OLRow>
-          </div>
-        )}
-      </OLModalBody>
-
-      <OLModalFooter>
-        <OLButton variant="secondary" onClick={handleHide}>
-          {t('close')}
-        </OLButton>
-      </OLModalFooter>
-    </>
-  )
-}
-
-WordCountModalContent.propTypes = {
-  handleHide: PropTypes.func.isRequired,
-}

services/web/frontend/js/features/word-count-modal/components/word-count-modal-content.tsx

@@ -0,0 +1,42 @@
+import { useTranslation } from 'react-i18next'
+import {
+  OLModalBody,
+  OLModalFooter,
+  OLModalHeader,
+  OLModalTitle,
+} from '@/features/ui/components/ol/ol-modal'
+import OLButton from '@/features/ui/components/ol/ol-button'
+import { WordCountServer } from './word-count-server'
+import { WordCountClient } from './word-count-client'
+import { isSplitTestEnabled } from '@/utils/splitTestUtils'
+
+// NOTE: this component is only mounted when the modal is open
+export default function WordCountModalContent({
+  handleHide,
+}: {
+  handleHide: () => void
+}) {
+  const { t } = useTranslation()
+
+  return (
+    <>
+      <OLModalHeader closeButton>
+        <OLModalTitle>{t('word_count')}</OLModalTitle>
+      </OLModalHeader>
+
+      <OLModalBody>
+        {isSplitTestEnabled('word-count-client') ? (
+          <WordCountClient />
+        ) : (
+          <WordCountServer />
+        )}
+      </OLModalBody>
+
+      <OLModalFooter>
+        <OLButton variant="secondary" onClick={handleHide}>
+          {t('close')}
+        </OLButton>
+      </OLModalFooter>
+    </>
+  )
+}

services/web/frontend/js/features/word-count-modal/components/word-count-modal.jsx → services/web/frontend/js/features/word-count-modal/components/word-count-modal.tsx

@@ -1,12 +1,14 @@
-import React from 'react'
-import PropTypes from 'prop-types'
+import { memo } from 'react'
 import WordCountModalContent from './word-count-modal-content'
 import withErrorBoundary from '../../../infrastructure/error-boundary'
 import OLModal from '@/features/ui/components/ol/ol-modal'
 
-const WordCountModal = React.memo(function WordCountModal({
+const WordCountModal = memo(function WordCountModal({
   show,
   handleHide,
+}: {
+  show: boolean
+  handleHide: () => void
 }) {
   return (
     <OLModal animation show={show} onHide={handleHide} id="word-count-modal">
@@ -15,9 +17,4 @@ const WordCountModal = React.memo(function WordCountModal({
   )
 })
 
-WordCountModal.propTypes = {
-  show: PropTypes.bool,
-  handleHide: PropTypes.func.isRequired,
-}
-
 export default withErrorBoundary(WordCountModal)

services/web/frontend/js/features/word-count-modal/components/word-count-server.tsx

@@ -0,0 +1,48 @@
+import { FC, useEffect, useState } from 'react'
+import { ServerWordCountData } from '@/features/word-count-modal/components/word-count-data'
+import { WordCountLoading } from '@/features/word-count-modal/components/word-count-loading'
+import { WordCountError } from '@/features/word-count-modal/components/word-count-error'
+import { useProjectContext } from '@/shared/context/project-context'
+import { useLocalCompileContext } from '@/shared/context/local-compile-context'
+import useAbortController from '@/shared/hooks/use-abort-controller'
+import { getJSON } from '@/infrastructure/fetch-json'
+import { debugConsole } from '@/utils/debugging'
+import { WordCounts } from '@/features/word-count-modal/components/word-counts'
+
+export const WordCountServer: FC = () => {
+  const { _id: projectId } = useProjectContext()
+  const { clsiServerId } = useLocalCompileContext()
+
+  const [loading, setLoading] = useState(true)
+  const [error, setError] = useState(false)
+  const [data, setData] = useState<ServerWordCountData | null>(null)
+
+  const { signal } = useAbortController()
+
+  useEffect(() => {
+    const url = new URL(`/project/${projectId}/wordcount`, window.location.href)
+    if (clsiServerId) {
+      url.searchParams.set('clsiserverid', clsiServerId)
+    }
+
+    getJSON(url.toString(), { signal })
+      .then(data => {
+        setData(data.texcount)
+      })
+      .catch(error => {
+        debugConsole.error(error)
+        setError(true)
+      })
+      .finally(() => {
+        setLoading(false)
+      })
+  }, [projectId, clsiServerId, signal])
+
+  return (
+    <>
+      {loading && !error && <WordCountLoading />}
+      {error && <WordCountError />}
+      {data && <WordCounts data={data} source="server" />}
+    </>
+  )
+}

services/web/frontend/js/features/word-count-modal/components/word-counts.tsx

@@ -0,0 +1,102 @@
+import {
+  ServerWordCountData,
+  WordCountData,
+} from '@/features/word-count-modal/components/word-count-data'
+import { useTranslation } from 'react-i18next'
+import { FC } from 'react'
+import { Container, Row, Col } from 'react-bootstrap-5'
+import OLNotification from '@/features/ui/components/ol/ol-notification'
+
+export const WordCounts: FC<
+  | {
+      data: ServerWordCountData
+      source: 'server'
+    }
+  | {
+      data: WordCountData
+      source: 'client'
+    }
+> = ({ data, source }) => {
+  const { t } = useTranslation()
+
+  return (
+    <Container fluid>
+      {data.messages && (
+        <Row>
+          <Col xs={12}>
+            <OLNotification
+              type="error"
+              content={
+                <p style={{ whiteSpace: 'pre-wrap' }}>{data.messages}</p>
+              }
+            />
+          </Col>
+        </Row>
+      )}
+
+      {source === 'client' ? (
+        <>
+          <Row>
+            <Col xs={4}>
+              <div className="pull-right">Text:</div>
+            </Col>
+            <Col xs={6}>{data.textWords}</Col>
+          </Row>
+
+          <Row>
+            <Col xs={4}>
+              <div className="pull-right">Headers:</div>
+            </Col>
+            <Col xs={6}>{data.headWords}</Col>
+          </Row>
+
+          <Row>
+            <Col xs={4}>
+              <div className="pull-right">Captions:</div>
+            </Col>
+            <Col xs={6}>{data.captionWords}</Col>
+          </Row>
+
+          <Row>
+            <Col xs={4}>
+              <div className="pull-right">Footnotes:</div>
+            </Col>
+            <Col xs={6}>{data.footnoteWords}</Col>
+          </Row>
+        </>
+      ) : (
+        <Row>
+          <Col xs={4}>
+            <div className="pull-right">{t('total_words')}:</div>
+          </Col>
+          <Col xs={6}>{data.textWords}</Col>
+        </Row>
+      )}
+
+      {source === 'server' && (
+        <>
+          <Row>
+            <Col xs={4}>
+              <div className="pull-right">{t('headers')}:</div>
+            </Col>
+            <Col xs={6}>{data.headers}</Col>
+          </Row>
+
+          <Row>
+            <Col xs={4}>
+              <div className="pull-right">{t('math_inline')}:</div>
+            </Col>
+            <Col xs={6}>{data.mathInline}</Col>
+          </Row>
+
+          <Row>
+            <Col xs={4}>
+              <div className="pull-right">{t('math_display')}:</div>
+            </Col>
+            <Col xs={6}>{data.mathDisplay}</Col>
+          </Row>
+        </>
+      )}
+    </Container>
+  )
+}

services/web/frontend/js/features/word-count-modal/hooks/use-word-count.js

@@ -1,26 +0,0 @@
-import useAbortController from '../../../shared/hooks/use-abort-controller'
-import { fetchWordCount } from '../utils/api'
-import { useEffect, useState } from 'react'
-
-export function useWordCount(projectId, clsiServerId) {
-  const [loading, setLoading] = useState(true)
-  const [error, setError] = useState(false)
-  const [data, setData] = useState()
-
-  const { signal } = useAbortController()
-
-  useEffect(() => {
-    fetchWordCount(projectId, clsiServerId, { signal })
-      .then(data => {
-        setData(data.texcount)
-      })
-      .catch(() => {
-        setError(true)
-      })
-      .finally(() => {
-        setLoading(false)
-      })
-  }, [signal, clsiServerId, projectId])
-
-  return { data, error, loading }
-}

services/web/frontend/js/features/word-count-modal/utils/api.js

@@ -1,10 +0,0 @@
-import { getJSON } from '../../../infrastructure/fetch-json'
-
-export function fetchWordCount(projectId, clsiServerId, options) {
-  let query = ''
-  if (clsiServerId) {
-    query = `?clsiserverid=${clsiServerId}`
-  }
-
-  return getJSON(`/project/${projectId}/wordcount${query}`, options)
-}

services/web/frontend/js/features/word-count-modal/utils/count-words-in-file.ts

@@ -0,0 +1,256 @@
+import { ProjectSnapshot } from '@/infrastructure/project-snapshot'
+import { LaTeXLanguage } from '@/features/source-editor/languages/latex/latex-language'
+import { WordCountData } from '@/features/word-count-modal/components/word-count-data'
+import { NodeType, SyntaxNodeRef } from '@lezer/common'
+import { debugConsole } from '@/utils/debugging'
+import { findPreambleExtent } from '@/features/word-count-modal/utils/find-preamble-extent'
+import { Segmenters } from './segmenters'
+
+const whiteSpaceRe = /^\s$/
+
+type Context = 'text' | 'header' | 'abstract' | 'caption' | 'footnote'
+
+const counters: Record<
+  Context,
+  {
+    word: keyof WordCountData
+    character: keyof WordCountData
+  }
+> = {
+  text: {
+    word: 'textWords',
+    character: 'textCharacters',
+  },
+  header: {
+    word: 'headWords',
+    character: 'headCharacters',
+  },
+  abstract: {
+    word: 'abstractWords',
+    character: 'abstractCharacters',
+  },
+  caption: {
+    word: 'captionWords',
+    character: 'captionCharacters',
+  },
+  footnote: {
+    word: 'footnoteWords',
+    character: 'footnoteCharacters',
+  },
+}
+
+const replacementsMap: Map<string, string> = new Map([
+  // LaTeX commands that create part of a word
+  ['aa', 'å'],
+  ['AA', 'Å'],
+  ['ae', 'æ'],
+  ['AE', 'Æ'],
+  ['oe', 'œ'],
+  ['OE', 'Œ'],
+  ['o', 'ø'],
+  ['O', 'Ø'],
+  ['ss', 'ß'],
+  ['SS', 'SS'],
+  ['l', 'ł'],
+  ['L', 'Ł'],
+  ['dh', 'ð'],
+  ['DH', 'Ð'],
+  ['dj', 'đ'],
+  ['DJ', 'Ð'],
+  ['th', 'þ'],
+  ['TH', 'Þ'],
+  ['ng', 'ŋ'],
+  ['NG', 'Ŋ'],
+  ['i', 'ı'],
+  ['j', 'ȷ'],
+  ['_', '_'],
+  // modifier commands for the character in the arguments
+  ['H', 'a'],
+  ['c', 'a'],
+  ['d', 'a'],
+  ['k', 'a'],
+  ['v', 'a'],
+  // modifier symbols for the subsequent character
+  ["'", ''],
+  ['^', ''],
+  ['"', ''],
+  ['=', ''],
+  ['.', ''],
+])
+
+type TextNode = {
+  from: number
+  to: number
+  text: string
+  context: Context
+}
+
+export const countWordsInFile = (
+  data: WordCountData,
+  projectSnapshot: ProjectSnapshot,
+  docPath: string,
+  segmenters: Segmenters
+) => {
+  debugConsole.log(`Counting words in ${docPath}`)
+
+  const content = projectSnapshot.getDocContents(docPath) // TODO: try with extensions
+  if (!content) return
+
+  // TODO: language from file extension
+  const tree = LaTeXLanguage.parser.parse(content)
+
+  let currentContext: Context = 'text'
+
+  const textNodes: TextNode[] = []
+
+  const iterateNode = (nodeRef: SyntaxNodeRef, context: Context = 'text') => {
+    const previousContext = currentContext
+    currentContext = context
+    const { node } = nodeRef
+    node.cursor().iterate(childNodeRef => {
+      // TODO: a better way to iterate only descendants?
+      if (childNodeRef.node !== node) {
+        return bodyMatcher(childNodeRef.type)?.(childNodeRef)
+      }
+    })
+    currentContext = previousContext
+  }
+
+  const headMatcher = NodeType.match<
+    (nodeRef: SyntaxNodeRef) => boolean | void
+  >({
+    Title(nodeRef) {
+      data.headers++
+      iterateNode(nodeRef, 'header')
+      return false
+    },
+  })
+
+  const bodyMatcher = NodeType.match<
+    (nodeRef: SyntaxNodeRef) => boolean | void
+  >({
+    Normal(nodeRef) {
+      textNodes.push({
+        from: nodeRef.from,
+        to: nodeRef.to,
+        text: content.substring(nodeRef.from, nodeRef.to),
+        context: currentContext,
+      })
+    },
+    Command(nodeRef) {
+      const child = nodeRef.node.getChild('UnknownCommand')
+      if (!child) return
+
+      const grandchild = child.getChild('CtrlSeq') ?? child.getChild('CtrlSym')
+      if (!grandchild) return
+
+      const commandName = content.substring(grandchild.from + 1, grandchild.to)
+      if (!commandName) return
+
+      if (!replacementsMap.has(commandName)) return
+
+      const text = replacementsMap.get(commandName)!
+      textNodes.push({
+        from: nodeRef.from,
+        to: nodeRef.to,
+        text,
+        context: currentContext,
+      })
+      return false
+    },
+    BeginEnv(nodeRef) {
+      const envName = content
+        ?.substring(nodeRef.from + '\\begin{'.length, nodeRef.to - 1)
+        .replace(/\*$/, '')
+
+      if (envName === 'abstract') {
+        data.headers++
+        iterateNode(nodeRef, 'abstract')
+        return false
+      }
+    },
+    'ShortTextArgument ShortOptionalArg'() {
+      return false
+    },
+    SectioningArgument(nodeRef) {
+      data.headers++
+      iterateNode(nodeRef, 'header')
+      return false
+    },
+    'DisplayMath BracketMath'() {
+      data.mathDisplay++
+    },
+    'InlineMath ParenMath'() {
+      data.mathInline++
+    },
+    Caption(nodeRef) {
+      iterateNode(nodeRef, 'caption')
+      return false
+    },
+    'FootnoteCommand EndnoteCommand'(nodeRef) {
+      iterateNode(nodeRef, 'footnote')
+      return false
+    },
+    'IncludeArgument InputArgument'(nodeRef) {
+      let path = content.substring(nodeRef.from + 1, nodeRef.to - 1)
+      if (!/\.\w+$/.test(path)) {
+        path += '.tex'
+      }
+      debugConsole.log(path)
+      if (path) {
+        countWordsInFile(data, projectSnapshot, path, segmenters)
+      }
+    },
+  })
+
+  const preambleExtent = findPreambleExtent(tree)
+
+  tree.iterate({
+    from: 0,
+    to: preambleExtent.to,
+    enter(nodeRef) {
+      return headMatcher(nodeRef.type)?.(nodeRef)
+    },
+  })
+
+  tree.iterate({
+    from: preambleExtent.to,
+    enter(nodeRef) {
+      return bodyMatcher(nodeRef.type)?.(nodeRef)
+    },
+  })
+
+  const texts: Record<Context, string> = {
+    abstract: '',
+    header: '',
+    caption: '',
+    text: '',
+    footnote: '',
+  }
+
+  let pos = 0
+  for (const textNode of textNodes) {
+    if (textNode.from !== pos) {
+      texts[textNode.context] += ' '
+    }
+    texts[textNode.context] += textNode.text
+    pos = textNode.to
+  }
+
+  for (const [context, text] of Object.entries(texts)) {
+    const counter = counters[context as Context]
+
+    for (const value of segmenters.word.segment(text)) {
+      if (value.isWordLike) {
+        data[counter.word]++
+      }
+    }
+
+    for (const value of segmenters.character.segment(text)) {
+      // TODO: option for whether to include whitespace?
+      if (!whiteSpaceRe.test(value.segment)) {
+        data[counter.character]++
+      }
+    }
+  }
+}

services/web/frontend/js/features/word-count-modal/utils/find-preamble-extent.ts

@@ -0,0 +1,40 @@
+import { NodeType, SyntaxNodeRef, Tree } from '@lezer/common'
+import { ancestorOfNodeWithType } from '@/features/source-editor/utils/tree-operations/ancestors'
+
+export const findPreambleExtent = (tree: Tree) => {
+  const preamble = { to: 0 }
+
+  let seenDocumentEnvironment = false
+
+  const preambleMatcher = NodeType.match<(nodeRef: SyntaxNodeRef) => void>({
+    'Title Author Affil Affiliation'(nodeRef) {
+      preamble.to = nodeRef.node.to
+    },
+    DocumentEnvironment(nodeRef) {
+      // only count the first instance of DocumentEnvironment
+      if (!seenDocumentEnvironment) {
+        preamble.to =
+          nodeRef.node.getChild('Content')?.from ?? nodeRef.node.from
+        seenDocumentEnvironment = true
+      }
+    },
+    Maketitle(nodeRef) {
+      // count \maketitle inside DocumentEnvironment
+      if (
+        ancestorOfNodeWithType(nodeRef.node, '$Environment')?.type.is(
+          'DocumentEnvironment'
+        )
+      ) {
+        preamble.to = nodeRef.node.from
+      }
+    },
+  })
+
+  tree.iterate({
+    enter(nodeRef) {
+      return preambleMatcher(nodeRef.type)?.(nodeRef)
+    },
+  })
+
+  return preamble
+}

services/web/frontend/js/features/word-count-modal/utils/segmenters.ts

@@ -0,0 +1,64 @@
+const wordRe = /['\-.\p{L}]+/gu
+const wordLikeRe = /\p{L}/gu // must contain at least one "letter" to be a word
+const characterRe = /\S/gu
+
+type SegmentDataLike = {
+  segment: string
+  isWordLike?: boolean
+}
+
+type SegmenterLike = {
+  segment(input: string): {
+    [Symbol.iterator](): IterableIterator<SegmentDataLike>
+  }
+}
+
+export type Segmenters = {
+  word: SegmenterLike
+  character: SegmenterLike
+}
+
+export const createSegmenters = (locale?: string): Segmenters => {
+  if (!Intl.Segmenter) {
+    return genericSegmenters
+  }
+
+  try {
+    return {
+      word: new Intl.Segmenter(locale, {
+        granularity: 'word', // TODO: count hyphenated words as a single word
+      }),
+      character: new Intl.Segmenter(locale, {
+        granularity: 'grapheme',
+      }),
+    }
+  } catch {
+    return genericSegmenters
+  }
+}
+
+const genericSegmenters = {
+  word: {
+    segment(input: string) {
+      const segments: SegmentDataLike[] = []
+      for (const match of input.matchAll(wordRe)) {
+        segments.push({
+          segment: match[0],
+          isWordLike: wordLikeRe.test(match[0]),
+        })
+      }
+      return segments
+    },
+  },
+  character: {
+    segment(input: string) {
+      const segments: SegmentDataLike[] = []
+      for (const match of input.matchAll(characterRe)) {
+        segments.push({
+          segment: match[0],
+        })
+      }
+      return segments
+    },
+  },
+}

services/web/frontend/stories/word-count-modal.stories.jsx

@@ -1,63 +0,0 @@
-import useFetchMock from './hooks/use-fetch-mock'
-import WordCountModal from '../js/features/word-count-modal/components/word-count-modal'
-import { ScopeDecorator } from './decorators/scope'
-
-const counts = {
-  headers: 4,
-  mathDisplay: 40,
-  mathInline: 400,
-  textWords: 4000,
-}
-
-const messages = [
-  'Lorem ipsum dolor sit amet.',
-  'Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.',
-].join('\n')
-
-export const WordCount = args => {
-  useFetchMock(fetchMock => {
-    fetchMock.get(
-      'express:/project/:projectId/wordcount',
-      { status: 200, body: { texcount: counts } },
-      { delay: 500 }
-    )
-  })
-
-  return <WordCountModal {...args} />
-}
-
-export const WordCountWithMessages = args => {
-  useFetchMock(fetchMock => {
-    fetchMock.get(
-      'express:/project/:projectId/wordcount',
-      { status: 200, body: { texcount: { ...counts, messages } } },
-      { delay: 500 }
-    )
-  })
-
-  return <WordCountModal {...args} />
-}
-
-export const ErrorResponse = args => {
-  useFetchMock(fetchMock => {
-    fetchMock.get(
-      'express:/project/:projectId/wordcount',
-      { status: 500 },
-      { delay: 500 }
-    )
-  })
-
-  return <WordCountModal {...args} />
-}
-
-export default {
-  title: 'Editor / Modals / Word Count',
-  component: WordCountModal,
-  args: {
-    show: true,
-  },
-  argTypes: {
-    handleHide: { action: 'close modal' },
-  },
-  decorators: [ScopeDecorator],
-}

services/web/frontend/stories/word-count-modal.stories.tsx

@@ -0,0 +1,80 @@
+import { Meta, StoryObj } from '@storybook/react'
+import WordCountModal from '@/features/word-count-modal/components/word-count-modal'
+import { ScopeDecorator } from './decorators/scope'
+import useFetchMock from './hooks/use-fetch-mock'
+
+export default {
+  title: 'Editor / Modals / Word Count',
+  component: WordCountModal,
+  args: {
+    show: true,
+  },
+  argTypes: {
+    handleHide: {
+      action: 'close modal',
+    },
+  },
+  decorators: [Story => ScopeDecorator(Story)],
+} satisfies Meta
+
+type Story = StoryObj<typeof WordCountModal>
+
+const counts = {
+  headers: 4,
+  mathDisplay: 40,
+  mathInline: 400,
+  textWords: 4000,
+}
+
+const messages = [
+  'Lorem ipsum dolor sit amet.',
+  'Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.',
+].join('\n')
+
+export const WordCount: Story = {
+  decorators: [
+    Story => {
+      useFetchMock(fetchMock => {
+        fetchMock.get(
+          'express:/project/:projectId/wordcount',
+          { status: 200, body: { texcount: counts } },
+          { delay: 500 }
+        )
+      })
+
+      return <Story />
+    },
+  ],
+}
+
+export const WordCountWithMessages: Story = {
+  decorators: [
+    Story => {
+      useFetchMock(fetchMock => {
+        fetchMock.get(
+          'express:/project/:projectId/wordcount',
+          { status: 200, body: { texcount: { ...counts, messages } } },
+          { delay: 500 }
+        )
+      })
+
+      return <Story />
+    },
+  ],
+}
+
+export const ErrorResponse: Story = {
+  decorators: [
+    Story => {
+      useFetchMock(fetchMock => {
+        fetchMock.get(
+          'express:/project/:projectId/wordcount',
+          { status: 500 },
+          { delay: 500 }
+        )
+      })
+
+      return <Story />
+    },
+  ],
+}

services/web/test/frontend/features/word-count-modal/components/word-count-modal.spec.tsx

@@ -0,0 +1,111 @@
+import WordCountModal from '@/features/word-count-modal/components/word-count-modal'
+import { EditorProviders } from '../../../helpers/editor-providers'
+
+describe('<WordCountModal />', function () {
+  beforeEach(function () {
+    cy.interceptCompile()
+  })
+
+  it('renders the translated modal title', function () {
+    cy.intercept('/project/*/wordcount*', {
+      body: { texcount: { messages: 'This is a test' } },
+    })
+
+    cy.mount(
+      <EditorProviders projectId="project-1" clsiServerId="clsi-server-1">
+        <WordCountModal show handleHide={cy.stub()} />
+      </EditorProviders>
+    )
+
+    cy.findByText('Word Count')
+    cy.findByText(/something went wrong/).should('not.exist')
+  })
+
+  it('renders a loading message when loading', function () {
+    const { promise, resolve } = Promise.withResolvers<void>()
+
+    cy.intercept('/project/*/wordcount*', async req => {
+      await promise
+      req.reply({ texcount: { messages: 'This is a test' } })
+    })
+
+    cy.mount(
+      <EditorProviders projectId="project-1" clsiServerId="clsi-server-1">
+        <WordCountModal show handleHide={cy.stub()} />
+      </EditorProviders>
+    )
+
+    cy.findByText('Loading…').then(() => {
+      resolve()
+    })
+
+    cy.findByText('This is a test')
+  })
+
+  it('renders an error message and hides loading message on error', function () {
+    cy.intercept('/project/*/wordcount?*', {
+      statusCode: 500,
+    })
+
+    cy.mount(
+      <EditorProviders projectId="project-1" clsiServerId="clsi-server-1">
+        <WordCountModal show handleHide={cy.stub()} />
+      </EditorProviders>
+    )
+
+    cy.findByText('Sorry, something went wrong')
+
+    cy.findByText('Loading').should('not.exist')
+  })
+
+  it('displays messages', function () {
+    cy.intercept('/project/*/wordcount*', {
+      body: { texcount: { messages: 'This is a test' } },
+    })
+
+    cy.mount(
+      <EditorProviders projectId="project-1" clsiServerId="clsi-server-1">
+        <WordCountModal show handleHide={cy.stub()} />
+      </EditorProviders>
+    )
+
+    cy.findByText('This is a test')
+  })
+
+  it('displays counts data', function () {
+    cy.intercept('/project/*/wordcount*', {
+      body: {
+        texcount: {
+          textWords: 500,
+          headWords: 100,
+          outside: 200,
+          mathDisplay: 2,
+          mathInline: 3,
+          headers: 4,
+        },
+      },
+    })
+
+    cy.mount(
+      <EditorProviders projectId="project-1" clsiServerId="clsi-server-1">
+        <WordCountModal show handleHide={cy.stub()} />
+      </EditorProviders>
+    )
+
+    cy.findByText((content, element) => {
+      return /^Total Words\s*:\s*500$/.test(element!.textContent!.trim())
+    })
+
+    cy.findByText((content, element) => {
+      return /^Math Display\s*:\s*2$/.test(element!.textContent!.trim())
+    })
+
+    cy.findByText((content, element) => {
+      return /^Math Inline\s*:\s*3$/.test(element!.textContent!.trim())
+    })
+
+    cy.findByText((content, element) => {
+      return /^Headers\s*:\s*4$/.test(element!.textContent!.trim())
+    })
+  })
+})

services/web/test/frontend/features/word-count-modal/components/word-count-modal.test.jsx

@@ -1,122 +0,0 @@
-import { screen } from '@testing-library/react'
-import { expect } from 'chai'
-import sinon from 'sinon'
-import fetchMock from 'fetch-mock'
-import { renderWithEditorContext } from '../../../helpers/render-with-context'
-import WordCountModal from '../../../../../frontend/js/features/word-count-modal/components/word-count-modal'
-
-describe('<WordCountModal />', function () {
-  afterEach(function () {
-    fetchMock.removeRoutes().clearHistory()
-  })
-
-  const contextProps = {
-    projectId: 'project-1',
-    clsiServerId: 'clsi-server-1',
-  }
-
-  it('renders the translated modal title', async function () {
-    fetchMock.get('express:/project/:projectId/wordcount', () => {
-      return { status: 200, body: { texcount: { messages: 'This is a test' } } }
-    })
-
-    const handleHide = sinon.stub()
-
-    renderWithEditorContext(
-      <WordCountModal show handleHide={handleHide} />,
-      contextProps
-    )
-
-    await screen.findByText('Word Count')
-  })
-
-  it('renders a loading message when loading', async function () {
-    fetchMock.get('express:/project/:projectId/wordcount', () => {
-      return { status: 200, body: { texcount: { messages: 'This is a test' } } }
-    })
-
-    const handleHide = sinon.stub()
-
-    renderWithEditorContext(
-      <WordCountModal show handleHide={handleHide} />,
-      contextProps
-    )
-
-    await screen.findByText('Loading…')
-
-    await screen.findByText('This is a test')
-  })
-
-  it('renders an error message and hides loading message on error', async function () {
-    fetchMock.get('express:/project/:projectId/wordcount', 500)
-
-    const handleHide = sinon.stub()
-
-    renderWithEditorContext(
-      <WordCountModal show handleHide={handleHide} />,
-      contextProps
-    )
-
-    await screen.findByText('Sorry, something went wrong')
-
-    expect(screen.queryByText(/Loading/)).to.not.exist
-  })
-
-  it('displays messages', async function () {
-    fetchMock.get('express:/project/:projectId/wordcount', () => {
-      return {
-        status: 200,
-        body: {
-          texcount: {
-            messages: 'This is a test',
-          },
-        },
-      }
-    })
-
-    const handleHide = sinon.stub()
-
-    renderWithEditorContext(
-      <WordCountModal show handleHide={handleHide} />,
-      contextProps
-    )
-
-    await screen.findByText('This is a test')
-  })
-
-  it('displays counts data', async function () {
-    fetchMock.get('express:/project/:projectId/wordcount', () => {
-      return {
-        status: 200,
-        body: {
-          texcount: {
-            textWords: 100,
-            mathDisplay: 200,
-            mathInline: 300,
-            headers: 400,
-          },
-        },
-      }
-    })
-
-    const handleHide = sinon.stub()
-
-    renderWithEditorContext(
-      <WordCountModal show handleHide={handleHide} />,
-      contextProps
-    )
-
-    await screen.findByText((content, element) =>
-      element.textContent.trim().match(/^Total Words\s*:\s*100$/)
-    )
-    await screen.findByText((content, element) =>
-      element.textContent.trim().match(/^Math Display\s*:\s*200$/)
-    )
-    await screen.findByText((content, element) =>
-      element.textContent.trim().match(/^Math Inline\s*:\s*300$/)
-    )
-    await screen.findByText((content, element) =>
-      element.textContent.trim().match(/^Headers\s*:\s*400$/)
-    )
-  })
-})