Merge pull request #24468 from overleaf/mj-client-side-references

[web] Perform ARS on client-side

GitOrigin-RevId: 19703c82758cae450fe52463ad9612d3a2383ba0

services/web/frontend/js/features/ide-react/context/react-context-root.tsx

@@ -87,13 +87,13 @@ export const ReactContextRoot: FC<
                             <Providers.EditorProvider>
                               <Providers.FileTreeDataProvider>
                                 <Providers.FileTreePathProvider>
-                                  <Providers.ReferencesProvider>
-                                    <Providers.UserFeaturesProvider>
-                                      <Providers.PermissionsProvider>
-                                        <Providers.RailProvider>
-                                          <Providers.LayoutProvider>
-                                            <Providers.ProjectSettingsProvider>
-                                              <Providers.EditorManagerProvider>
+                                  <Providers.UserFeaturesProvider>
+                                    <Providers.PermissionsProvider>
+                                      <Providers.RailProvider>
+                                        <Providers.LayoutProvider>
+                                          <Providers.ProjectSettingsProvider>
+                                            <Providers.EditorManagerProvider>
+                                              <Providers.ReferencesProvider>
                                                 <Providers.LocalCompileProvider>
                                                   <Providers.DetachCompileProvider>
                                                     <Providers.ChatProvider>
@@ -113,13 +113,13 @@ export const ReactContextRoot: FC<
                                                     </Providers.ChatProvider>
                                                   </Providers.DetachCompileProvider>
                                                 </Providers.LocalCompileProvider>
-                                              </Providers.EditorManagerProvider>
-                                            </Providers.ProjectSettingsProvider>
-                                          </Providers.LayoutProvider>
-                                        </Providers.RailProvider>
-                                      </Providers.PermissionsProvider>
-                                    </Providers.UserFeaturesProvider>
-                                  </Providers.ReferencesProvider>
+                                              </Providers.ReferencesProvider>
+                                            </Providers.EditorManagerProvider>
+                                          </Providers.ProjectSettingsProvider>
+                                        </Providers.LayoutProvider>
+                                      </Providers.RailProvider>
+                                    </Providers.PermissionsProvider>
+                                  </Providers.UserFeaturesProvider>
                                 </Providers.FileTreePathProvider>
                               </Providers.FileTreeDataProvider>
                             </Providers.EditorProvider>

services/web/frontend/js/features/ide-react/context/references-context.tsx

@@ -7,21 +7,31 @@ import {
   useCallback,
   useMemo,
   useState,
+  useRef,
 } from 'react'
 import { useIdeReactContext } from '@/features/ide-react/context/ide-react-context'
 import { useConnectionContext } from '@/features/ide-react/context/connection-context'
-import { postJSON } from '@/infrastructure/fetch-json'
 import { ShareJsDoc } from '@/features/ide-react/editor/share-js-doc'
 import { useFileTreeData } from '@/shared/context/file-tree-data-context'
 import { findDocEntityById } from '@/features/ide-react/util/find-doc-entity-by-id'
 import { IdeEvents } from '@/features/ide-react/create-ide-event-emitter'
-import { debugConsole } from '@/utils/debugging'
 import useEventListener from '@/shared/hooks/use-event-listener'
+import { useProjectContext } from '@/shared/context/project-context'
+import { useEditorManagerContext } from './editor-manager-context'
+import { signalWithTimeout } from '@/utils/abort-signal'
+import { postJSON } from '@/infrastructure/fetch-json'
+import { debugConsole } from '@/utils/debugging'
+import { useFeatureFlag } from '@/shared/context/split-test-context'
+import type { ReferenceIndexer } from '../references/reference-indexer'
+import { AdvancedReferenceSearchResult } from '@/features/ide-react/references/types'
 
 export const ReferencesContext = createContext<
   | {
       referenceKeys: Set<string>
       indexAllReferences: (shouldBroadcast: boolean) => Promise<void>
+      searchLocalReferences: (
+        query: string
+      ) => Promise<AdvancedReferenceSearchResult>
     }
   | undefined
 >(undefined)
