fix(ui): edge cases in quick switch, simpler logic

invokeai/frontend/web/src/features/controlLayers/components/common/CanvasEntityIsBookmarkedForQuickSwitchToggle.tsx

@@ -2,7 +2,7 @@ import { IconButton } from '@invoke-ai/ui-library';
 import { useAppDispatch } from 'app/store/storeHooks';
 import { useEntityIdentifierContext } from 'features/controlLayers/contexts/EntityIdentifierContext';
 import { useEntityIsBookmarkedForQuickSwitch } from 'features/controlLayers/hooks/useEntityIsBookmarkedForQuickSwitch';
-import { entityIsBookmarkedForQuickSwitchChanged } from 'features/controlLayers/store/canvasSlice';
+import { bookmarkedEntityChanged } from 'features/controlLayers/store/canvasSlice';
 import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiBookmarkSimpleBold, PiBookmarkSimpleFill } from 'react-icons/pi';
@@ -14,9 +14,9 @@ export const CanvasEntityIsBookmarkedForQuickSwitchToggle = memo(() => {
   const dispatch = useAppDispatch();
   const onClick = useCallback(() => {
     if (isBookmarked) {
-      dispatch(entityIsBookmarkedForQuickSwitchChanged({ entityIdentifier: null }));
+      dispatch(bookmarkedEntityChanged({ entityIdentifier: null }));
     } else {
-      dispatch(entityIsBookmarkedForQuickSwitchChanged({ entityIdentifier }));
+      dispatch(bookmarkedEntityChanged({ entityIdentifier }));
     }
   }, [dispatch, entityIdentifier, isBookmarked]);
 

invokeai/frontend/web/src/features/controlLayers/hooks/useCanvasEntityQuickSwitchHotkey.ts

@@ -1,83 +1,47 @@
-import { useStore } from '@nanostores/react';
-import { createSelector } from '@reduxjs/toolkit';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import { useAssertSingleton } from 'common/hooks/useAssertSingleton';
 import { entitySelected } from 'features/controlLayers/store/canvasSlice';
 import {
-  selectCanvasSlice,
-  selectEntity,
-  selectQuickSwitchEntityIdentifier,
+  selectBookmarkedEntityIdentifier,
   selectSelectedEntityIdentifier,
 } from 'features/controlLayers/store/selectors';
 import type { CanvasEntityIdentifier } from 'features/controlLayers/store/types';
-import { atom } from 'nanostores';
-import { useCallback, useEffect, useMemo } from 'react';
+import { useCallback, useEffect, useState } from 'react';
 import { useHotkeys } from 'react-hotkeys-hook';
 
-const $selectedEntityBuffer = atom<CanvasEntityIdentifier | null>(null);
-
 export const useCanvasEntityQuickSwitchHotkey = () => {
-  useAssertSingleton('useCanvasEntityQuickSwitch');
-
   const dispatch = useAppDispatch();
-  const selectedEntityBuffer = useStore($selectedEntityBuffer);
-  const quickSwitchEntityIdentifier = useAppSelector(selectQuickSwitchEntityIdentifier);
-  const selectedEntityIdentifier = useAppSelector(selectSelectedEntityIdentifier);
-  const selectBufferEntityIdentifier = useMemo(
-    () =>
-      createSelector(selectCanvasSlice, (canvas) =>
-        selectedEntityBuffer ? (selectEntity(canvas, selectedEntityBuffer) ?? null) : null
-      ),
-    [selectedEntityBuffer]
-  );
-  const bufferEntityIdentifier = useAppSelector(selectBufferEntityIdentifier);
-
-  const quickSwitch = useCallback(() => {
-    if (quickSwitchEntityIdentifier !== null) {
-      // If there is a quick switch entity, we should switch between it and the buffer
-      if (quickSwitchEntityIdentifier.id !== selectedEntityIdentifier?.id) {
-        // The quick switch entity is not selected - select it
-        dispatch(entitySelected({ entityIdentifier: quickSwitchEntityIdentifier }));
-        $selectedEntityBuffer.set(selectedEntityIdentifier);
-      } else if (bufferEntityIdentifier !== null) {
-        // The quick switch entity is already selected - select the buffer
-        dispatch(entitySelected({ entityIdentifier: bufferEntityIdentifier }));
-        $selectedEntityBuffer.set(quickSwitchEntityIdentifier);
-      }
-    } else {
-      // No quick switch entity, so we should switch between buffer and selected entity
-      // If there is no selected entity or buffer, we should not do anything
-      if (selectedEntityIdentifier === null && bufferEntityIdentifier === null) {
-        return;
-      }
-      // If there is no selected entity but we do have a buffer, we should select the buffer
-      if (selectedEntityIdentifier === null && bufferEntityIdentifier !== null) {
-        dispatch(entitySelected({ entityIdentifier: bufferEntityIdentifier }));
-        return;
-      }
-      // If there is a selected entity but no buffer, we should buffer the selected entity
-      if (selectedEntityIdentifier !== null && bufferEntityIdentifier === null) {
-        $selectedEntityBuffer.set(selectedEntityIdentifier);
-        return;
-      }
-      // If there is a selected entity and a buffer, and they are different, we should swap the selected entity and the buffer
-      if (
-        selectedEntityIdentifier !== null &&
-        bufferEntityIdentifier !== null &&
-        selectedEntityIdentifier.id !== bufferEntityIdentifier.id
-      ) {
-        $selectedEntityBuffer.set(selectedEntityIdentifier);
-        dispatch(entitySelected({ entityIdentifier: bufferEntityIdentifier }));
-        return;
-      }
-    }
-  }, [bufferEntityIdentifier, dispatch, quickSwitchEntityIdentifier, selectedEntityIdentifier]);
+  const [prev, setPrev] = useState<CanvasEntityIdentifier | null>(null);
+  const [current, setCurrent] = useState<CanvasEntityIdentifier | null>(null);
+  const selected = useAppSelector(selectSelectedEntityIdentifier);
+  const bookmarked = useAppSelector(selectBookmarkedEntityIdentifier);
 
+  // Update prev and current when selected entity changes
   useEffect(() => {
-    if (!bufferEntityIdentifier) {
-      $selectedEntityBuffer.set(null);
+    if (current?.id !== selected?.id) {
+      setPrev(current);
+      setCurrent(selected);
     }
-  }, [bufferEntityIdentifier]);
+  }, [current, selected]);
 
-  useHotkeys('q', quickSwitch, { enabled: true, preventDefault: true }, [quickSwitch]);
+  const onQuickSwitch = useCallback(() => {
+    if (bookmarked) {
+      if (current?.id !== bookmarked.id) {
+        // Switch between current (non-bookmarked) and bookmarked
+        setPrev(current);
+        setCurrent(bookmarked);
+        dispatch(entitySelected({ entityIdentifier: bookmarked }));
+      } else if (prev) {
+        // Switch back to the last non-bookmarked entity
+        setCurrent(prev);
+        dispatch(entitySelected({ entityIdentifier: prev }));
+      }
+    } else if (prev !== null && current !== null) {
+      // Switch between prev and current if no bookmarked entity
+      setPrev(current);
+      setCurrent(prev);
+      dispatch(entitySelected({ entityIdentifier: prev }));
+    }
+  }, [bookmarked, current, dispatch, prev]);
+
+  useHotkeys('q', onQuickSwitch);
 };

invokeai/frontend/web/src/features/controlLayers/hooks/useEntityIsBookmarkedForQuickSwitch.ts

@@ -8,7 +8,7 @@ export const useEntityIsBookmarkedForQuickSwitch = (entityIdentifier: CanvasEnti
   const selectIsBookmarkedForQuickSwitch = useMemo(
     () =>
       createSelector(selectCanvasSlice, (canvas) => {
-        return canvas.quickSwitchEntityIdentifier?.id === entityIdentifier.id;
+        return canvas.bookmarkedEntityIdentifier?.id === entityIdentifier.id;
       }),
     [entityIdentifier]
   );

invokeai/frontend/web/src/features/controlLayers/store/canvasSlice.ts

@@ -84,7 +84,7 @@ const getRGMaskFill = (state: CanvasState): RgbColor => {
 const initialState: CanvasState = {
   _version: 3,
   selectedEntityIdentifier: null,
-  quickSwitchEntityIdentifier: null,
+  bookmarkedEntityIdentifier: null,
   rasterLayers: {
     isHidden: false,
     entities: [],
@@ -792,13 +792,10 @@ export const canvasSlice = createSlice({
       }
       state.selectedEntityIdentifier = entityIdentifier;
     },
-    entityIsBookmarkedForQuickSwitchChanged: (
-      state,
-      action: PayloadAction<{ entityIdentifier: CanvasEntityIdentifier | null }>
-    ) => {
+    bookmarkedEntityChanged: (state, action: PayloadAction<{ entityIdentifier: CanvasEntityIdentifier | null }>) => {
       const { entityIdentifier } = action.payload;
       if (!entityIdentifier) {
-        state.quickSwitchEntityIdentifier = null;
+        state.bookmarkedEntityIdentifier = null;
         return;
       }
       const entity = selectEntity(state, entityIdentifier);
@@ -806,7 +803,7 @@ export const canvasSlice = createSlice({
         // Cannot select a non-existent entity
         return;
       }
-      state.quickSwitchEntityIdentifier = entityIdentifier;
+      state.bookmarkedEntityIdentifier = entityIdentifier;
     },
     entityNameChanged: (state, action: PayloadAction<EntityIdentifierPayload<{ name: string | null }>>) => {
       const { entityIdentifier, name } = action.payload;
@@ -1122,7 +1119,7 @@ export const {
   canvasClearHistory,
   // All entities
   entitySelected,
-  entityIsBookmarkedForQuickSwitchChanged,
+  bookmarkedEntityChanged,
   entityNameChanged,
   entityReset,
   entityIsEnabledToggled,

invokeai/frontend/web/src/features/controlLayers/store/selectors.ts

@@ -181,9 +181,9 @@ export const selectSelectedEntityIdentifier = createSelector(
   (canvas) => canvas.selectedEntityIdentifier
 );
 
-export const selectQuickSwitchEntityIdentifier = createSelector(
+export const selectBookmarkedEntityIdentifier = createSelector(
   selectCanvasSlice,
-  (canvas) => canvas.quickSwitchEntityIdentifier
+  (canvas) => canvas.bookmarkedEntityIdentifier
 );
 
 export const selectIsSelectedEntityDrawable = createSelector(

invokeai/frontend/web/src/features/controlLayers/store/types.ts

@@ -688,7 +688,7 @@ export type StagingAreaImage = {
 export type CanvasState = {
   _version: 3;
   selectedEntityIdentifier: CanvasEntityIdentifier | null;
-  quickSwitchEntityIdentifier: CanvasEntityIdentifier | null;
+  bookmarkedEntityIdentifier: CanvasEntityIdentifier | null;
   inpaintMasks: {
     isHidden: boolean;
     entities: CanvasInpaintMaskState[];