1use std::{iter::Peekable, ops::Range, slice::Iter};
2
3use ratatui::widgets::ListState;
4
5use crate::{
6 config::Symbols,
7 note_editor::ast::{HeadingLevel, Node},
8};
9
10use super::item::{FindItem, Flatten, Item};
11
12#[derive(Debug, Default, Clone, PartialEq)]
13pub struct OutlineState {
14 pub(crate) selected_item_index: Option<usize>,
15 pub(crate) items: Vec<Item>,
16 pub(crate) open: bool,
17 pub(crate) list_state: ListState,
18 pub(crate) active: bool,
19 pub(crate) symbols: Symbols,
20}
21
22#[derive(Debug, Clone, PartialEq)]
23struct Heading {
24 index: usize,
25 level: HeadingLevel,
26 content: String,
27}
28
29#[derive(Debug, Clone, PartialEq)]
30struct HeadingEntry {
31 range: Range<usize>,
32 level: HeadingLevel,
33 content: String,
34 children: Vec<HeadingEntry>,
35}
36
37impl From<HeadingEntry> for Item {
38 fn from(value: HeadingEntry) -> Self {
39 if value.children.is_empty() {
40 Item::Heading {
41 range: value.range,
42 content: value.content,
43 }
44 } else {
45 Item::HeadingEntry {
46 range: value.range,
47 content: value.content,
48 children: value.children.into_iter().map(Item::from).collect(),
49 expanded: false,
50 }
51 }
52 }
53}
54
55fn build_outline_tree(headings: &[Heading], max_end: usize) -> Vec<HeadingEntry> {
56 fn build_outline_tree_rec(
57 headings: &mut Peekable<Iter<Heading>>,
58 parent_level: Option<HeadingLevel>,
59 max_end: usize,
60 ) -> Vec<HeadingEntry> {
61 let mut result: Vec<HeadingEntry> = vec![];
62
63 while let Some(next_heading) = headings.peek() {
64 if parent_level.is_some_and(|parent_level| next_heading.level <= parent_level) {
65 break;
66 }
67
68 if let Some(heading) = headings.next() {
69 let next_heading = headings.peek();
70 let range_start = heading.index;
71 let range_end = next_heading
72 .map(|next_heading| next_heading.index)
73 .unwrap_or(max_end);
74
75 let children = match next_heading {
76 Some(next_heading) if next_heading.level > heading.level => {
77 build_outline_tree_rec(headings, Some(heading.level), max_end)
78 }
79 _ => vec![],
80 };
81
82 result.push(HeadingEntry {
83 range: range_start..range_end,
84 level: heading.level,
85 content: heading.content.clone(),
86 children,
87 });
88 }
89 }
90
91 result
92 }
93
94 build_outline_tree_rec(&mut headings.iter().peekable(), None, max_end)
95}
96
97trait NodesAsHeadings {
98 fn to_headings(&self) -> Vec<Heading>;
99}
100
101impl NodesAsHeadings for &[Node] {
102 fn to_headings(&self) -> Vec<Heading> {
103 self.iter()
104 .enumerate()
105 .filter_map(|(index, node)| {
106 if let Node::Heading { level, text, .. } = &node {
107 Some(Heading {
108 index,
109 level: *level,
110 content: text.to_string(),
111 })
112 } else {
113 None
114 }
115 })
116 .collect()
117 }
118}
119
120trait HeadingsAsItems {
121 fn to_items(&self, max_end: usize) -> Vec<Item>;
122}
123
124impl HeadingsAsItems for Vec<Heading> {
125 fn to_items(&self, max_end: usize) -> Vec<Item> {
126 build_outline_tree(self, max_end)
127 .into_iter()
128 .map(Item::from)
129 .collect()
130 }
131}
132
133impl OutlineState {
134 pub fn new(nodes: &[Node], index: usize, open: bool, symbols: &Symbols) -> Self {
135 let headings = nodes.to_headings();
136
137 let mut state = OutlineState {
138 open,
139 selected_item_index: None,
140 items: headings.to_items(nodes.len()),
141 list_state: ListState::default(),
142 symbols: symbols.clone(),
143 ..Default::default()
144 };
145 state.select_at(index);
146 state.expand_all();
147 state
148 }
149
150 pub fn set_nodes(&mut self, nodes: &[Node]) {
151 let headings = nodes.to_headings();
152 self.items = headings.to_items(nodes.len());
153 self.expand_all();
154 }
155
156 pub fn selected(&self) -> Option<Item> {
157 if let Some(selected) = self.list_state.selected() {
158 self.items.flatten().get(selected).cloned()
159 } else {
160 None
161 }
162 }
163
164 pub fn set_active(&mut self, active: bool) {
165 self.active = active;
166 }
167
168 pub fn toggle(&mut self) {
169 self.open = !self.open;
170 }
171
172 pub fn open(&mut self) {
173 self.open = true;
174 }
175
176 pub fn close(&mut self) {
177 self.open = false;
178 }
179
180 fn toggle_item_in_tree(item: &Item, target_range: &Range<usize>, should_toggle: bool) -> Item {
181 let item = item.clone();
182
183 match item {
184 Item::HeadingEntry {
185 range: heading_range,
186 content,
187 expanded,
188 children,
189 } => {
190 let expanded = if heading_range == *target_range && should_toggle {
191 !expanded
192 } else {
193 expanded
194 };
195
196 Item::HeadingEntry {
197 range: heading_range.clone(),
198 content,
199 expanded,
200 children: children
201 .iter()
202 .map(|item| Self::toggle_item_in_tree(item, target_range, should_toggle))
203 .collect(),
204 }
205 }
206 _ => item,
207 }
208 }
209
210 pub fn toggle_item(&mut self) {
211 let index = self.list_state.selected().unwrap_or_default();
212
213 let items = self.items.flatten();
214 let selected_item = items.get(index);
215
216 if let Some(Item::HeadingEntry { range, .. }) = selected_item {
217 let target_range = range.clone();
218
219 self.items = self
220 .items
221 .iter()
222 .map(|item| Self::toggle_item_in_tree(item, &target_range, true))
223 .collect();
224 };
225 }
226
227 pub fn select_at(&mut self, index: usize) {
228 let (selected_item_index, _) = self.items.find_item(index).unzip();
229 self.selected_item_index = selected_item_index;
230 self.list_state.select(selected_item_index);
231 }
232
233 fn expanded_to_all_items(items: &[Item], expanded: bool) -> Vec<Item> {
234 items
235 .iter()
236 .map(|item| match item {
237 Item::HeadingEntry {
238 range,
239 content,
240 children,
241 ..
242 } => Item::HeadingEntry {
243 range: range.clone(),
244 content: content.clone(),
245 children: Self::expanded_to_all_items(children, expanded),
246 expanded,
247 },
248 heading => heading.clone(),
249 })
250 .collect()
251 }
252
253 fn get_visible_item_count(&self) -> usize {
254 fn item_count(items: &[Item]) -> usize {
255 items.iter().fold(0, |acc, item| {
256 let next = acc + 1;
257 match item {
258 Item::HeadingEntry {
259 expanded: true,
260 children,
261 ..
262 } => next + item_count(children),
263 _ => next,
264 }
265 })
266 }
267
268 item_count(&self.items)
269 }
270
271 pub fn expand_all(&mut self) {
272 self.items = Self::expanded_to_all_items(self.items.as_slice(), true);
273 }
274
275 pub fn collapse_all(&mut self) {
276 self.items = Self::expanded_to_all_items(self.items.as_slice(), false);
277 }
278
279 pub fn is_open(&self) -> bool {
280 self.open
281 }
282
283 pub fn next(&mut self, amount: usize) {
284 let index = self
285 .list_state
286 .selected()
287 .map(|i| (i + amount).min(self.get_visible_item_count().saturating_sub(1)))
288 .unwrap_or_default();
289 self.list_state.select(Some(index));
290 }
291
292 pub fn previous(&mut self, amount: usize) {
293 let index = self.list_state.selected().map(|i| i.saturating_sub(amount));
294 self.list_state.select(index);
295 }
296}