From e39b16d0a72e727ac0fd549215ef39c0e25350e7 Mon Sep 17 00:00:00 2001
From: Connor Peet <connor@peet.io>
Date: Fri, 1 Nov 2024 15:36:37 -0700
Subject: [PATCH] fix: retain selection state when restoring webview (#544)

Fixes #529
---
 media/editor/dataDisplayContext.tsx | 21 ++++++++++++++++++---
 media/editor/hooks.ts               |  4 ++--
 media/editor/state.ts               | 10 +++++++++-
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/media/editor/dataDisplayContext.tsx b/media/editor/dataDisplayContext.tsx
index be5414e..3d2bd67 100644
--- a/media/editor/dataDisplayContext.tsx
+++ b/media/editor/dataDisplayContext.tsx
@@ -3,9 +3,9 @@ import { createContext, useContext, useEffect, useState } from "react";
 import { SetterOrUpdater } from "recoil";
 import { HexDocumentEdit } from "../../shared/hexDocumentModel";
 import { MessageType } from "../../shared/protocol";
-import { Range, getRangeSelectionsFromStack } from "../../shared/util/range";
+import { Range, RangeDirection, getRangeSelectionsFromStack } from "../../shared/util/range";
 import _style from "./dataDisplayContext.css";
-import { messageHandler, registerHandler } from "./state";
+import { getWebviewState, messageHandler, registerHandler, setWebviewState } from "./state";
 import { throwOnUndefinedAccessInDev } from "./util";
 
 const style = throwOnUndefinedAccessInDev(_style);
@@ -34,12 +34,18 @@ export class FocusedElement {
 	}
 }
 
+const selectionStateKey = "DisplayContextSelection";
+
+type SelectionState = { start: number; end: number; dir: RangeDirection }[];
+
 /**
  * Data management context component. Initially we used Recoil for this, but
  * this ended up introducing performance issues with very many components.
  */
 export class DisplayContext {
-	private _selection: Range[] = [];
+	private _selection: Range[] = getWebviewState<SelectionState>(selectionStateKey, []).map(
+		r => new Range(r.start, r.end, r.dir),
+	);
 	private _hoveredByte?: FocusedElement;
 	private _focusedByte?: FocusedElement;
 	private _unsavedRanges: readonly Range[] = [];
@@ -288,6 +294,15 @@ export class DisplayContext {
 			selected += range.size;
 		}
 
+		setWebviewState(
+			selectionStateKey,
+			this._selection.map(r => ({
+				start: r.start,
+				end: r.end,
+				dir: r.direction,
+			})) satisfies SelectionState,
+		);
+
 		messageHandler.sendEvent({
 			type: MessageType.SetSelectedCount,
 			selected: selected,
diff --git a/media/editor/hooks.ts b/media/editor/hooks.ts
index 8c6180d..356e0cc 100644
--- a/media/editor/hooks.ts
+++ b/media/editor/hooks.ts
@@ -44,10 +44,10 @@ export const usePersistedState = <T>(
 	key: string,
 	defaultValue: T,
 ): [T, React.Dispatch<React.SetStateAction<T>>] => {
-	const [value, setValue] = useState<T>(select.vscode.getState()?.[key] ?? defaultValue);
+	const [value, setValue] = useState<T>(select.getWebviewState(key, defaultValue));
 
 	useLazyEffect(() => {
-		select.vscode.setState({ ...select.vscode.getState(), [key]: value });
+		select.setWebviewState(key, value);
 	}, [value]);
 
 	return [value, setValue];
diff --git a/media/editor/state.ts b/media/editor/state.ts
index bad9444..05134f3 100644
--- a/media/editor/state.ts
+++ b/media/editor/state.ts
@@ -34,6 +34,14 @@ const acquireVsCodeApi: () => {
 
 export const vscode = acquireVsCodeApi?.();
 
+export const setWebviewState = (key: string, value: unknown) => {
+	vscode.setState?.({ ...vscode.getState(), [key]: value });
+};
+
+export const getWebviewState = <T>(key: string, defaultValue: T): T => {
+	return vscode.getState?.()[key] ?? defaultValue;
+};
+
 type HandlerFn = (message: ToWebviewMessage) => Promise<FromWebviewMessage> | undefined;
 
 const handles: { [T in ToWebviewMessage["type"]]?: HandlerFn | HandlerFn[] } = {};
@@ -469,7 +477,7 @@ export const decoratorsPage = selectorFamily({
 			const searcherByEnd = binarySearch<HexDecorator>(decorator => decorator.range.end);
 			const startIndex = searcherByEnd(pageSize * pageNumber, allDecorators);
 			const searcherByStart = binarySearch<HexDecorator>(d => d.range.start);
-			const endIndex = searcherByStart(pageSize * pageNumber + pageSize+1, allDecorators);
+			const endIndex = searcherByStart(pageSize * pageNumber + pageSize + 1, allDecorators);
 			return allDecorators.slice(startIndex, endIndex);
 		},
 });