diff --git a/app/frontend/src/common/yjs/event.ts b/app/frontend/src/common/yjs/event.ts index 8b3071c..0b4414b 100644 --- a/app/frontend/src/common/yjs/event.ts +++ b/app/frontend/src/common/yjs/event.ts @@ -28,6 +28,7 @@ function useComponentsListener() { useEffect(() => { const fn = (e: Y.YMapEvent) => { + console.log(e); for (const [cid, cc] of e.changes.keys) { switch (cc.action) { case 'add': { @@ -53,8 +54,8 @@ function useComponentsListener() { return; } - const {sliceIndex, entityId} = cc.oldValue; - transferComponent({sliceIndex, entityId, componentId: cid, targetEntityId: comp.eid}); + const {sidx, eid} = cc.oldValue as ComponentYjs; + transferComponent({sliceIndex: sidx, entityId: eid, componentId: cid, targetEntityId: comp.eid}); break; } case 'delete': { diff --git a/app/frontend/src/component/panel/ActionBar.tsx b/app/frontend/src/component/panel/ActionBar.tsx index 9c28502..af3f3de 100644 --- a/app/frontend/src/component/panel/ActionBar.tsx +++ b/app/frontend/src/component/panel/ActionBar.tsx @@ -28,7 +28,7 @@ import {ConfigContext} from 'common/context'; import {rectFitTransform} from 'common/geometry'; import {useInvertSelection, useFocusAreas, useDrawing} from 'common/hook'; import {useCanvasSize, leftSidebarWidth} from './layout'; -import {useTemporalAnnoStore} from 'state/annotate/annotation-temporal'; +import {useAnnoHistoryStore} from 'state/annotate/annotation-provider'; const ActionButton: FC<{helpCode: string; icon: React.ReactNode; hotKey?: string} & ButtonProps> = ({ helpCode, @@ -136,7 +136,10 @@ export const ActionBar: FC = ({...baseProps}) => { const focusAreas = useFocusAreas(canvasSize); // redo and undo - const {undoCount, redoCount, redo, undo} = useTemporalAnnoStore(); + const redo = useAnnoHistoryStore(s => s.redo); + const undo = useAnnoHistoryStore(s => s.undo); + const undoCount = useAnnoHistoryStore(s => s.index); + const redoCount = useAnnoHistoryStore(s => s.actions.length - s.index - 1); return (
diff --git a/app/frontend/src/component/panel/AnnotateLayer.tsx b/app/frontend/src/component/panel/AnnotateLayer.tsx index a816498..b10eb98 100644 --- a/app/frontend/src/component/panel/AnnotateLayer.tsx +++ b/app/frontend/src/component/panel/AnnotateLayer.tsx @@ -5,7 +5,7 @@ import intl from 'react-intl-universal'; import {useStore as useRenderStore} from 'state/annotate/render'; import {useStore as useUIStore} from 'state/annotate/ui'; -import {useTemporalAnnoStore} from 'state/annotate/annotation-temporal'; +import {useAnnoHistoryStore} from 'state/annotate/annotation-provider'; import {useStore as useDrawPolyStore} from 'state/annotate/polychain/draw'; import {useStore as useEditPolyStore} from 'state/annotate/polychain/edit'; import {useStore as useDrawRectStore} from 'state/annotate/rectangle/draw'; @@ -45,7 +45,7 @@ export const AnnotateLayer: FC> = ({...divProps}) const Loaded: FC> = ({...divProps}) => { // Reset history when a new annotation is started. - const {clear} = useTemporalAnnoStore(); + const clear = useAnnoHistoryStore(s => s.reset); useEffect(() => clear(), [clear]); const isDrawingPoly = useDrawPolyStore(s => s.vertices.length > 0); diff --git a/app/frontend/src/component/panel/layer/Idle.tsx b/app/frontend/src/component/panel/layer/Idle.tsx index 64e8806..6745dfa 100644 --- a/app/frontend/src/component/panel/layer/Idle.tsx +++ b/app/frontend/src/component/panel/layer/Idle.tsx @@ -4,7 +4,7 @@ import intl from 'react-intl-universal'; import {v4 as uuidv4} from 'uuid'; import {message} from 'antd'; -import {useTemporalAnnoStore} from 'state/annotate/annotation-temporal'; +import {useAnnoHistoryStore} from 'state/annotate/annotation-provider'; import {EntityComponentId, useStore as useRenderStore} from 'state/annotate/render'; import {useStore as useUIStore} from 'state/annotate/ui'; import {useStore as useEditPolyStore} from 'state/annotate/polychain/edit'; @@ -160,7 +160,8 @@ function useRenderSettings(): RenderSetting[] { } const IdleLayerTemporal: FC = () => { - const {undo, redo} = useTemporalAnnoStore(); + const undo = useAnnoHistoryStore(s => s.undo); + const redo = useAnnoHistoryStore(s => s.redo); useHotkeys('ctrl+z, meta+z', () => undo()); useHotkeys('ctrl+shift+z, meta+shift+z', () => redo()); return <>; diff --git a/app/frontend/src/state/annotate/annotation-broadcast.ts b/app/frontend/src/state/annotate/annotation-broadcast.ts index ea2937a..3fa7d9c 100644 --- a/app/frontend/src/state/annotate/annotation-broadcast.ts +++ b/app/frontend/src/state/annotate/annotation-broadcast.ts @@ -5,6 +5,8 @@ import type { MaskComponent, PolychainComponent, RectangleComponent, + Entity, + SliceIndex, } from 'type/annotation'; import {newComponentAdapter} from 'common/adapter'; import {initialVertexBezier} from 'common/annotation'; @@ -26,6 +28,7 @@ import { UpdatePolychainVerticesInput, UpdateRectangleAnchorsInput, UpdateSliceMasksInput, + getComponent, } from './annotation'; import {encodeEntityCategoryMapKey, yjsEntityCategoriesMap} from 'common/yjs/docs/entity'; import {useYjsContext} from 'common/yjs/context'; @@ -33,15 +36,39 @@ import {yjsComponentMap, Component as YjsComponent} from 'common/yjs/docs/compon import {yjsRectangleAnchorsMap} from 'common/yjs/docs/rectangle'; import {yjsPolychainVerticesMap} from 'common/yjs/docs/polychain'; import {yjsMaskMap} from 'common/yjs/docs/mask'; +import {useAnnoHistoryStore, useAnnoStoreRaw} from './annotation-provider'; +import {deepClone} from 'common/util'; export function useSetEntityCategory(): Pick { const {doc} = useYjsContext(); const cats = yjsEntityCategoriesMap(doc); + const annoStore = useAnnoStoreRaw(); + const pushAction = useAnnoHistoryStore(s => s.pushAction); const setEntityCategory = (input: SetEntityCategoryInput) => { const {sliceIndex: sidx, entityId: eid, category, entries} = input; + + // current + const entity = annoStore.getState().annotation.entities[eid]; + const oldEntries = getEntityCategory(entity, category, sidx); + + // update const key = encodeEntityCategoryMapKey({eid, sidx, category}); cats.set(key, entries); + + // history + const undo = () => { + setEntityCategory({ + sliceIndex: sidx, + entityId: eid, + category, + entries: oldEntries, + }); + }; + const redo = () => { + setEntityCategory(input); + }; + pushAction({undo, redo}); }; return {setEntityCategory}; @@ -51,10 +78,33 @@ export function useClearEntityCategory(): Pick s.pushAction); + const {setEntityCategory} = useSetEntityCategory(); const clearEntityCategory = (input: ClearEntityCategoryInput) => { const {sliceIndex: sidx, entityId: eid, category} = input; + + // current + const entity = annoStore.getState().annotation.entities[eid]; + const oldEntries = getEntityCategory(entity, category, sidx); + + // update const key = encodeEntityCategoryMapKey({eid, sidx, category}); cats.delete(key); + + // history + const undo = () => { + setEntityCategory({ + sliceIndex: sidx, + entityId: eid, + category, + entries: oldEntries, + }); + }; + const redo = () => { + clearEntityCategory(input); + }; + pushAction({undo, redo}); }; return {clearEntityCategory}; @@ -68,6 +118,8 @@ export function useSeparateComponent(): Pick s.pushAction); + const {transferComponent} = useTransferComponent(); const separateComponent = (input: SeparateComponentInput) => { const {sliceIndex: sidx, entityId: eid, componentId: cid, newEntityId, newComponentId} = input; @@ -127,6 +179,20 @@ export function useSeparateComponent(): Pick { + transferComponent({ + sliceIndex: sidx, + entityId: newEntityId, + componentId: newComponentId, + targetEntityId: eid, + }); + }; + const redo = () => { + // TODO(xu): check if this `newComponentId` indicates a bad design. + separateComponent({...input, componentId: newComponentId}); + }; + pushAction({undo, redo}); }; return {separateComponent}; @@ -135,8 +201,22 @@ export function useSeparateComponent(): Pick { const {doc} = useYjsContext(); + const pushAction = useAnnoHistoryStore(s => s.pushAction); + const {deleteComponents} = useDeleteComponents(); const addComponent = (input: AddComponentInput) => { + const {sliceIndex, entityId, component} = input; addComponentToDoc(doc, input); + + const undo = () => { + deleteComponents({ + sliceIndex, + components: [[entityId, component.id]], + }); + }; + const redo = () => { + addComponent(input); + }; + pushAction({undo, redo}); }; return {addComponent}; @@ -145,11 +225,26 @@ export function useAddComponent(): Pick { export function useAddComponents(): Pick { const {doc} = useYjsContext(); + const pushAction = useAnnoHistoryStore(s => s.pushAction); + const {deleteComponents} = useDeleteComponents(); const addComponents = (input: AddComponentsInput) => { const {entityId, components} = input; doc.transact(() => { components.forEach(component => addComponentToDoc(doc, {entityId, ...component})); }); + + const undo = () => { + components.forEach(({sliceIndex, component}) => { + deleteComponents({ + sliceIndex, + components: [[entityId, component.id]], + }); + }); + }; + const redo = () => { + addComponents(input); + }; + pushAction({undo, redo}); }; return {addComponents}; @@ -159,9 +254,30 @@ export function useUpdatePolychainVertices(): Pick s.pushAction); const updatePolychainVertices = (input: UpdatePolychainVerticesInput) => { - const {componentId: cid, vertices} = input; + const {componentId: cid, vertices, sliceIndex, entityId} = input; + const old = getComponent(annoStore.getState(), sliceIndex, entityId, cid); + if (old?.type !== 'polychain') { + return; + } + const oldVertices = deepClone(old.vertices); + verts.set(cid, Y.Array.from(vertices)); + + const undo = () => { + updatePolychainVertices({ + sliceIndex, + entityId, + componentId: cid, + vertices: oldVertices, + }); + }; + const redo = () => { + updatePolychainVertices(input); + }; + pushAction({undo, redo}); }; return {updatePolychainVertices}; @@ -171,10 +287,20 @@ export function useTransferComponent(): Pick s.pushAction); + const {addComponent} = useAddComponent(); + const {deleteComponents} = useDeleteComponents(); const transferComponent = (input: TransferComponentInput) => { const {sliceIndex: sidx, entityId: eid, componentId: cid, targetEntityId} = input; + + const c = deepClone(getComponent(annoStore.getState(), sidx, eid, cid)); + if (!c) { + return; + } + doc.transact(() => { - const old = comps.get(cid); + const old = deepClone(comps.get(cid)); if (!old) { console.warn(`component ${cid} not found`); return; @@ -186,6 +312,22 @@ export function useTransferComponent(): Pick { + deleteComponents({ + sliceIndex: sidx, + components: [[targetEntityId, cid]], + }); + addComponent({ + sliceIndex: sidx, + entityId: eid, + component: c, + }); + }; + const redo = () => { + transferComponent(input); + }; + pushAction({undo, redo}); }; return {transferComponent}; @@ -195,6 +337,10 @@ export function useDeleteEntities(): Pick { const {doc} = useYjsContext(); const comps = yjsComponentMap(doc); + const annoStore = useAnnoStoreRaw(); + const {addComponents} = useAddComponents(); + const {setEntityCategory} = useSetEntityCategory(); + const pushAction = useAnnoHistoryStore(s => s.pushAction); const deleteEntities = (input: DeleteEntitiesInput) => { const {entityIds} = input; const eids = new Set(entityIds); @@ -208,6 +354,52 @@ export function useDeleteEntities(): Pick { for (const [cid, comp] of victims.entries()) { deleteComponentFromDoc(doc, cid, comp.type); } + + // history + const s = annoStore.getState(); + const deleted = entityIds.map(eid => s.annotation.entities[eid]); + const entities = deepClone(deleted); + const undo = () => { + // add back components + const inputs: AddComponentsInput[] = []; + entities.forEach(e => { + Object.entries(e.geometry.slices).map(([sidx, sliceComponents]) => { + const components = Object.values(sliceComponents).map(component => ({ + sliceIndex: parseInt(sidx), + component, + })); + inputs.push({entityId: e.id, components}); + }); + }); + inputs.forEach(input => { + addComponents(input); + }); + + // add back categories + entities.forEach(e => { + Object.entries(e.globalCategories ?? {}).forEach(([category, es]) => { + setEntityCategory({ + entityId: e.id, + category, + entries: Object.keys(es), + }); + }); + Object.entries(e.sliceCategories ?? {}).forEach(([sidx, sliceCategories]) => { + Object.entries(sliceCategories).forEach(([category, es]) => { + setEntityCategory({ + entityId: e.id, + sliceIndex: parseInt(sidx), + category, + entries: Object.keys(es), + }); + }); + }); + }); + }; + const redo = () => { + deleteEntities(input); + }; + pushAction({undo, redo}); }); }; @@ -218,10 +410,18 @@ export function useTruncateEntities(): Pick s.pushAction); const truncateEntities = (input: TruncateEntitiesInput) => { // TODO(xu): can be inefficient when there are too many components const {entityIds, sinceSliceIndex} = input; + const s = annoStore.getState(); + const entities = deepClone(entityIds.map(eid => s.annotation.entities[eid])); + + // update const cids: string[] = []; for (const [cid, comp] of comps.entries()) { if (!entityIds.includes(comp.eid)) { @@ -243,6 +443,50 @@ export function useTruncateEntities(): Pick { + // add back components + const inputs: AddComponentsInput[] = []; + entities.forEach(e => { + Object.entries(e.geometry.slices).map(([sidxStr, sliceComponents]) => { + const sidx = parseInt(sidxStr); + if (sidx < sinceSliceIndex) { + return; + } + const components = Object.values(sliceComponents).map(component => ({ + sliceIndex: sidx, + component, + })); + inputs.push({entityId: e.id, components}); + }); + }); + inputs.forEach(input => { + addComponents(input); + }); + + // add back categories + entities.forEach(e => { + Object.entries(e.sliceCategories ?? {}).forEach(([sidxStr, sliceCategories]) => { + const sidx = parseInt(sidxStr); + if (sidx < sinceSliceIndex) { + return; + } + Object.entries(sliceCategories).forEach(([category, es]) => { + setEntityCategory({ + entityId: e.id, + sliceIndex: sidx, + category, + entries: Object.keys(es), + }); + }); + }); + }); + }; + const redo = () => { + truncateEntities(input); + }; + pushAction({undo, redo}); }; return {truncateEntities}; @@ -549,3 +793,13 @@ function deleteComponentFromDoc(doc: Y.Doc, cid: ComponentId, ctype: YjsComponen break; } } + +export function getEntityCategory(e: Entity, category: string, sliceIndex?: SliceIndex): string[] { + if (sliceIndex !== undefined) { + const val = e.sliceCategories?.[sliceIndex]?.[category]; + return Object.keys(val ?? {}); + } + + const val = e.globalCategories?.[category]; + return Object.keys(val ?? {}); +} diff --git a/app/frontend/src/state/annotate/annotation-history.tsx b/app/frontend/src/state/annotate/annotation-history.tsx new file mode 100644 index 0000000..1d84822 --- /dev/null +++ b/app/frontend/src/state/annotate/annotation-history.tsx @@ -0,0 +1,82 @@ +import {createStore} from 'zustand'; +import {immer} from 'zustand/middleware/immer'; + +type VoidAction = () => void; + +interface ActionConjugation { + undo: VoidAction; + redo: VoidAction; +} + +const noop = () => {}; +const initialAction: ActionConjugation = { + undo: noop, + redo: noop, +}; + +export type HistoryState = { + actions: ActionConjugation[]; + index: number; + + pushAction: (action: ActionConjugation) => void; + undo: () => void; + redo: () => void; + reset: () => void; +}; + +export function createAnnoHistoryStore() { + return createStore()( + immer(set => ({ + actions: [initialAction], + index: 0, + + pushAction: (action: ActionConjugation) => { + set(s => { + if (s.index === s.actions.length - 1) { + s.actions.push(action); + } else { + s.actions = [initialAction, action]; + } + s.index = s.actions.length - 1; + }); + }, + + undo: () => { + set(s => { + if (s.index === 0) { + return; + } + const {undo: fn} = s.actions[s.index]; + s.index--; + fn(); + }); + }, + + redo: () => { + set(s => { + if (s.index === s.actions.length - 1) { + return; + } + const {redo: fn} = s.actions[s.index + 1]; + s.index++; + fn(); + }); + }, + + reset: () => { + set(s => { + s.actions = [initialAction]; + s.index = 0; + }); + }, + })) + ); +} + +export interface AnnoHistoryStore { + undo: () => void; + redo: () => void; + clear: () => void; + redoCount: number; + undoCount: number; +} diff --git a/app/frontend/src/state/annotate/annotation-provider.tsx b/app/frontend/src/state/annotate/annotation-provider.tsx index 89ae367..07382cc 100644 --- a/app/frontend/src/state/annotate/annotation-provider.tsx +++ b/app/frontend/src/state/annotate/annotation-provider.tsx @@ -1,21 +1,32 @@ import {ReactNode, createContext, useContext} from 'react'; import {State, createAnnoStore} from './annotation'; import {useStore} from 'zustand'; +import {createAnnoHistoryStore, HistoryState} from './annotation-history'; -type T = ReturnType; -const Context = createContext(undefined!); +const AnnoContext = createContext>(undefined!); +const AnnoHistoryContext = createContext>(undefined!); export function useAnnoStoreRaw() { - const store = useContext(Context); + const store = useContext(AnnoContext); return store; } export function useAnnoStore(selector: (state: State) => T, equalityFn?: (a: T, b: T) => boolean) { - const store = useContext(Context); + const store = useContext(AnnoContext); + return useStore(store, selector, equalityFn); +} + +export function useAnnoHistoryStore(selector: (state: HistoryState) => T, equalityFn?: (a: T, b: T) => boolean) { + const store = useContext(AnnoHistoryContext); return useStore(store, selector, equalityFn); } export function AnnoProvider({children}: {children: ReactNode}): JSX.Element { - const store = createAnnoStore(); - return {children}; + const annoStore = createAnnoStore(); + const annoHistoryStore = createAnnoHistoryStore(); + return ( + + {children} + + ); } diff --git a/app/frontend/src/state/annotate/annotation-temporal.tsx b/app/frontend/src/state/annotate/annotation-temporal.tsx deleted file mode 100644 index e68cb59..0000000 --- a/app/frontend/src/state/annotate/annotation-temporal.tsx +++ /dev/null @@ -1,96 +0,0 @@ -import create from 'zustand'; -import {immer} from 'zustand/middleware/immer'; - -type VoidAction = () => void; - -interface ActionConjugation { - undo: VoidAction; - redo: VoidAction; -} - -const noop = () => {}; -const initialAction: ActionConjugation = { - undo: noop, - redo: noop, -}; - -type HistoryState = { - actions: ActionConjugation[]; - index: number; - - pushAction: (action: ActionConjugation) => void; - undo: () => void; - redo: () => void; - reset: () => void; -}; - -const useHistoryStore = create()( - immer(set => ({ - actions: [initialAction], - index: 0, - - pushAction: (action: ActionConjugation) => { - set(s => { - if (s.index === s.actions.length - 1) { - s.actions.push(action); - } else { - s.actions = [initialAction, action]; - } - s.index = s.actions.length - 1; - }); - }, - - undo: () => { - set(s => { - if (s.index === 0) { - return; - } - const {undo: fn} = s.actions[s.index]; - s.index--; - fn(); - }); - }, - - redo: () => { - set(s => { - if (s.index === s.actions.length - 1) { - return; - } - const {redo: fn} = s.actions[s.index + 1]; - s.index++; - fn(); - }); - }, - - reset: () => { - set(s => { - s.actions = [initialAction]; - s.index = 0; - }); - }, - })) -); - -export interface TemporalStore { - undo: () => void; - redo: () => void; - clear: () => void; - redoCount: number; - undoCount: number; -} - -export function useTemporalAnnoStore(): TemporalStore { - const undo = useHistoryStore(s => s.undo); - const redo = useHistoryStore(s => s.redo); - const undoCount = useHistoryStore(s => s.index); - const redoCount = useHistoryStore(s => s.actions.length - s.index - 1); - const clear = useHistoryStore(s => s.reset); - - return { - undo, - redo, - clear, - redoCount, - undoCount, - }; -}