syntax/syntax_editor/
mapping.rs

1//! Maps syntax elements through disjoint syntax nodes.
2//!
3//! [`SyntaxMappingBuilder`] should be used to create mappings to add to a [`SyntaxEditor`]
4
5use itertools::Itertools;
6use rustc_hash::FxHashMap;
7
8use crate::{SyntaxElement, SyntaxNode};
9
10#[derive(Debug, Default)]
11pub struct SyntaxMapping {
12    // important information to keep track of:
13    // node -> node
14    // token -> token (implicit in mappings)
15    // input parent -> output parent (for deep lookups)
16
17    // mappings ->  parents
18    entry_parents: Vec<SyntaxNode>,
19    node_mappings: FxHashMap<SyntaxNode, MappingEntry>,
20}
21
22impl SyntaxMapping {
23    /// Like [`SyntaxMapping::upmap_child`] but for syntax elements.
24    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    /// Maps a child node of the input ancestor to the corresponding node in
50    /// the output ancestor.
51    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        // Build a list mapping up to the first mappable ancestor
63        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        // Progressively up-map the input ancestor until we get to the output ancestor
80        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        // Try to follow the mapping tree, if it exists
162        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                // A mapping exists at the input, follow along the tree
169                Some(self.upmap_child(&input_mapping, &input_mapping, output_root))
170            }
171            (None, Some(input_ancestor)) => {
172                // A mapping exists at an ancestor, follow along the tree
173                Some(self.upmap_child(input, &input_ancestor, output_root))
174            }
175            (None, None) => {
176                // No mapping exists at all, is the same position in the final tree
177                None
178            }
179        }
180    }
181
182    pub fn merge(&mut self, mut other: SyntaxMapping) {
183        // Remap other's entry parents to be after the current list of entry parents
184        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    /// Follows the input one step along the syntax mapping tree
193    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}