1  import type { SessionEntry } from "@mariozechner/pi-coding-agent";
  2  
  3  export const ANYCOPY_FOLD_STATE_CUSTOM_TYPE = "anycopy-fold-state";
  4  const FOLD_STATE_SCHEMA_VERSION = 1;
  5  
  6  type TreeSelectorLike = {
  7  	getTreeList(): unknown;
  8  };
  9  
 10  type TreeListInternals = {
 11  	foldedNodes?: Set<string>;
 12  	flatNodes?: Array<{ node?: { entry?: { id?: string } } }>;
 13  	applyFilter?: () => void;
 14  	isFoldable?: (entryId: string) => boolean;
 15  };
 16  
 17  export type AnycopyFoldStateEntryData = {
 18  	v: 1;
 19  	foldedNodeIds: string[];
 20  };
 21  
 22  type MergeExplicitFoldMutationArgs = {
 23  	durableFoldedNodeIds: Iterable<string>;
 24  	beforeTransientFoldedNodeIds: Iterable<string>;
 25  	afterTransientFoldedNodeIds: Iterable<string>;
 26  	validNodeIds?: ReadonlySet<string>;
 27  };
 28  
 29  export const normalizeFoldedNodeIds = (
 30  	nodeIds: Iterable<unknown>,
 31  	validNodeIds?: ReadonlySet<string>,
 32  ): string[] => {
 33  	const uniqueNodeIds = new Set<string>();
 34  	for (const nodeId of nodeIds) {
 35  		if (typeof nodeId !== "string" || nodeId.length === 0) continue;
 36  		if (validNodeIds && !validNodeIds.has(nodeId)) continue;
 37  		uniqueNodeIds.add(nodeId);
 38  	}
 39  	return [...uniqueNodeIds].sort();
 40  };
 41  
 42  export const foldStateNodeIdListsEqual = (left: readonly string[], right: readonly string[]): boolean =>
 43  	left.length === right.length && left.every((value, index) => value === right[index]);
 44  
 45  export const createFoldStateEntryData = (
 46  	foldedNodeIds: Iterable<string>,
 47  	validNodeIds?: ReadonlySet<string>,
 48  ): AnycopyFoldStateEntryData => ({
 49  	v: FOLD_STATE_SCHEMA_VERSION,
 50  	foldedNodeIds: normalizeFoldedNodeIds(foldedNodeIds, validNodeIds),
 51  });
 52  
 53  export const parseFoldStateEntryData = (
 54  	data: unknown,
 55  	validNodeIds?: ReadonlySet<string>,
 56  ): AnycopyFoldStateEntryData | null => {
 57  	if (!data || typeof data !== "object") return null;
 58  
 59  	const candidate = data as { v?: unknown; foldedNodeIds?: unknown };
 60  	if (candidate.v !== FOLD_STATE_SCHEMA_VERSION || !Array.isArray(candidate.foldedNodeIds)) {
 61  		return null;
 62  	}
 63  
 64  	return createFoldStateEntryData(candidate.foldedNodeIds, validNodeIds);
 65  };
 66  
 67  const getTreeListInternals = (selector: TreeSelectorLike): TreeListInternals | null => {
 68  	const treeList = selector.getTreeList();
 69  	if (!treeList || typeof treeList !== "object") return null;
 70  	return treeList as TreeListInternals;
 71  };
 72  
 73  export const getSelectorFoldedNodeIds = (selector: TreeSelectorLike): string[] => {
 74  	const internals = getTreeListInternals(selector);
 75  	if (!internals?.foldedNodes) return [];
 76  	return normalizeFoldedNodeIds(internals.foldedNodes);
 77  };
 78  
 79  export const setSelectorFoldedNodeIds = (selector: TreeSelectorLike, entryIds: Iterable<string>): string[] => {
 80  	const internals = getTreeListInternals(selector);
 81  	if (!internals?.foldedNodes || !internals.applyFilter || !Array.isArray(internals.flatNodes)) {
 82  		return [];
 83  	}
 84  
 85  	const validNodeIds = new Set(
 86  		internals.flatNodes
 87  			.map((flatNode) => flatNode.node?.entry?.id)
 88  			.filter((entryId): entryId is string => typeof entryId === "string"),
 89  	);
 90  
 91  	internals.foldedNodes.clear();
 92  	internals.applyFilter();
 93  
 94  	for (const entryId of normalizeFoldedNodeIds(entryIds, validNodeIds)) {
 95  		if (!internals.isFoldable || internals.isFoldable(entryId)) {
 96  			internals.foldedNodes.add(entryId);
 97  		}
 98  	}
 99  
100  	internals.applyFilter();
101  	return normalizeFoldedNodeIds(internals.foldedNodes);
102  };
103  
104  export const loadLatestFoldStateFromEntries = (
105  	entries: Iterable<SessionEntry>,
106  	validNodeIds?: ReadonlySet<string>,
107  ): AnycopyFoldStateEntryData | null => {
108  	const sessionEntries = [...entries];
109  	for (let index = sessionEntries.length - 1; index >= 0; index -= 1) {
110  		const entry = sessionEntries[index];
111  		if (entry.type !== "custom" || entry.customType !== ANYCOPY_FOLD_STATE_CUSTOM_TYPE) {
112  			continue;
113  		}
114  
115  		const parsedEntry = parseFoldStateEntryData(entry.data, validNodeIds);
116  		if (parsedEntry) {
117  			return parsedEntry;
118  		}
119  	}
120  
121  	return null;
122  };
123  
124  export const mergeExplicitFoldMutation = ({
125  	durableFoldedNodeIds,
126  	beforeTransientFoldedNodeIds,
127  	afterTransientFoldedNodeIds,
128  	validNodeIds,
129  }: MergeExplicitFoldMutationArgs): string[] => {
130  	const normalizedDurableFoldedNodeIds = normalizeFoldedNodeIds(durableFoldedNodeIds, validNodeIds);
131  	const normalizedBeforeTransientFoldedNodeIds = normalizeFoldedNodeIds(beforeTransientFoldedNodeIds, validNodeIds);
132  	const normalizedAfterTransientFoldedNodeIds = normalizeFoldedNodeIds(afterTransientFoldedNodeIds, validNodeIds);
133  
134  	const removedNodeIds = new Set(
135  		normalizedBeforeTransientFoldedNodeIds.filter((nodeId) => !normalizedAfterTransientFoldedNodeIds.includes(nodeId)),
136  	);
137  	const addedNodeIds = new Set(
138  		normalizedAfterTransientFoldedNodeIds.filter((nodeId) => !normalizedBeforeTransientFoldedNodeIds.includes(nodeId)),
139  	);
140  	const nextDurableFoldedNodeIds = normalizedDurableFoldedNodeIds.filter((nodeId) => !removedNodeIds.has(nodeId));
141  	for (const nodeId of addedNodeIds) {
142  		nextDurableFoldedNodeIds.push(nodeId);
143  	}
144  
145  	return normalizeFoldedNodeIds(nextDurableFoldedNodeIds, validNodeIds);
146  };