Fix sample selection with map panel (#4739)

* clear records on page refresh

* refactor as independent hook, add hook test

* rm useRef

* fixing selections with partial index records

* add range tests

* relay cleanup tweak

* typos

app/packages/core/src/components/Grid/Grid.tsx

@@ -1,9 +1,9 @@
 import styles from "./Grid.module.css";
 
-import type { Lookers } from "@fiftyone/state";
-
 import { freeVideos } from "@fiftyone/looker";
-import Spotlight, { type ID } from "@fiftyone/spotlight";
+import type { ID } from "@fiftyone/spotlight";
+import Spotlight from "@fiftyone/spotlight";
+import type { Lookers } from "@fiftyone/state";
 import * as fos from "@fiftyone/state";
 import React, {
   useEffect,
@@ -17,6 +17,7 @@ import { v4 as uuid } from "uuid";
 import { gridCrop, gridSpacing, pageParameters } from "./recoil";
 import useAt from "./useAt";
 import useEscape from "./useEscape";
+import useRecords from "./useRecords";
 import useRefreshers from "./useRefreshers";
 import useSelect from "./useSelect";
 import useSelectSample from "./useSelectSample";
@@ -26,19 +27,22 @@ import useThreshold from "./useThreshold";
 function Grid() {
   const id = useMemo(() => uuid(), []);
   const lookerStore = useMemo(() => new WeakMap<ID, Lookers>(), []);
-  const selectSample = useRef<ReturnType<typeof useSelectSample>>();
-  const [resizing, setResizing] = useState(false);
-
   const spacing = useRecoilValue(gridSpacing);
 
+  const selectSample = useRef<ReturnType<typeof useSelectSample>>();
   const { pageReset, reset } = useRefreshers();
-  const { get, set } = useAt(pageReset);
+  const [resizing, setResizing] = useState(false);
   const threshold = useThreshold();
 
-  const { page, records, store } = useSpotlightPager({
-    pageSelector: pageParameters,
-    zoomSelector: gridCrop,
-  });
+  const records = useRecords(pageReset);
+  const { page, store } = useSpotlightPager(
+    {
+      pageSelector: pageParameters,
+      zoomSelector: gridCrop,
+    },
+    records
+  );
+  const { get, set } = useAt(pageReset);
 
   const lookerOptions = fos.useLookerOptions(false);
   const createLooker = fos.useCreateLooker(false, true, lookerOptions);
@@ -71,7 +75,7 @@ function Grid() {
 
         const init = (l) => {
           l.addEventListener("selectthumbnail", ({ detail }: CustomEvent) => {
-            selectSample.current?.(records, detail);
+            selectSample.current?.(detail);
           });
           lookerStore.set(id, l);
           l.attach(element, dimensions);
@@ -97,7 +101,7 @@ function Grid() {
     store,
     threshold,
   ]);
-  selectSample.current = useSelectSample();
+  selectSample.current = useSelectSample(records);
   useSelect(lookerOptions, lookerStore, spotlight);
 
   useLayoutEffect(() => {

app/packages/core/src/components/Grid/useRecords.test.ts

@@ -0,0 +1,23 @@
+import { act, renderHook } from "@testing-library/react-hooks";
+import { describe, expect, it } from "vitest";
+import useRecords from "./useRecords";
+
+describe("useRecords", () => {
+  it("return new records when clear string changes", () => {
+    const { result, rerender } = renderHook(
+      (clear: string) => useRecords(clear),
+      { initialProps: "one" }
+    );
+    expect(result.current.size).toBe(0);
+
+    act(() => {
+      result.current.set("one", 1);
+    });
+
+    expect(result.current.size).toBe(1);
+
+    rerender("two");
+
+    expect(result.current.size).toBe(0);
+  });
+});

app/packages/core/src/components/Grid/useRecords.ts

@@ -0,0 +1,10 @@
+import { useMemo } from "react";
+
+export type Records = Map<string, number>;
+
+export default (clear: string) => {
+  return useMemo(() => {
+    clear;
+    return new Map<string, number>();
+  }, [clear]);
+};

app/packages/core/src/components/Grid/useSelect.ts

@@ -1,5 +1,5 @@
 import type Spotlight from "@fiftyone/spotlight";
-import { ID } from "@fiftyone/spotlight";
+import type { ID } from "@fiftyone/spotlight";
 import * as fos from "@fiftyone/state";
 import { useEffect } from "react";
 import { useRecoilValue } from "recoil";

app/packages/core/src/components/Grid/useSelectSample.test.ts

@@ -0,0 +1,29 @@
+import { describe, expect, it } from "vitest";
+import { addRange, removeRange } from "./useSelectSample";
+
+describe("range selection tests", () => {
+  it("adds a range, and includes selections without an index record", () => {
+    const result = addRange(
+      2,
+      new Set(["0", "other"]),
+      new Map([
+        ["0", 0],
+        ["1", 1],
+        ["2", 2],
+      ])
+    );
+    expect(result).toStrictEqual(new Set(["0", "1", "2", "other"]));
+  });
+
+  it("removes a range, and includes selections without an index record", () => {
+    const result = removeRange(
+      1,
+      new Set(["0", "1", "other"]),
+      new Map([
+        ["0", 0],
+        ["1", 1],
+      ])
+    );
+    expect(result).toStrictEqual(new Set(["other"]));
+  });
+});

app/packages/core/src/components/Grid/useSelectSample.ts

@@ -1,13 +1,12 @@
-import { ID } from "@fiftyone/spotlight";
+import type { ID } from "@fiftyone/spotlight";
 import type { Sample } from "@fiftyone/state";
-
 import {
   selectedSampleObjects,
   selectedSamples,
   useSetSelected,
 } from "@fiftyone/state";
-import { MutableRefObject } from "react";
 import { useRecoilCallback } from "recoil";
+import type { Records } from "./useRecords";
 
 export interface SelectThumbnailData {
   shiftKey: boolean;
@@ -16,50 +15,62 @@ export interface SelectThumbnailData {
   symbol: ID;
 }
 
-type Records = MutableRefObject<Map<string, number>>;
-
-const argFact = (compareFn) => (array) =>
-  array.map((el, idx) => [el, idx]).reduce(compareFn)[1];
-
-const argMin = argFact((max, el) => (el[0] < max[0] ? el : max));
-
-const addRange = (index: number, items: string[], records: Records) => {
+export const addRange = (
+  index: number,
+  selected: Set<string>,
+  records: Records
+) => {
+  // filter selections without an index record
+  const items = [...selected].filter((i) => records.has(i));
   const reverse = Object.fromEntries(
-    Array.from(records.current.entries()).map(([k, v]) => [v, k])
-  );
-
-  const min = argMin(
-    items.map((id) => Math.abs(records.current.get(id) - index))
+    Array.from(records.entries()).map(([k, v]) => [v, k])
   );
+  const min = argMin(items.map((id) => Math.abs(get(records, id) - index)));
 
-  const close = records.current.get(items[min]);
+  const close = get(records, items[min]);
 
   const [start, end] = index < close ? [index, close] : [close, index];
 
   const added = new Array(end - start + 1)
     .fill(0)
     .map((_, i) => reverse[i + start]);
 
-  return new Set([...items, ...added]);
+  return new Set([...selected, ...added]);
 };
 
-const removeRange = (
+const argFact = (compareFn) => (array) =>
+  array.map((el, idx) => [el, idx]).reduce(compareFn)[1];
+
+const argMin = argFact((max, el) => (el[0] < max[0] ? el : max));
+
+const get = (records: Records, id: string) => {
+  const index = records.get(id);
+  if (index !== undefined) {
+    return index;
+  }
+
+  throw new Error(`record '${id}' not found`);
+};
+
+export const removeRange = (
   index: number,
   selected: Set<string>,
   records: Records
 ) => {
+  // filter selections without an index record
+  const items = new Set([...selected].filter((i) => records.has(i)));
   const reverse = Object.fromEntries(
-    Array.from(records.current.entries()).map(([k, v]) => [v, k])
+    Array.from(records.entries()).map(([k, v]) => [v, k])
   );
 
   let before = index;
-  while (selected.has(reverse[before])) {
+  while (items.has(reverse[before])) {
     before--;
   }
   before += 1;
 
   let after = index;
-  while (selected.has(reverse[after])) {
+  while (items.has(reverse[after])) {
     after++;
   }
   after -= 1;
@@ -73,31 +84,44 @@ const removeRange = (
       ? [before, index]
       : [index, after];
 
-  return new Set(
-    Array.from(selected).filter(
-      (s) => records.current.get(s) < start || records.current.get(s) > end
+  const next = new Set(
+    Array.from(items).filter(
+      (s) => get(records, s) < start || get(records, s) > end
     )
   );
+
+  for (const id of selected) {
+    if (records.has(id)) continue;
+    // not in index records so it was not removed, add it back
+    next.add(id);
+  }
+
+  return next;
 };
 
-export default () => {
+export default (records: Records) => {
   const setSelected = useSetSelected();
 
   return useRecoilCallback(
     ({ set, snapshot }) =>
-      async (
-        records: Records,
-        { shiftKey, id: sampleId, sample, symbol }: SelectThumbnailData
-      ) => {
-        let selected = new Set(await snapshot.getPromise(selectedSamples));
+      async ({
+        shiftKey,
+        id: sampleId,
+        sample,
+        symbol,
+      }: SelectThumbnailData) => {
+        const current = new Set(await snapshot.getPromise(selectedSamples));
+        let selected = new Set(current);
         const selectedObjects = new Map(
           await snapshot.getPromise(selectedSampleObjects)
         );
 
-        const items = Array.from(selected);
-        const index = records.current.get(symbol.description);
+        const index = get(records, symbol.description);
         if (shiftKey && !selected.has(sampleId)) {
-          selected = addRange(index, items, records);
+          selected = new Set([
+            ...selected,
+            ...addRange(index, selected, records),
+          ]);
         } else if (shiftKey) {
           selected = removeRange(index, selected, records);
         } else {
@@ -116,6 +140,6 @@ export default () => {
         set(selectedSampleObjects, selectedObjects);
         setSelected(new Set(selected));
       },
-    [setSelected]
+    [records, setSelected]
   );
 };

app/packages/core/src/components/Grid/useSpotlightPager.ts

@@ -14,6 +14,7 @@ import {
 import type { RecoilValueReadOnly } from "recoil";
 import { useRecoilCallback, useRecoilValue } from "recoil";
 import type { Subscription } from "relay-runtime";
+import type { Records } from "./useRecords";
 
 export const PAGE_SIZE = 20;
 
@@ -45,15 +46,18 @@ const processSamplePageData = (
   });
 };
 
-const useSpotlightPager = ({
-  pageSelector,
-  zoomSelector,
-}: {
-  pageSelector: RecoilValueReadOnly<
-    (page: number, pageSize: number) => VariablesOf<foq.paginateSamplesQuery>
-  >;
-  zoomSelector: RecoilValueReadOnly<boolean>;
-}) => {
+const useSpotlightPager = (
+  {
+    pageSelector,
+    zoomSelector,
+  }: {
+    pageSelector: RecoilValueReadOnly<
+      (page: number, pageSize: number) => VariablesOf<foq.paginateSamplesQuery>
+    >;
+    zoomSelector: RecoilValueReadOnly<boolean>;
+  },
+  records: Records
+) => {
   const environment = useRelayEnvironment();
   const pager = useRecoilValue(pageSelector);
   const zoom = useRecoilValue(zoomSelector);
@@ -62,7 +66,8 @@ const useSpotlightPager = ({
     () => new WeakMap<ID, { sample: fos.Sample; index: number }>(),
     []
   );
-  const records = useRef(new Map<string, number>());
+
+  const keys = useRef(new Set<string>());
 
   const page = useRecoilCallback(
     ({ snapshot }) => {
@@ -89,8 +94,9 @@ const useSpotlightPager = ({
                 data,
                 schema,
                 zoom,
-                records.current
+                records
               );
+              for (const item of items) keys.current.add(item.id.description);
 
               resolve({
                 items,
@@ -110,16 +116,16 @@ const useSpotlightPager = ({
   );
 
   const refresher = useRecoilValue(fos.refresher);
+
   useEffect(() => {
     refresher;
-    const current = records.current;
     const clear = () => {
       commitLocalUpdate(fos.getCurrentEnvironment(), (store) => {
-        for (const id of Array.from(current.keys())) {
+        for (const id of keys.current) {
           store.get(id)?.invalidateRecord();
         }
-        current.clear();
       });
+      keys.current.clear();
     };
 
     const unsubscribe = foq.subscribe(