1use itertools::Itertools;
6use rustc_hash::FxHashMap;
7
8use crate::{SyntaxElement, SyntaxNode};
9
10#[derive(Debug, Default)]
11pub struct SyntaxMapping {
12 entry_parents: Vec<SyntaxNode>,
19 node_mappings: FxHashMap<SyntaxNode, MappingEntry>,
20}
21
22impl SyntaxMapping {
23 pub fn upmap_child_element(
25 &self,
26 child: &SyntaxElement,
27 input_ancestor: &SyntaxNode,
28 output_ancestor: &SyntaxNode,
29 ) -> Result<SyntaxElement, MissingMapping> {
30 match child {
31 SyntaxElement::Node(node) => {
32 self.upmap_child(node, input_ancestor, output_ancestor).map(SyntaxElement::Node)
33 }
34 SyntaxElement::Token(token) => {
35 let upmap_parent =
36 self.upmap_child(&token.parent().unwrap(), input_ancestor, output_ancestor)?;
37
38 let element = upmap_parent.children_with_tokens().nth(token.index()).unwrap();
39 debug_assert!(
40 element.as_token().is_some_and(|it| it.kind() == token.kind()),
41 "token upmapping mapped to the wrong node ({token:?} -> {element:?})"
42 );
43
44 Ok(element)
45 }
46 }
47 }
48
49 pub fn upmap_child(
52 &self,
53 child: &SyntaxNode,
54 input_ancestor: &SyntaxNode,
55 output_ancestor: &SyntaxNode,
56 ) -> Result<SyntaxNode, MissingMapping> {
57 debug_assert!(
58 child == input_ancestor
59 || child.ancestors().any(|ancestor| &ancestor == input_ancestor)
60 );
61
62 let to_first_upmap = if child != input_ancestor {
64 std::iter::successors(Some((child.index(), child.clone())), |(_, current)| {
65 let parent = current.parent().unwrap();
66
67 if &parent == input_ancestor {
68 return None;
69 }
70
71 Some((parent.index(), parent))
72 })
73 .map(|(i, _)| i)
74 .collect::<Vec<_>>()
75 } else {
76 vec![]
77 };
78
79 let to_output_ancestor = if input_ancestor != output_ancestor {
81 self.upmap_to_ancestor(input_ancestor, output_ancestor)?
82 } else {
83 vec![]
84 };
85
86 let to_map_down =
87 to_output_ancestor.into_iter().rev().chain(to_first_upmap.into_iter().rev());
88
89 let mut target = output_ancestor.clone();
90
91 for index in to_map_down {
92 target = target
93 .children_with_tokens()
94 .nth(index)
95 .and_then(|it| it.into_node())
96 .expect("equivalent ancestor node should be present in target tree");
97 }
98
99 debug_assert_eq!(child.kind(), target.kind());
100
101 Ok(target)
102 }
103
104 fn upmap_to_ancestor(
105 &self,
106 input_ancestor: &SyntaxNode,
107 output_ancestor: &SyntaxNode,
108 ) -> Result<Vec<usize>, MissingMapping> {
109 let mut current =
110 self.upmap_node_single(input_ancestor).unwrap_or_else(|| input_ancestor.clone());
111 let mut upmap_chain = vec![current.index()];
112
113 loop {
114 let Some(parent) = current.parent() else { break };
115
116 if &parent == output_ancestor {
117 return Ok(upmap_chain);
118 }
119
120 current = match self.upmap_node_single(&parent) {
121 Some(next) => next,
122 None => parent,
123 };
124 upmap_chain.push(current.index());
125 }
126
127 Err(MissingMapping(current))
128 }
129
130 pub fn upmap_element(
131 &self,
132 input: &SyntaxElement,
133 output_root: &SyntaxNode,
134 ) -> Option<Result<SyntaxElement, MissingMapping>> {
135 match input {
136 SyntaxElement::Node(node) => {
137 Some(self.upmap_node(node, output_root)?.map(SyntaxElement::Node))
138 }
139 SyntaxElement::Token(token) => {
140 let upmap_parent = match self.upmap_node(&token.parent().unwrap(), output_root)? {
141 Ok(it) => it,
142 Err(err) => return Some(Err(err)),
143 };
144
145 let element = upmap_parent.children_with_tokens().nth(token.index()).unwrap();
146 debug_assert!(
147 element.as_token().is_some_and(|it| it.kind() == token.kind()),
148 "token upmapping mapped to the wrong node ({token:?} -> {element:?})"
149 );
150
151 Some(Ok(element))
152 }
153 }
154 }
155
156 pub fn upmap_node(
157 &self,
158 input: &SyntaxNode,
159 output_root: &SyntaxNode,
160 ) -> Option<Result<SyntaxNode, MissingMapping>> {
161 let input_mapping = self.upmap_node_single(input);
163 let input_ancestor =
164 input.ancestors().find_map(|ancestor| self.upmap_node_single(&ancestor));
165
166 match (input_mapping, input_ancestor) {
167 (Some(input_mapping), _) => {
168 Some(self.upmap_child(&input_mapping, &input_mapping, output_root))
170 }
171 (None, Some(input_ancestor)) => {
172 Some(self.upmap_child(input, &input_ancestor, output_root))
174 }
175 (None, None) => {
176 None
178 }
179 }
180 }
181
182 pub fn merge(&mut self, mut other: SyntaxMapping) {
183 let remap_base: u32 = self.entry_parents.len().try_into().unwrap();
185
186 self.entry_parents.append(&mut other.entry_parents);
187 self.node_mappings.extend(other.node_mappings.into_iter().map(|(node, entry)| {
188 (node, MappingEntry { parent: entry.parent + remap_base, ..entry })
189 }));
190 }
191
192 fn upmap_node_single(&self, input: &SyntaxNode) -> Option<SyntaxNode> {
194 let MappingEntry { parent, child_slot } = self.node_mappings.get(input)?;
195
196 let output = self.entry_parents[*parent as usize]
197 .children_with_tokens()
198 .nth(*child_slot as usize)
199 .and_then(SyntaxElement::into_node)
200 .unwrap();
201
202 debug_assert_eq!(input.kind(), output.kind());
203 Some(output)
204 }
205
206 pub fn add_mapping(&mut self, syntax_mapping: SyntaxMappingBuilder) {
207 let SyntaxMappingBuilder { parent_node, node_mappings } = syntax_mapping;
208
209 let parent_entry: u32 = self.entry_parents.len().try_into().unwrap();
210 self.entry_parents.push(parent_node);
211
212 let node_entries = node_mappings
213 .into_iter()
214 .map(|(node, slot)| (node, MappingEntry { parent: parent_entry, child_slot: slot }));
215
216 self.node_mappings.extend(node_entries);
217 }
218}
219
220#[derive(Debug)]
221pub struct SyntaxMappingBuilder {
222 parent_node: SyntaxNode,
223 node_mappings: Vec<(SyntaxNode, u32)>,
224}
225
226impl SyntaxMappingBuilder {
227 pub fn new(parent_node: SyntaxNode) -> Self {
228 Self { parent_node, node_mappings: vec![] }
229 }
230
231 pub fn map_node(&mut self, input: SyntaxNode, output: SyntaxNode) {
232 debug_assert_eq!(output.parent().as_ref(), Some(&self.parent_node));
233 self.node_mappings.push((input, output.index() as u32));
234 }
235
236 pub fn map_children(
237 &mut self,
238 input: impl IntoIterator<Item = SyntaxNode>,
239 output: impl IntoIterator<Item = SyntaxNode>,
240 ) {
241 for pairs in input.into_iter().zip_longest(output) {
242 let (input, output) = match pairs {
243 itertools::EitherOrBoth::Both(l, r) => (l, r),
244 itertools::EitherOrBoth::Left(_) => {
245 unreachable!("mapping more input nodes than there are output nodes")
246 }
247 itertools::EitherOrBoth::Right(_) => break,
248 };
249
250 self.map_node(input, output);
251 }
252 }
253
254 pub fn finish(self, mappings: &mut SyntaxMapping) {
255 mappings.add_mapping(self);
256 }
257}
258
259#[derive(Debug)]
260pub struct MissingMapping(pub SyntaxNode);
261
262#[derive(Debug, Clone, Copy)]
263struct MappingEntry {
264 parent: u32,
265 child_slot: u32,
266}