@@ -32,14 +42,18 @@ export const ReferencesProvider: FC<React.PropsWithChildren> = ({
   const { fileTreeData } = useFileTreeData()
   const { eventEmitter, projectId } = useIdeReactContext()
   const { socket } = useConnectionContext()
+  const { projectSnapshot } = useProjectContext()
+  const { openDocs } = useEditorManagerContext()
+  const abortControllerRef = useRef<AbortController | null>(null)
 
   const [referenceKeys, setReferenceKeys] = useState(new Set<string>())
+  const clientSideReferences = useFeatureFlag('client-side-references')
 
   const [existingIndexHash, setExistingIndexHash] = useState<
     Record<string, { hash: string; timestamp: number }>
   >({})
 
-  const indexAllReferences = useCallback(
+  const indexAllReferencesServerside = useCallback(
     async (shouldBroadcast: boolean) => {
       return postJSON(`/project/${projectId}/references/indexAll`, {
         body: {
@@ -57,6 +71,54 @@ export const ReferencesProvider: FC<React.PropsWithChildren> = ({
     [projectId]
   )
 
+  const indexerRef = useRef<Promise<ReferenceIndexer> | null>(null)
+  if (clientSideReferences && indexerRef.current === null) {
+    indexerRef.current = import('../references/reference-indexer').then(
+      m => new m.ReferenceIndexer()
+    )
+  }
+
+  const indexAllReferencesLocally = useCallback(
+    async (shouldBroadcast: boolean) => {
+      abortControllerRef.current?.abort()
+
+      if (!indexerRef.current) {
+        return
+      }
+
+      abortControllerRef.current = new AbortController()
+      const signal = abortControllerRef.current.signal
+
+      await openDocs.awaitBufferedOps(signalWithTimeout(signal, 5000))
+      await projectSnapshot.refresh()
+
+      if (signal.aborted) {
+        return
+      }
+
+      const indexer = await indexerRef.current
+      const keys = await indexer.updateFromSnapshot(projectSnapshot, { signal })
+      if (signal.aborted) {
+        return
+      }
+      setReferenceKeys(keys)
+      if (shouldBroadcast) {
+        // Inform other clients about change in keys
+        await postJSON(`/project/${projectId}/references/indexAll`, {
+          body: { shouldBroadcast: true },
+        }).catch(error => {
+          // allow the request to fail
+          debugConsole.error(error)
+        })
+      }
+    },
+    [projectSnapshot, openDocs, projectId]
+  )
+
+  const indexAllReferences = clientSideReferences
+    ? indexAllReferencesLocally
+    : indexAllReferencesServerside
+
   const indexReferencesIfDocModified = useCallback(
     (doc: ShareJsDoc, shouldBroadcast: boolean) => {
       // avoid reindexing references if the bib file has not changed since the
@@ -115,9 +177,13 @@ export const ReferencesProvider: FC<React.PropsWithChildren> = ({
       // We only need to grab the references when the editor first loads,
       // not on every reconnect
       socket.on('references:keys:updated', (keys, allDocs) => {
-        setReferenceKeys(oldKeys =>
-          allDocs ? new Set(keys) : new Set([...oldKeys, ...keys])
-        )
+        if (clientSideReferences) {
+          indexAllReferences(false)
+        } else {
+          setReferenceKeys(oldDocs =>
+            allDocs ? new Set(keys) : new Set([...oldDocs, ...keys])
+          )
+        }
       })
       indexAllReferences(false)
     }
@@ -127,14 +193,26 @@ export const ReferencesProvider: FC<React.PropsWithChildren> = ({
     return () => {
       eventEmitter.off('project:joined', handleProjectJoined)
     }
-  }, [eventEmitter, indexAllReferences, socket])
+  }, [eventEmitter, indexAllReferences, socket, clientSideReferences])
+
+  const searchLocalReferences = useCallback(
+    async (query: string): Promise<AdvancedReferenceSearchResult> => {
+      if (!indexerRef.current) {
+        return { hits: [] }
+      }
+      const indexer = await indexerRef.current
+      return await indexer.search(query)
+    },
+    []
+  )
 
   const value = useMemo(
     () => ({
       referenceKeys,
       indexAllReferences,
+      searchLocalReferences,
     }),
-    [indexAllReferences, referenceKeys]
+    [indexAllReferences, referenceKeys, searchLocalReferences]
   )
 
   return (