import type { DragEndEvent, DragMoveEvent, DragStartEvent } from '@dnd-kit/core'; import { DndContext, MeasuringStrategy, PointerSensor, pointerWithin, useDroppable, useSensor, useSensors, } from '@dnd-kit/core'; import { type } from '@tauri-apps/plugin-os'; import classNames from 'classnames'; import type { ComponentType, MouseEvent, ReactElement, Ref, RefAttributes } from 'react'; import { forwardRef, memo, useCallback, useEffect, useImperativeHandle, useMemo, useRef, useState, } from 'react'; import { useKey, useKeyPressEvent } from 'react-use'; import type { HotKeyOptions, HotkeyAction } from '../../../hooks/useHotKey'; import { useHotKey } from '../../../hooks/useHotKey'; import { computeSideForDragMove } from '../../../lib/dnd'; import { jotaiStore } from '../../../lib/jotai'; import type { ContextMenuProps, DropdownItem } from '../Dropdown'; import { ContextMenu } from '../Dropdown'; import { collapsedFamily, draggingIdsFamily, focusIdsFamily, hoveredParentFamily, isCollapsedFamily, selectedIdsFamily, } from './atoms'; import type { SelectableTreeNode, TreeNode } from './common'; import { closestVisibleNode, equalSubtree, getSelectedItems, hasAncestor } from './common'; import { TreeDragOverlay } from './TreeDragOverlay'; import type { TreeItemClickEvent, TreeItemHandle, TreeItemProps } from './TreeItem'; import type { TreeItemListProps } from './TreeItemList'; import { TreeItemList } from './TreeItemList'; import { useSelectableItems } from './useSelectableItems'; /** So we re-calculate after expanding a folder during drag */ const measuring = { droppable: { strategy: MeasuringStrategy.Always } }; export interface TreeProps { root: TreeNode; treeId: string; getItemKey: (item: T) => string; getContextMenu?: (items: T[]) => ContextMenuProps['items'] | Promise; ItemInner: ComponentType<{ treeId: string; item: T }>; ItemLeftSlotInner?: ComponentType<{ treeId: string; item: T }>; ItemRightSlot?: ComponentType<{ treeId: string; item: T }>; className?: string; onActivate?: (item: T) => void; onDragEnd?: (opt: { items: T[]; parent: T; children: T[]; insertAt: number }) => void; hotkeys?: { actions: Partial void } & HotKeyOptions>>; }; getEditOptions?: (item: T) => { defaultValue: string; placeholder?: string; onChange: (item: T, text: string) => void; }; } export interface TreeHandle { treeId: string; focus: () => boolean; hasFocus: () => boolean; selectItem: (id: string, focus?: boolean) => void; renameItem: (id: string) => void; showContextMenu: () => void; } function TreeInner( { className, getContextMenu, getEditOptions, getItemKey, hotkeys, onActivate, onDragEnd, ItemInner, ItemLeftSlotInner, ItemRightSlot, root, treeId, }: TreeProps, ref: Ref, ) { const treeRef = useRef(null); const selectableItems = useSelectableItems(root); const [showContextMenu, setShowContextMenu] = useState<{ items: DropdownItem[]; x: number; y: number; } | null>(null); const treeItemRefs = useRef>({}); const handleAddTreeItemRef = useCallback((item: T, r: TreeItemHandle | null) => { if (r == null) { delete treeItemRefs.current[item.id]; } else { treeItemRefs.current[item.id] = r; } }, []); // Select the first item on first render // biome-ignore lint/correctness/useExhaustiveDependencies: Only used for initial render useEffect(() => { const ids = jotaiStore.get(selectedIdsFamily(treeId)); const fallback = selectableItems[0]; if (ids.length === 0 && fallback != null) { jotaiStore.set(selectedIdsFamily(treeId), [fallback.node.item.id]); jotaiStore.set(focusIdsFamily(treeId), { anchorId: fallback.node.item.id, lastId: fallback.node.item.id, }); } }, [treeId]); const handleCloseContextMenu = useCallback(() => { setShowContextMenu(null); }, []); const isTreeFocused = useCallback(() => { return treeRef.current?.contains(document.activeElement); }, []); const tryFocus = useCallback(() => { const $el = treeRef.current?.querySelector( '.tree-item button[tabindex="0"]', ); if ($el == null) { return false; } $el.focus(); $el.scrollIntoView({ block: 'nearest' }); return true; }, []); const ensureTabbableItem = useCallback(() => { const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId; const lastSelectedItem = selectableItems.find( (i) => i.node.item.id === lastSelectedId && !i.node.hidden, ); // If no item found, default to selecting the first item (prefer leaf node); if (lastSelectedItem == null) { const firstLeafItem = selectableItems.find((i) => !i.node.hidden && i.node.children == null); const firstItem = firstLeafItem ?? selectableItems.find((i) => !i.node.hidden); if (firstItem != null) { const id = firstItem.node.item.id; jotaiStore.set(selectedIdsFamily(treeId), [id]); jotaiStore.set(focusIdsFamily(treeId), { anchorId: id, lastId: id }); } return; } const closest = closestVisibleNode(treeId, lastSelectedItem.node); if (closest != null) { const id = closest.item.id; jotaiStore.set(selectedIdsFamily(treeId), [id]); jotaiStore.set(focusIdsFamily(treeId), { anchorId: id, lastId: id }); } }, [selectableItems, treeId]); // Ensure there's always a tabbable item after collapsed state changes useEffect(() => { const unsub = jotaiStore.sub(collapsedFamily(treeId), ensureTabbableItem); return unsub; }, [ensureTabbableItem, treeId]); // Ensure there's always a tabbable item after render useEffect(() => { requestAnimationFrame(ensureTabbableItem); }); const hasFocus = useCallback(() => { return treeRef.current?.contains(document.activeElement) ?? false; }, []); const setSelected = useCallback( (ids: string[], focus: boolean) => { jotaiStore.set(selectedIdsFamily(treeId), ids); // TODO: Figure out a better way than timeout if (!focus) return; setTimeout(tryFocus, 50); }, [treeId, tryFocus], ); const treeHandle = useMemo( () => ({ treeId, focus: tryFocus, hasFocus: hasFocus, renameItem: (id) => treeItemRefs.current[id]?.rename(), selectItem: (id, focus) => { if (jotaiStore.get(selectedIdsFamily(treeId)).includes(id)) { // Already selected return; } jotaiStore.set(focusIdsFamily(treeId), { anchorId: id, lastId: id }); setSelected([id], focus === true); }, showContextMenu: async () => { if (getContextMenu == null) return; const items = getSelectedItems(treeId, selectableItems); const menuItems = await getContextMenu(items); const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId; const rect = lastSelectedId ? treeItemRefs.current[lastSelectedId]?.rect() : null; if (rect == null) return; setShowContextMenu({ items: menuItems, x: rect.x, y: rect.y }); }, }), [getContextMenu, hasFocus, selectableItems, setSelected, treeId, tryFocus], ); useImperativeHandle(ref, (): TreeHandle => treeHandle, [treeHandle]); const handleGetContextMenu = useMemo(() => { if (getContextMenu == null) return; return (item: T) => { const items = getSelectedItems(treeId, selectableItems); const isSelected = items.find((i) => i.id === item.id); if (isSelected) { // If right-clicked an item that was in the multiple-selection, use the entire selection return getContextMenu(items); } // If right-clicked an item that was NOT in the multiple-selection, just use that one // Also update the selection with it setSelected([item.id], false); jotaiStore.set(focusIdsFamily(treeId), (prev) => ({ ...prev, lastId: item.id })); return getContextMenu([item]); }; }, [getContextMenu, selectableItems, setSelected, treeId]); const handleSelect = useCallback['onClick']>>( (item, { shiftKey, metaKey, ctrlKey }) => { const anchorSelectedId = jotaiStore.get(focusIdsFamily(treeId)).anchorId; const selectedIdsAtom = selectedIdsFamily(treeId); const selectedIds = jotaiStore.get(selectedIdsAtom); // Mark the item as the last one selected jotaiStore.set(focusIdsFamily(treeId), (prev) => ({ ...prev, lastId: item.id })); if (shiftKey) { const validSelectableItems = getValidSelectableItems(treeId, selectableItems); const anchorIndex = validSelectableItems.findIndex( (i) => i.node.item.id === anchorSelectedId, ); const currIndex = validSelectableItems.findIndex((v) => v.node.item.id === item.id); // Nothing was selected yet, so just select this item if (selectedIds.length === 0 || anchorIndex === -1 || currIndex === -1) { setSelected([item.id], true); jotaiStore.set(focusIdsFamily(treeId), (prev) => ({ ...prev, anchorId: item.id })); return; } if (currIndex > anchorIndex) { // Selecting down const itemsToSelect = validSelectableItems.slice(anchorIndex, currIndex + 1); setSelected( itemsToSelect.map((v) => v.node.item.id), true, ); } else if (currIndex < anchorIndex) { // Selecting up const itemsToSelect = validSelectableItems.slice(currIndex, anchorIndex + 1); setSelected( itemsToSelect.map((v) => v.node.item.id), true, ); } else { setSelected([item.id], true); } } else if (type() === 'macos' ? metaKey : ctrlKey) { const withoutCurr = selectedIds.filter((id) => id !== item.id); if (withoutCurr.length === selectedIds.length) { // It wasn't in there, so add it setSelected([...selectedIds, item.id], true); } else { // It was in there, so remove it setSelected(withoutCurr, true); } } else { // Select single setSelected([item.id], true); jotaiStore.set(focusIdsFamily(treeId), (prev) => ({ ...prev, anchorId: item.id })); } }, [selectableItems, setSelected, treeId], ); const handleClick = useCallback['onClick']>>( (item, e) => { if (e.shiftKey || e.ctrlKey || e.metaKey) { handleSelect(item, e); } else { handleSelect(item, e); onActivate?.(item); } }, [handleSelect, onActivate], ); const selectPrevItem = useCallback( (e: TreeItemClickEvent) => { const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId; const validSelectableItems = getValidSelectableItems(treeId, selectableItems); const index = validSelectableItems.findIndex((i) => i.node.item.id === lastSelectedId); const item = validSelectableItems[index - 1]; if (item != null) { handleSelect(item.node.item, e); } }, [handleSelect, selectableItems, treeId], ); const selectNextItem = useCallback( (e: TreeItemClickEvent) => { const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId; const validSelectableItems = getValidSelectableItems(treeId, selectableItems); const index = validSelectableItems.findIndex((i) => i.node.item.id === lastSelectedId); const item = validSelectableItems[index + 1]; if (item != null) { handleSelect(item.node.item, e); } }, [handleSelect, selectableItems, treeId], ); const selectParentItem = useCallback( (e: TreeItemClickEvent) => { const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId; const lastSelectedItem = selectableItems.find((i) => i.node.item.id === lastSelectedId)?.node ?? null; if (lastSelectedItem?.parent != null) { handleSelect(lastSelectedItem.parent.item, e); } }, [handleSelect, selectableItems, treeId], ); useKey( (e) => e.key === 'ArrowUp' || e.key.toLowerCase() === 'k', (e) => { if (!isTreeFocused()) return; e.preventDefault(); selectPrevItem(e); }, undefined, [selectableItems, handleSelect], ); useKey( (e) => e.key === 'ArrowDown' || e.key.toLowerCase() === 'j', (e) => { if (!isTreeFocused()) return; e.preventDefault(); selectNextItem(e); }, undefined, [selectableItems, handleSelect], ); // If the selected item is a collapsed folder, expand it. Otherwise, select next item useKey( (e) => e.key === 'ArrowRight' || e.key === 'l', (e) => { if (!isTreeFocused()) return; e.preventDefault(); const collapsed = jotaiStore.get(collapsedFamily(treeId)); const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId; const lastSelectedItem = selectableItems.find((i) => i.node.item.id === lastSelectedId); if ( lastSelectedId && lastSelectedItem?.node.children != null && collapsed[lastSelectedItem.node.item.id] === true ) { jotaiStore.set(isCollapsedFamily({ treeId, itemId: lastSelectedId }), false); } else { selectNextItem(e); } }, undefined, [selectableItems, handleSelect], ); // If the selected item is in a folder, select its parent. // If the selected item is an expanded folder, collapse it. useKey( (e) => e.key === 'ArrowLeft' || e.key === 'h', (e) => { if (!isTreeFocused()) return; e.preventDefault(); const collapsed = jotaiStore.get(collapsedFamily(treeId)); const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId; const lastSelectedItem = selectableItems.find((i) => i.node.item.id === lastSelectedId); if ( lastSelectedId && lastSelectedItem?.node.children != null && collapsed[lastSelectedItem.node.item.id] !== true ) { jotaiStore.set(isCollapsedFamily({ treeId, itemId: lastSelectedId }), true); } else { selectParentItem(e); } }, { options: {} }, [selectableItems, handleSelect], ); useKeyPressEvent('Escape', async () => { if (!treeRef.current?.contains(document.activeElement)) return; clearDragState(); const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId; if (lastSelectedId == null) return; setSelected([lastSelectedId], false); }); const handleDragMove = useCallback( function handleDragMove(e: DragMoveEvent) { const over = e.over; if (!over) { // Clear the drop indicator when hovering outside the tree jotaiStore.set(hoveredParentFamily(treeId), { parentId: null, parentDepth: null, childIndex: null, index: null, }); return; } // Not sure when or if this happens if (e.active.rect.current.initial == null) { return; } // Root is anything past the end of the list, so set it to the end const hoveringRoot = over.id === root.item.id; if (hoveringRoot) { jotaiStore.set(hoveredParentFamily(treeId), { parentId: root.item.id, parentDepth: root.depth, index: selectableItems.length, childIndex: selectableItems.length, }); return; } const overSelectableItem = selectableItems.find((i) => i.node.item.id === over.id) ?? null; if (overSelectableItem == null) { return; } const draggingItems = jotaiStore.get(draggingIdsFamily(treeId)); for (const id of draggingItems) { const item = selectableItems.find((i) => i.node.item.id === id)?.node ?? null; if (item == null) { return; } const isSameParent = item.parent?.item.id === overSelectableItem.node.parent?.item.id; if (item.localDrag && !isSameParent) { return; } } const node = overSelectableItem.node; const side = computeSideForDragMove(node.item.id, e); const item = node.item; let hoveredParent = node.parent; const dragIndex = selectableItems.findIndex((n) => n.node.item.id === item.id) ?? -1; const hovered = selectableItems[dragIndex]?.node ?? null; const hoveredIndex = dragIndex + (side === 'before' ? 0 : 1); let hoveredChildIndex = overSelectableItem.index + (side === 'before' ? 0 : 1); // Move into the folder if it's open and we're moving after it if (hovered?.children != null && side === 'after') { hoveredParent = hovered; hoveredChildIndex = 0; } const parentId = hoveredParent?.item.id ?? null; const parentDepth = hoveredParent?.depth ?? null; const index = hoveredIndex; const childIndex = hoveredChildIndex; const existing = jotaiStore.get(hoveredParentFamily(treeId)); if ( !( parentId === existing.parentId && parentDepth === existing.parentDepth && index === existing.index && childIndex === existing.childIndex ) ) { jotaiStore.set(hoveredParentFamily(treeId), { parentId, parentDepth, index, childIndex, }); } }, [root.depth, root.item.id, selectableItems, treeId], ); const handleDragStart = useCallback( function handleDragStart(e: DragStartEvent) { const selectedItems = getSelectedItems(treeId, selectableItems); const isDraggingSelectedItem = selectedItems.find((i) => i.id === e.active.id); // If we started dragging an already-selected item, we'll use that if (isDraggingSelectedItem) { jotaiStore.set( draggingIdsFamily(treeId), selectedItems.map((i) => i.id), ); } else { // If we started dragging a non-selected item, only drag that item const activeItem = selectableItems.find((i) => i.node.item.id === e.active.id)?.node.item; if (activeItem != null) { jotaiStore.set(draggingIdsFamily(treeId), [activeItem.id]); // Also update selection to just be this one handleSelect(activeItem, { shiftKey: false, metaKey: false, ctrlKey: false, }); } } }, [handleSelect, selectableItems, treeId], ); const clearDragState = useCallback(() => { jotaiStore.set(hoveredParentFamily(treeId), { parentId: null, parentDepth: null, index: null, childIndex: null, }); jotaiStore.set(draggingIdsFamily(treeId), []); }, [treeId]); const handleDragEnd = useCallback( function handleDragEnd(e: DragEndEvent) { // Get this from the store so our callback doesn't change all the time const { index: hoveredIndex, parentId: hoveredParentId, childIndex: hoveredChildIndex, } = jotaiStore.get(hoveredParentFamily(treeId)); const draggingItems = jotaiStore.get(draggingIdsFamily(treeId)); clearDragState(); // Dropped outside the tree? if (e.over == null) { return; } const hoveredParentS = hoveredParentId === root.item.id ? { node: root, depth: 0, index: 0 } : (selectableItems.find((i) => i.node.item.id === hoveredParentId) ?? null); const hoveredParent = hoveredParentS?.node ?? null; if (hoveredParent == null || hoveredIndex == null || !draggingItems?.length) { return; } // Resolve the actual tree nodes for each dragged item (keeps order of draggingItems) const draggedNodes: TreeNode[] = draggingItems .map((id) => { return selectableItems.find((i) => i.node.item.id === id)?.node ?? null; }) .filter((n) => n != null) // Filter out invalid drags (dragging into descendant) .filter( (n) => hoveredParent.item.id !== n.item.id && !hasAncestor(hoveredParent, n.item.id), ); // Work on a local copy of target children const nextChildren = [...(hoveredParent.children ?? [])]; // Remove any of the dragged nodes already in the target, adjusting hoveredIndex let insertAt = hoveredChildIndex ?? 0; for (const node of draggedNodes) { const i = nextChildren.findIndex((n) => n.item.id === node.item.id); if (i !== -1) { nextChildren.splice(i, 1); if (i < insertAt) insertAt -= 1; // account for removed-before } } // Batch callback onDragEnd?.({ items: draggedNodes.map((n) => n.item), parent: hoveredParent.item, children: nextChildren.map((c) => c.item), insertAt, }); }, [treeId, clearDragState, selectableItems, root, onDragEnd], ); const treeItemListProps: Omit< TreeItemListProps, 'nodes' | 'treeId' | 'activeIdAtom' | 'hoveredParent' | 'hoveredIndex' > = { getItemKey, getContextMenu: handleGetContextMenu, onClick: handleClick, getEditOptions, ItemInner, ItemLeftSlotInner, ItemRightSlot, }; const handleContextMenu = useCallback( async (e: MouseEvent) => { if (getContextMenu == null) return; e.preventDefault(); e.stopPropagation(); const items = await getContextMenu([]); setShowContextMenu({ items, x: e.clientX, y: e.clientY }); }, [getContextMenu], ); const sensors = useSensors(useSensor(PointerSensor, { activationConstraint: { distance: 6 } })); return ( <> {showContextMenu && ( )}
{/* Assign root ID so we can reuse our same move/end logic */}
); } // 1) Preserve generics through forwardRef: const Tree_ = forwardRef(TreeInner) as ( props: TreeProps & RefAttributes, ) => ReactElement | null; export const Tree = memo( Tree_, ({ root: prevNode, ...prevProps }, { root: nextNode, ...nextProps }) => { for (const key of Object.keys(prevProps)) { if (prevProps[key as keyof typeof prevProps] !== nextProps[key as keyof typeof nextProps]) { return false; } } return equalSubtree(prevNode, nextNode, nextProps.getItemKey); }, ) as typeof Tree_; function DropRegionAfterList({ id, onContextMenu, }: { id: string; onContextMenu?: (e: MouseEvent) => void; }) { const { setNodeRef } = useDroppable({ id }); // biome-ignore lint/a11y/noStaticElementInteractions: Meh return
; } interface TreeHotKeyProps { action: HotkeyAction; selectableItems: SelectableTreeNode[]; treeId: string; onDone: (items: T[]) => void; priority?: number; enable?: boolean | (() => boolean); } function TreeHotKey({ treeId, action, onDone, selectableItems, enable, ...options }: TreeHotKeyProps) { useHotKey( action, () => { onDone(getSelectedItems(treeId, selectableItems)); }, { ...options, enable: () => { if (enable == null) return true; if (typeof enable === 'function') return enable(); return enable; }, }, ); return null; } function TreeHotKeys({ treeId, hotkeys, selectableItems, }: { treeId: string; hotkeys: TreeProps['hotkeys']; selectableItems: SelectableTreeNode[]; }) { if (hotkeys == null) return null; return ( <> {Object.entries(hotkeys.actions).map(([hotkey, { cb, ...options }]) => ( ))} ); } function getValidSelectableItems( treeId: string, selectableItems: SelectableTreeNode[], ) { const collapsed = jotaiStore.get(collapsedFamily(treeId)); return selectableItems.filter((i) => { if (i.node.hidden) return false; let p = i.node.parent; while (p) { if (collapsed[p.item.id]) return false; p = p.parent; } return true; }); }