Add eviction to cache to safe memory on document with many pages

Author: bobsingor

packages/plugin-selection/src/lib/actions.ts

@@ -11,6 +11,7 @@ export const END_SELECTION = 'SELECTION/END_SELECTION';
 export const CLEAR_SELECTION = 'SELECTION/CLEAR_SELECTION';
 export const SET_RECTS = 'SELECTION/SET_RECTS';
 export const SET_SLICES = 'SELECTION/SET_SLICES';
+export const EVICT_PAGE_GEOMETRY = 'SELECTION/EVICT_PAGE_GEOMETRY';
 export const RESET = 'SELECTION/RESET'; // This might be obsolete, but we'll keep it for now
 
 export interface InitSelectionStateAction extends Action {
@@ -60,6 +61,11 @@ export interface SetSlicesAction extends Action {
   payload: { documentId: string; slices: Record<number, { start: number; count: number }> };
 }
 
+export interface EvictPageGeometryAction extends Action {
+  type: typeof EVICT_PAGE_GEOMETRY;
+  payload: { documentId: string; pages: number[] };
+}
+
 export interface ResetAction extends Action {
   type: typeof RESET;
   payload: { documentId: string };
@@ -75,6 +81,7 @@ export type SelectionAction =
   | ClearSelectionAction
   | SetRectsAction
   | SetSlicesAction
+  | EvictPageGeometryAction
   | ResetAction;
 
 export const initSelectionState = (
@@ -132,6 +139,14 @@ export const setSlices = (
   slices: Record<number, { start: number; count: number }>,
 ): SetSlicesAction => ({ type: SET_SLICES, payload: { documentId, slices } });
 
+export const evictPageGeometry = (
+  documentId: string,
+  pages: number[],
+): EvictPageGeometryAction => ({
+  type: EVICT_PAGE_GEOMETRY,
+  payload: { documentId, pages },
+});
+
 export const reset = (documentId: string): ResetAction => ({
   type: RESET,
   payload: { documentId },

packages/plugin-selection/src/lib/reducer.ts

@@ -11,6 +11,7 @@ import {
   SET_RECTS,
   INIT_SELECTION_STATE,
   CLEANUP_SELECTION_STATE,
+  EVICT_PAGE_GEOMETRY,
 } from './actions';
 
 export const initialSelectionDocumentState: SelectionDocumentState = {
@@ -121,6 +122,21 @@ export const selectionReducer = (state = initialState, action: SelectionAction):
       return updateDocState(state, documentId, { ...docState, slices });
     }
 
+    case EVICT_PAGE_GEOMETRY: {
+      const { documentId, pages } = action.payload;
+      const docState = state.documents[documentId];
+      if (!docState) return state;
+      const geometry = { ...docState.geometry };
+      const rects = { ...docState.rects };
+      const slices = { ...docState.slices };
+      for (const p of pages) {
+        delete geometry[p];
+        delete rects[p];
+        delete slices[p];
+      }
+      return updateDocState(state, documentId, { ...docState, geometry, rects, slices });
+    }
+
     case RESET: {
       const { documentId } = action.payload;
       const docState = state.documents[documentId];

packages/plugin-selection/src/lib/selection-plugin.ts

@@ -28,6 +28,7 @@ import { ScrollCapability, ScrollPlugin } from '@embedpdf/plugin-scroll';
 
 import {
   cachePageGeometry,
+  evictPageGeometry,
   setSelection,
   SelectionAction,
   endSelection,
@@ -93,6 +94,9 @@ export class SelectionPlugin extends BasePlugin<
   /** Page callbacks for rect updates, per document */
   private pageCallbacks = new Map<string, Map<number, (data: SelectionRectsCallback) => void>>();
 
+  /** LRU access order for geometry cache, per document (oldest first) */
+  private geoAccessOrder = new Map<string, number[]>();
+
   private readonly menuPlacement$ = createScopedEmitter<
     SelectionMenuPlacement | null,
     SelectionMenuPlacementEvent,
@@ -229,6 +233,7 @@ export class SelectionPlugin extends BasePlugin<
       ]),
     );
     this.pageCallbacks.set(documentId, new Map());
+    this.geoAccessOrder.set(documentId, []);
     this.selecting.set(documentId, false);
     this.anchor.set(documentId, undefined);
     this.hasTextAnchor.set(documentId, false);
@@ -238,6 +243,7 @@ export class SelectionPlugin extends BasePlugin<
     this.dispatch(cleanupSelectionState(documentId));
     this.enabledModesPerDoc.delete(documentId);
     this.pageCallbacks.delete(documentId);
+    this.geoAccessOrder.delete(documentId);
     this.selecting.delete(documentId);
     this.hasTextAnchor.delete(documentId);
     this.anchor.delete(documentId);
@@ -395,6 +401,27 @@ export class SelectionPlugin extends BasePlugin<
       boundingRect: selector.selectBoundingRectForPage(docState, pageIndex),
     });
 
+    // When geometry arrives (possibly re-fetched after eviction), recompute
+    // rects for this page if it falls within an active selection.
+    geoTask.wait((geo) => {
+      const currentState = this.getDocumentState(documentId);
+      const sel = currentState.selection;
+      if (!sel || pageIndex < sel.start.page || pageIndex > sel.end.page) return;
+
+      const sb = sliceBounds(sel, geo, pageIndex);
+      if (!sb) return;
+
+      const pageRects = rectsWithinSlice(geo, sb.from, sb.to);
+      this.dispatch(setRects(documentId, { ...currentState.rects, [pageIndex]: pageRects }));
+      this.dispatch(
+        setSlices(documentId, {
+          ...currentState.slices,
+          [pageIndex]: { start: sb.from, count: sb.to - sb.from + 1 },
+        }),
+      );
+      this.notifyPage(documentId, pageIndex);
+    }, ignore);
+
     // Create text selection handler
     const textHandler = createTextSelectionHandler({
       getGeometry: () => this.getDocumentState(documentId).geometry[pageIndex],
@@ -690,18 +717,59 @@ export class SelectionPlugin extends BasePlugin<
     const task = this.engine.getPageGeometry(coreDoc.document, page);
     task.wait((geo) => {
       this.dispatch(cachePageGeometry(documentId, pageIdx, geo));
+      this.touchGeometry(documentId, pageIdx);
     }, ignore);
     return task;
   }
 
   /* ── geometry cache ───────────────────────────────────── */
   private getOrLoadGeometry(documentId: string, pageIdx: number): PdfTask<PdfPageGeometry> {
     const cached = this.getDocumentState(documentId).geometry[pageIdx];
-    if (cached) return PdfTaskHelper.resolve(cached);
+    if (cached) {
+      this.touchGeometry(documentId, pageIdx);
+      return PdfTaskHelper.resolve(cached);
+    }
 
     return this.getNewPageGeometryAndCache(documentId, pageIdx);
   }
 
+  /* ── geometry LRU eviction ──────────────────────────────── */
+
+  private touchGeometry(documentId: string, pageIdx: number): void {
+    const order = this.geoAccessOrder.get(documentId);
+    if (!order) return;
+
+    const idx = order.indexOf(pageIdx);
+    if (idx > -1) order.splice(idx, 1);
+    order.push(pageIdx);
+
+    this.evictGeometryIfNeeded(documentId);
+  }
+
+  private evictGeometryIfNeeded(documentId: string): void {
+    const max = this.config.maxCachedGeometries ?? 50;
+    const order = this.geoAccessOrder.get(documentId);
+    if (!order || order.length <= max) return;
+
+    const pinned = this.pageCallbacks.get(documentId);
+    const toEvict: number[] = [];
+
+    while (order.length - toEvict.length > max) {
+      const candidate = order.find((p) => !toEvict.includes(p) && !pinned?.has(p));
+      if (candidate === undefined) break;
+      toEvict.push(candidate);
+    }
+
+    if (toEvict.length === 0) return;
+
+    for (const p of toEvict) {
+      const idx = order.indexOf(p);
+      if (idx > -1) order.splice(idx, 1);
+    }
+
+    this.dispatch(evictPageGeometry(documentId, toEvict));
+  }
+
   /* ── selection state updates ───────────────────────────── */
   private beginSelection(documentId: string, page: number, index: number, modeId: string) {
     this.selecting.set(documentId, true);

packages/plugin-selection/src/lib/types.ts

@@ -31,6 +31,12 @@ export interface SelectionPluginConfig extends BasePluginConfig {
    * @default 3
    */
   minSelectionDragDistance?: number;
+  /**
+   * Maximum number of pages whose geometry data is kept in memory per document.
+   * Oldest unused pages are evicted when this limit is exceeded.
+   * @default 50
+   */
+  maxCachedGeometries?: number;
 }
 
 export interface SelectionMenuPlacement {