propagate annotations to mapped elements

3 files changed, 168 insertions, 52 deletions

diff --git a/src/tools/rust-analyzer/crates/syntax/src/syntax_editor.rs b/src/tools/rust-analyzer/crates/syntax/src/syntax_editor.rs
index 42373eba827..75347aeeed9 100644
--- a/src/tools/rust-analyzer/crates/syntax/src/syntax_editor.rs
+++ b/src/tools/rust-analyzer/crates/syntax/src/syntax_editor.rs
@@ -75,8 +75,8 @@ impl SyntaxEdit {
         &self.root
     }
 
-    /// Which syntax elements in the modified syntax tree were modified as part
-    /// of the edit.
+    /// Which syntax elements in the modified syntax tree were inserted or
+    /// modified as part of the edit.
     ///
     /// Note that for syntax nodes, only the upper-most parent of a set of
     /// changes is included, not any child elements that may have been modified.
@@ -343,11 +343,22 @@ mod tests {
         editor.replace(to_wrap.syntax(), new_block.syntax());
         // editor.replace(to_replace.syntax(), name_ref.syntax());
 
-        // dbg!(&editor.mappings);
         let edit = editor.finish();
 
-        let expect = expect![];
+        dbg!(&edit.annotations);
+
+        let expect = expect![[r#"
+            _ => {
+                let var_name = 2 + 2;
+                (var_name, true)
+            }"#]];
         expect.assert_eq(&edit.root.to_string());
+
         assert_eq!(edit.find_annotation(placeholder_snippet).len(), 2);
+        assert!(edit
+            .annotations
+            .iter()
+            .flat_map(|(_, elements)| elements)
+            .all(|element| element.ancestors().any(|it| &it == edit.root())))
     }
 }
diff --git a/src/tools/rust-analyzer/crates/syntax/src/syntax_editor/edit_algo.rs b/src/tools/rust-analyzer/crates/syntax/src/syntax_editor/edit_algo.rs
index f8354c5eb67..734a26b6003 100644
--- a/src/tools/rust-analyzer/crates/syntax/src/syntax_editor/edit_algo.rs
+++ b/src/tools/rust-analyzer/crates/syntax/src/syntax_editor/edit_algo.rs
@@ -1,9 +1,10 @@
 use std::{collections::VecDeque, ops::RangeInclusive};
 
 use rowan::TextRange;
+use rustc_hash::{FxHashMap, FxHashSet};
 
 use crate::{
-    syntax_editor::{Change, ChangeKind},
+    syntax_editor::{mapping::MissingMapping, Change, ChangeKind},
     ted, SyntaxElement, SyntaxNode, SyntaxNodePtr,
 };
 
@@ -29,9 +30,6 @@ pub(super) fn apply_edits(editor: SyntaxEditor) -> SyntaxEdit {
 
     let SyntaxEditor { root, mut changes, mappings, annotations } = editor;
 
-    dbg!(("initial: ", &root));
-    dbg!(&changes);
-
     // Sort changes by range then change kind, so that we can:
     // - ensure that parent edits are ordered before child edits
     // - ensure that inserts will be guaranteed to be inserted at the right range
@@ -102,15 +100,18 @@ pub(super) fn apply_edits(editor: SyntaxEditor) -> SyntaxEdit {
         }
     }
 
-    dbg!(("before: ", &changes, &dependent_changes, &independent_changes));
-
     // Map change targets to the correct syntax nodes
     let tree_mutator = TreeMutator::new(&root);
+    let mut changed_elements = vec![];
 
     for index in independent_changes {
         match &mut changes[index as usize] {
-            Change::Replace(target, _) => {
+            Change::Replace(target, new_node) => {
                 *target = tree_mutator.make_element_mut(target);
+
+                if let Some(new_node) = new_node {
+                    changed_elements.push(new_node.clone());
+                }
             }
         }
     }
@@ -124,15 +125,20 @@ pub(super) fn apply_edits(editor: SyntaxEditor) -> SyntaxEdit {
             Change::Replace(_, None) => continue, // silently drop outdated change
         };
 
+        let upmap_target = |target: &SyntaxElement| {
+            match mappings.upmap_child_element(target, &input_ancestor, &output_ancestor) {
+                Ok(it) => it,
+                Err(MissingMapping(current)) => unreachable!("no mappings exist between {current:?} (ancestor of {input_ancestor:?}) and {output_ancestor:?}"),
+            }
+        };
+
         match &mut changes[child as usize] {
             Change::Replace(target, _) => {
-                *target = mappings.upmap_child_element(target, &input_ancestor, output_ancestor)
+                *target = upmap_target(&target);
             }
         }
     }
 
-    dbg!(("after: ", &changes));
-
     // Apply changes
     for change in changes {
         match change {
@@ -141,9 +147,29 @@ pub(super) fn apply_edits(editor: SyntaxEditor) -> SyntaxEdit {
         }
     }
 
-    dbg!(("modified:", tree_mutator.mutable_clone));
+    // Propagate annotations
+    let annotations = annotations.into_iter().filter_map(|(element, annotation)| {
+        match mappings.upmap_element(&element, &tree_mutator.mutable_clone) {
+            // Needed to follow the new tree to find the resulting element
+            Some(Ok(mapped)) => Some((mapped, annotation)),
+            // Element did not need to be mapped
+            None => Some((element, annotation)),
+            // Element did not make it to the final tree
+            Some(Err(_)) => None,
+        }
+    });
+
+    let mut annotation_groups = FxHashMap::default();
 
-    todo!("draw the rest of the owl")
+    for (element, annotation) in annotations {
+        annotation_groups.entry(annotation).or_insert(vec![]).push(element);
+    }
+
+    SyntaxEdit {
+        root: tree_mutator.mutable_clone,
+        changed_elements,
+        annotations: annotation_groups,
+    }
 }
 
 fn to_owning_node(element: &SyntaxElement) -> SyntaxNode {
diff --git a/src/tools/rust-analyzer/crates/syntax/src/syntax_editor/mapping.rs b/src/tools/rust-analyzer/crates/syntax/src/syntax_editor/mapping.rs
index 8a79f7e186e..f14d0b347d7 100644
--- a/src/tools/rust-analyzer/crates/syntax/src/syntax_editor/mapping.rs
+++ b/src/tools/rust-analyzer/crates/syntax/src/syntax_editor/mapping.rs
@@ -22,19 +22,20 @@ impl SyntaxMapping {
         Self::default()
     }
 
+    /// Like [`SyntaxMapping::upmap_child`] but for syntax elements.
     pub fn upmap_child_element(
         &self,
         child: &SyntaxElement,
         input_ancestor: &SyntaxNode,
-        output_ancestor: SyntaxNode,
-    ) -> SyntaxElement {
+        output_ancestor: &SyntaxNode,
+    ) -> Result<SyntaxElement, MissingMapping> {
         match child {
             SyntaxElement::Node(node) => {
-                SyntaxElement::Node(self.upmap_child(node, input_ancestor, output_ancestor))
+                self.upmap_child(node, input_ancestor, output_ancestor).map(SyntaxElement::Node)
             }
             SyntaxElement::Token(token) => {
                 let upmap_parent =
-                    self.upmap_child(&token.parent().unwrap(), input_ancestor, output_ancestor);
+                    self.upmap_child(&token.parent().unwrap(), input_ancestor, output_ancestor)?;
 
                 let element = upmap_parent.children_with_tokens().nth(token.index()).unwrap();
                 debug_assert!(
@@ -42,21 +43,26 @@ impl SyntaxMapping {
                     "token upmapping mapped to the wrong node ({token:?} -> {element:?})"
                 );
 
-                element
+                Ok(element)
             }
         }
     }
 
+    /// Maps a child node of the input ancestor to the corresponding node in
+    /// the output ancestor.
     pub fn upmap_child(
         &self,
         child: &SyntaxNode,
         input_ancestor: &SyntaxNode,
-        output_ancestor: SyntaxNode,
-    ) -> SyntaxNode {
-        debug_assert!(child.ancestors().any(|ancestor| &ancestor == input_ancestor));
+        output_ancestor: &SyntaxNode,
+    ) -> Result<SyntaxNode, MissingMapping> {
+        debug_assert!(
+            child == input_ancestor
+                || child.ancestors().any(|ancestor| &ancestor == input_ancestor)
+        );
 
         // Build a list mapping up to the first mappable ancestor
-        let to_first_upmap =
+        let to_first_upmap = if child != input_ancestor {
             std::iter::successors(Some((child.index(), child.clone())), |(_, current)| {
                 let parent = current.parent().unwrap();
 
@@ -67,24 +73,14 @@ impl SyntaxMapping {
                 Some((parent.index(), parent))
             })
             .map(|(i, _)| i)
-            .collect::<Vec<_>>();
+            .collect::<Vec<_>>()
+        } else {
+            vec![]
+        };
 
         // Progressively up-map the input ancestor until we get to the output ancestor
-        let to_output_ancestor = if input_ancestor != &output_ancestor {
-            std::iter::successors(Some((input_ancestor.index(), self.upmap_node(input_ancestor).unwrap_or_else(|| input_ancestor.clone()))), |(_,  current)| {
-                let Some(parent) = current.parent() else {
-                    unreachable!("no mappings exist between {current:?} (ancestor of {input_ancestor:?}) and {output_ancestor:?}")
-                };
-
-                if &parent == &output_ancestor {
-                    return None;
-                }
-
-                Some((parent.index(), match self.upmap_node(&parent) {
-                    Some(next) => next,
-                    None => parent
-                }))
-            }).map(|(i, _)| i).collect::<Vec<_>>()
+        let to_output_ancestor = if input_ancestor != output_ancestor {
+            self.upmap_to_ancestor(input_ancestor, output_ancestor)?
         } else {
             vec![]
         };
@@ -92,7 +88,7 @@ impl SyntaxMapping {
         let to_map_down =
             to_output_ancestor.into_iter().rev().chain(to_first_upmap.into_iter().rev());
 
-        let mut target = output_ancestor;
+        let mut target = output_ancestor.clone();
 
         for index in to_map_down {
             target = target
@@ -104,20 +100,86 @@ impl SyntaxMapping {
 
         debug_assert_eq!(child.kind(), target.kind());
 
-        target
+        Ok(target)
     }
 
-    pub fn upmap_node(&self, input: &SyntaxNode) -> Option<SyntaxNode> {
-        let MappingEntry { parent, child_slot } = self.node_mappings.get(input)?;
+    fn upmap_to_ancestor(
+        &self,
+        input_ancestor: &SyntaxNode,
+        output_ancestor: &SyntaxNode,
+    ) -> Result<Vec<usize>, MissingMapping> {
+        eprintln!("mapping ancestor {input_ancestor:#?} to {output_ancestor:#?}");
+        let mut current =
+            self.upmap_node_single(input_ancestor).unwrap_or_else(|| input_ancestor.clone());
+        let mut upmap_chain = vec![current.index()];
+
+        loop {
+            let Some(parent) = current.parent() else { break };
+
+            if &parent == output_ancestor {
+                return Ok(upmap_chain);
+            }
 
-        let output = self.entry_parents[*parent as usize]
-            .children_with_tokens()
-            .nth(*child_slot as usize)
-            .and_then(SyntaxElement::into_node)
-            .unwrap();
+            current = match self.upmap_node_single(&parent) {
+                Some(next) => next,
+                None => parent,
+            };
+            upmap_chain.push(current.index());
+        }
 
-        debug_assert_eq!(input.kind(), output.kind());
-        Some(output)
+        Err(MissingMapping(current))
+    }
+
+    pub fn upmap_element(
+        &self,
+        input: &SyntaxElement,
+        output_root: &SyntaxNode,
+    ) -> Option<Result<SyntaxElement, MissingMapping>> {
+        match input {
+            SyntaxElement::Node(node) => {
+                Some(self.upmap_node(node, output_root)?.map(SyntaxElement::Node))
+            }
+            SyntaxElement::Token(token) => {
+                let upmap_parent = match self.upmap_node(&token.parent().unwrap(), output_root)? {
+                    Ok(it) => it,
+                    Err(err) => return Some(Err(err)),
+                };
+
+                let element = upmap_parent.children_with_tokens().nth(token.index()).unwrap();
+                debug_assert!(
+                    element.as_token().is_some_and(|it| it.kind() == token.kind()),
+                    "token upmapping mapped to the wrong node ({token:?} -> {element:?})"
+                );
+
+                Some(Ok(element))
+            }
+        }
+    }
+
+    pub fn upmap_node(
+        &self,
+        input: &SyntaxNode,
+        output_root: &SyntaxNode,
+    ) -> Option<Result<SyntaxNode, MissingMapping>> {
+        // Try to follow the mapping tree, if it exists
+        let input_mapping = self.upmap_node_single(input);
+        let input_ancestor =
+            input.ancestors().find_map(|ancestor| self.upmap_node_single(&ancestor));
+
+        match (input_mapping, input_ancestor) {
+            (Some(input_mapping), _) => {
+                // A mapping exists at the input, follow along the tree
+                Some(self.upmap_child(&input_mapping, &input_mapping, &output_root))
+            }
+            (None, Some(input_ancestor)) => {
+                // A mapping exists at an ancestor, follow along the tree
+                Some(self.upmap_child(input, &input_ancestor, &output_root))
+            }
+            (None, None) => {
+                // No mapping exists at all, is the same position in the final tree
+                None
+            }
+        }
     }
 
     pub fn merge(&mut self, mut other: SyntaxMapping) {
@@ -130,6 +192,20 @@ impl SyntaxMapping {
         }));
     }
 
+    /// Follows the input one step along the syntax mapping tree
+    fn upmap_node_single(&self, input: &SyntaxNode) -> Option<SyntaxNode> {
+        let MappingEntry { parent, child_slot } = self.node_mappings.get(input)?;
+
+        let output = self.entry_parents[*parent as usize]
+            .children_with_tokens()
+            .nth(*child_slot as usize)
+            .and_then(SyntaxElement::into_node)
+            .unwrap();
+
+        debug_assert_eq!(input.kind(), output.kind());
+        Some(output)
+    }
+
     fn add_mapping(&mut self, syntax_mapping: SyntaxMappingBuilder) {
         let SyntaxMappingBuilder { parent_node, node_mappings } = syntax_mapping;
 
@@ -183,6 +259,9 @@ impl SyntaxMappingBuilder {
     }
 }
 
+#[derive(Debug)]
+pub struct MissingMapping(pub SyntaxNode);
+
 #[derive(Debug, Clone, Copy)]
 struct MappingEntry {
     parent: u32,