1package picker
2
3import (
4 "io/fs"
5 "path"
6 "sort"
7)
8
9// CheckState represents the tri-state checkbox value for a node.
10type CheckState int
11
12const (
13 // CheckNone means no descendants are selected.
14 CheckNone CheckState = iota
15 // CheckPartial means some but not all descendants are selected.
16 CheckPartial
17 // CheckAll means all selectable descendants are selected.
18 CheckAll
19)
20
21// nodeKind distinguishes leaf-selectable nodes from structural ones.
22type nodeKind int
23
24const (
25 kindFile nodeKind = iota // regular file — always selectable
26 kindEmptyDir // directory with no children — selectable leaf
27 kindDir // directory with children — state is derived
28)
29
30// Selection tracks multi-select state over an fs.FS tree. It maintains
31// a flat set of selected leaves (files and empty dirs) and cached
32// subtree counts for O(1) tri-state lookups.
33//
34// The tree index is built once from the FS and is immutable; only the
35// selected set and subtree counts change on toggle.
36type Selection struct {
37 selected map[string]bool // selected leaves (files + empty dirs)
38
39 // Immutable tree index, built once from the FS.
40 children map[string][]string // parent → sorted child paths
41 parent map[string]string // child → parent path
42 kind map[string]nodeKind
43 totalSelectable map[string]int // total selectable leaves in subtree
44
45 // Mutable counter, updated incrementally on toggle.
46 selectedCount map[string]int // selected leaves in subtree
47}
48
49// NewSelection builds a Selection from the given fs.FS by walking the
50// entire tree once. The FS must implement fs.ReadDirFS.
51func NewSelection(fsys fs.FS) *Selection {
52 s := &Selection{
53 selected: make(map[string]bool),
54 children: make(map[string][]string),
55 parent: make(map[string]string),
56 kind: make(map[string]nodeKind),
57 totalSelectable: make(map[string]int),
58 selectedCount: make(map[string]int),
59 }
60
61 s.buildIndex(fsys, ".")
62 return s
63}
64
65// buildIndex recursively walks the FS to populate the tree index.
66// Returns the total selectable count for the subtree rooted at dir.
67func (s *Selection) buildIndex(fsys fs.FS, dir string) int {
68 entries, err := fs.ReadDir(fsys, dir)
69 if err != nil {
70 return 0
71 }
72
73 total := 0
74 for _, e := range entries {
75 child := path.Join(dir, e.Name())
76
77 s.parent[child] = dir
78 if e.IsDir() {
79 childTotal := s.buildIndex(fsys, child)
80 if childTotal == 0 {
81 // Empty directory — treat as a selectable leaf.
82 s.kind[child] = kindEmptyDir
83 s.totalSelectable[child] = 1
84 total++
85 } else {
86 s.kind[child] = kindDir
87 s.totalSelectable[child] = childTotal
88 total += childTotal
89 }
90 } else {
91 s.kind[child] = kindFile
92 s.totalSelectable[child] = 1
93 total++
94 }
95 s.children[dir] = append(s.children[dir], child)
96 }
97
98 s.totalSelectable[dir] = total
99 s.kind[dir] = kindDir
100
101 return total
102}
103
104// State returns the tri-state checkbox value for the given path.
105func (s *Selection) State(p string) CheckState {
106 k := s.kind[p]
107 switch k {
108 case kindFile, kindEmptyDir:
109 if s.selected[p] {
110 return CheckAll
111 }
112 return CheckNone
113 default:
114 total := s.totalSelectable[p]
115 count := s.selectedCount[p]
116 switch count {
117 case 0:
118 return CheckNone
119 case total:
120 return CheckAll
121 default:
122 return CheckPartial
123 }
124 }
125}
126
127// Toggle flips the selection of the given path. For files and empty
128// dirs, this is a simple toggle. For directories, it selects all
129// descendants if the current state is none or partial, or deselects
130// all if fully selected.
131func (s *Selection) Toggle(p string) {
132 state := s.State(p)
133
134 k := s.kind[p]
135 switch k {
136 case kindFile, kindEmptyDir:
137 if s.selected[p] {
138 delete(s.selected, p)
139 s.adjustCounts(p, -1)
140 } else {
141 s.selected[p] = true
142 s.adjustCounts(p, +1)
143 }
144 default:
145 if state == CheckAll {
146 s.setSubtree(p, false)
147 } else {
148 s.setSubtree(p, true)
149 }
150 }
151}
152
153// setSubtree selects or deselects all selectable leaves under p.
154// It walks the subtree once, flipping leaves and returning the count
155// of changes, then adjusts ancestor counts in a single upward pass.
156func (s *Selection) setSubtree(p string, sel bool) {
157 changed := s.applySubtree(p, sel)
158 if changed == 0 {
159 return
160 }
161
162 // Adjust ancestors of p (not p itself — applySubtree already
163 // updated interior counts within the subtree).
164 delta := changed
165 if !sel {
166 delta = -changed
167 }
168 cur := p
169 for {
170 par, ok := s.parent[cur]
171 if !ok {
172 break
173 }
174 s.selectedCount[par] += delta
175 cur = par
176 }
177}
178
179// applySubtree recursively sets leaves under p to the desired state.
180// For interior nodes within the subtree, it updates selectedCount
181// directly. Returns the number of leaves that actually changed.
182func (s *Selection) applySubtree(p string, sel bool) int {
183 k := s.kind[p]
184 switch k {
185 case kindFile, kindEmptyDir:
186 was := s.selected[p]
187 if sel == was {
188 return 0
189 }
190 if sel {
191 s.selected[p] = true
192 } else {
193 delete(s.selected, p)
194 }
195 return 1
196 default:
197 changed := 0
198 for _, child := range s.children[p] {
199 changed += s.applySubtree(child, sel)
200 }
201 // Update this interior node's count directly.
202 if sel {
203 s.selectedCount[p] += changed
204 } else {
205 s.selectedCount[p] -= changed
206 }
207 return changed
208 }
209}
210
211// adjustCounts increments or decrements selectedCount for all ancestors
212// of p (not p itself, since leaves don't have subtree counts).
213func (s *Selection) adjustCounts(p string, delta int) {
214 cur := p
215 for {
216 par, ok := s.parent[cur]
217 if !ok {
218 break
219 }
220 s.selectedCount[par] += delta
221 cur = par
222 }
223}
224
225// AllSelected reports whether every selectable leaf in the tree is
226// selected. When true, the caller should perform a full restore
227// rather than emitting --include flags.
228func (s *Selection) AllSelected() bool {
229 return s.State(".") == CheckAll
230}
231
232// SelectedPaths returns the minimal set of FS-relative paths covering
233// the current selection. Fully-selected subtrees are compressed to
234// their root directory. Returns nil if nothing is selected or if
235// everything is selected (check AllSelected() to distinguish).
236func (s *Selection) SelectedPaths() []string {
237 if s.State(".") == CheckNone || s.State(".") == CheckAll {
238 return nil
239 }
240 var paths []string
241 s.collectPaths(".", &paths)
242 sort.Strings(paths)
243 return paths
244}
245
246// collectPaths walks the tree top-down, emitting fully-selected
247// subtrees as single paths and recursing into partial ones.
248func (s *Selection) collectPaths(p string, out *[]string) {
249 for _, child := range s.children[p] {
250 state := s.State(child)
251 switch state {
252 case CheckNone:
253 continue
254 case CheckAll:
255 // Fully selected — emit just this path, don't recurse.
256 *out = append(*out, child)
257 case CheckPartial:
258 s.collectPaths(child, out)
259 }
260 }
261}