1use crate::cst::{Item, Node};
2use crate::document::Document;
3
4fn item_index_of_data(items: &[Item], data_index: usize) -> Option<usize> {
5 let mut seen = 0;
6 for (i, item) in items.iter().enumerate() {
7 if matches!(item, Item::Node(_)) {
8 if seen == data_index {
9 return Some(i);
10 }
11 seen += 1;
12 }
13 }
14 None
15}
16
17fn separator_before(items: &[Item], item_idx: usize) -> String {
20 if item_idx > 0 {
21 if let Some(Item::Ws(w)) = items.get(item_idx - 1) {
22 if let Some(nl) = w.rfind('\n') {
23 return format!("\n{}", &w[nl + 1..]);
24 }
25 }
26 }
27 " ".to_string()
28}
29
30fn reflow(node: Node, indent: usize) -> Node {
33 let text = node.to_pretty(indent);
34 match Document::parse(&text) {
35 Ok(doc) => doc
36 .items
37 .into_iter()
38 .find_map(|i| match i {
39 Item::Node(n) => Some(n),
40 _ => None,
41 })
42 .unwrap_or(node),
43 Err(_) => node,
44 }
45}
46
47fn indent_of(sep: &str) -> usize {
48 sep.strip_prefix('\n').map_or(0, str::len)
49}
50
51impl Node {
52 fn items_mut(&mut self) -> Option<&mut Vec<Item>> {
53 match self {
54 Node::List { items, .. } => Some(items),
55 _ => None,
56 }
57 }
58
59 pub fn data_len(&self) -> usize {
61 self.list_nodes().count()
62 }
63
64 pub fn data_child(&self, data_index: usize) -> Option<&Node> {
65 self.list_nodes().nth(data_index)
66 }
67
68 pub fn data_child_mut(&mut self, data_index: usize) -> Option<&mut Node> {
69 let items = self.items_mut()?;
70 let idx = item_index_of_data(items, data_index)?;
71 match &mut items[idx] {
72 Item::Node(n) => Some(n),
73 _ => None,
74 }
75 }
76
77 pub fn position_of<F: Fn(&Node) -> bool>(&self, pred: F) -> Option<usize> {
79 self.list_nodes().position(pred)
80 }
81
82 pub fn insert_child(&mut self, data_index: usize, node: Node) {
85 let Some(items) = self.items_mut() else {
86 return;
87 };
88 let Some(idx) = item_index_of_data(items, data_index) else {
89 self.push_child(node);
90 return;
91 };
92 let sep = separator_before(items, idx);
93 let node = reflow(node, indent_of(&sep));
94 items.insert(idx, Item::Node(node));
95 items.insert(idx + 1, Item::Ws(sep));
96 }
97
98 pub fn push_child(&mut self, node: Node) {
100 let last = self.data_len().checked_sub(1);
101 let Some(items) = self.items_mut() else {
102 return;
103 };
104 match last.and_then(|d| item_index_of_data(items, d)) {
105 Some(idx) => {
106 let sep = separator_before(items, idx);
107 let node = reflow(node, indent_of(&sep));
108 items.insert(idx + 1, Item::Ws(sep));
109 items.insert(idx + 2, Item::Node(node));
110 }
111 None => items.push(Item::Node(reflow(node, 0))),
112 }
113 }
114
115 pub fn replace_child(&mut self, data_index: usize, node: Node) -> bool {
116 let Some(items) = self.items_mut() else {
117 return false;
118 };
119 let Some(idx) = item_index_of_data(items, data_index) else {
120 return false;
121 };
122 let sep = separator_before(items, idx);
123 items[idx] = Item::Node(reflow(node, indent_of(&sep)));
124 true
125 }
126
127 pub fn remove_child(&mut self, data_index: usize, take_leading_trivia: bool) -> Option<Node> {
131 let items = self.items_mut()?;
132 let idx = item_index_of_data(items, data_index)?;
133 let mut start = idx;
134 while start > 0 {
135 match &items[start - 1] {
136 Item::Ws(_) => start -= 1,
137 Item::LineComment(_) | Item::BlockComment(_) | Item::DatumComment(_)
138 if take_leading_trivia =>
139 {
140 start -= 1
141 }
142 _ => break,
143 }
144 }
145 let mut removed = None;
146 for item in items.drain(start..=idx) {
147 if let Item::Node(n) = item {
148 removed = Some(n);
149 }
150 }
151 removed
152 }
153}
154
155#[cfg(test)]
156mod tests {
157 use crate::{list, sym, Document};
158
159 #[test]
160 fn push_child_indents_like_siblings() {
161 let mut doc = Document::parse("(list\n (channel a)\n (channel b))").unwrap();
162 let form = doc.forms_mut().next().unwrap();
163 form.push_child(list(vec![sym("channel"), sym("c")]));
164 assert_eq!(
165 doc.to_string(),
166 "(list\n (channel a)\n (channel b)\n (channel c))"
167 );
168 }
169
170 #[test]
171 fn remove_child_takes_attached_comment() {
172 let mut doc =
173 Document::parse("(list\n (channel a)\n ;; b's comment\n (channel b))").unwrap();
174 let form = doc.forms_mut().next().unwrap();
175 let idx = form
176 .position_of(|n| n.to_source().contains("channel b"))
177 .unwrap();
178 form.remove_child(idx, true);
179 assert_eq!(doc.to_string(), "(list\n (channel a))");
180 }
181
182 #[test]
183 fn remove_child_keeps_unrelated_regions_verbatim() {
184 let src = "(list\n (channel a) ; weird spacing preserved\n (channel b))";
185 let mut doc = Document::parse(src).unwrap();
186 let form = doc.forms_mut().next().unwrap();
187 let idx = form
188 .position_of(|n| n.to_source().contains("channel b"))
189 .unwrap();
190 form.remove_child(idx, false);
191 assert!(doc
192 .to_string()
193 .contains("(channel a) ; weird spacing preserved"));
194 }
195
196 #[test]
197 fn insert_before_tail_symbol() {
198 let mut doc = Document::parse("(cons* (channel a)\n %default-channels)").unwrap();
199 let form = doc.forms_mut().next().unwrap();
200 let tail = form
201 .position_of(|n| n.as_symbol() == Some("%default-channels"))
202 .unwrap();
203 form.insert_child(tail, list(vec![sym("channel"), sym("b")]));
204 assert_eq!(
205 doc.to_string(),
206 "(cons* (channel a)\n (channel b)\n %default-channels)"
207 );
208 }
209}