astral-sh · MichaReiser · Apr 2, 2025 · Mar 21, 2025 · Mar 23, 2025 · Mar 25, 2025
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/red_knot_ide/Cargo.toml b/crates/red_knot_ide/Cargo.toml
@@ -12,13 +12,19 @@ license = { workspace = true }
 
 [dependencies]
 ruff_db = { workspace = true }
+ruff_python_ast = { workspace = true }
+ruff_python_parser = { workspace = true }
+ruff_text_size = { workspace = true }
 red_knot_python_semantic = { workspace = true }
 
 salsa = { workspace = true }
+smallvec = { workspace = true }
 tracing = { workspace = true }
 
 [dev-dependencies]
 red_knot_vendored = { workspace = true }
 
+insta = { workspace = true, features = ["filters"] }
+
 [lints]
 workspace = true
diff --git a/crates/red_knot_ide/src/db.rs b/crates/red_knot_ide/src/db.rs
@@ -2,7 +2,7 @@ use red_knot_python_semantic::Db as SemanticDb;
 use ruff_db::{Db as SourceDb, Upcast};
 
 #[salsa::db]
-pub trait Db: SemanticDb + Upcast<dyn SourceDb> {}
+pub trait Db: SemanticDb + Upcast<dyn SemanticDb> + Upcast<dyn SourceDb> {}
 
 #[cfg(test)]
 pub(crate) mod tests {
@@ -94,6 +94,16 @@ pub(crate) mod tests {
         }
     }
 
+    impl Upcast<dyn SemanticDb> for TestDb {
+        fn upcast(&self) -> &(dyn SemanticDb + 'static) {
+            self
+        }
+
+        fn upcast_mut(&mut self) -> &mut dyn SemanticDb {
+            self
+        }
+    }
+
     #[salsa::db]
     impl SemanticDb for TestDb {
         fn is_file_open(&self, file: File) -> bool {

diff --git a/crates/red_knot_ide/src/find_node.rs b/crates/red_knot_ide/src/find_node.rs
@@ -0,0 +1,106 @@
+use ruff_python_ast::visitor::source_order::{SourceOrderVisitor, TraversalSignal};
+use ruff_python_ast::AnyNodeRef;
+use ruff_text_size::{Ranged, TextRange};
+use std::fmt;
+use std::fmt::Formatter;
+
+/// Returns the node with a minimal range that fully contains `range`.
+///
+/// If `range` is empty and falls within a parser *synthesized* node generated during error recovery,
+/// then the first node with the given range is returned.
+///
+/// ## Panics
+/// Panics if `range` is not contained within `root`.
+pub(crate) fn covering_node(root: AnyNodeRef, range: TextRange) -> CoveringNode {
+    struct Visitor<'a> {
+        range: TextRange,
+        found: bool,
+        ancestors: Vec<AnyNodeRef<'a>>,
+    }
+
+    impl<'a> SourceOrderVisitor<'a> for Visitor<'a> {
+        fn enter_node(&mut self, node: AnyNodeRef<'a>) -> TraversalSignal {
+            // If the node fully contains the range, than it is a possible match but traverse into its children
+            // to see if there's a node with a narrower range.
+            if !self.found && node.range().contains_range(self.range) {
+                self.ancestors.push(node);
+                TraversalSignal::Traverse
+            } else {
+                TraversalSignal::Skip
+            }
+        }
+
+        fn leave_node(&mut self, node: AnyNodeRef<'a>) {
+            if !self.found && self.ancestors.last() == Some(&node) {
+                self.found = true;
+            }
+        }
+    }
+
+    assert!(
+        root.range().contains_range(range),
+        "Range is not contained within root"
+    );
+
+    let mut visitor = Visitor {
+        range,
+        found: false,
+        ancestors: Vec::new(),
+    };
+
+    root.visit_source_order(&mut visitor);
+
+    let minimal = visitor.ancestors.pop().unwrap_or(root);
+    CoveringNode {
+        node: minimal,
+        ancestors: visitor.ancestors,
+    }
+}
+
+/// The node with a minimal range that fully contains the search range.
+pub(crate) struct CoveringNode<'a> {
+    /// The node with a minimal range that fully contains the search range.
+    node: AnyNodeRef<'a>,
+
+    /// The node's ancestor (the spine up to the root).
+    ancestors: Vec<AnyNodeRef<'a>>,
+}
+
+impl<'a> CoveringNode<'a> {
+    pub(crate) fn node(&self) -> AnyNodeRef<'a> {
+        self.node
+    }
+
+    /// Returns the node's parent.
+    pub(crate) fn parent(&self) -> Option<AnyNodeRef<'a>> {
+        self.ancestors.last().copied()
+    }
+
+    /// Finds the minimal node that fully covers the range and fulfills the given predicate.
+    pub(crate) fn find(mut self, f: impl Fn(AnyNodeRef<'a>) -> bool) -> Result<Self, Self> {
+        if f(self.node) {
+            return Ok(self);
+        }
+
+        match self.ancestors.iter().rposition(|node| f(*node)) {
+            Some(index) => {
+                let node = self.ancestors[index];
+                self.ancestors.truncate(index);
+
+                Ok(Self {
+                    node,
+                    ancestors: self.ancestors,
+                })
+            }
+            None => Err(self),
+        }
+    }
+}
+
+impl fmt::Debug for CoveringNode<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.debug_tuple("NodeWithAncestors")
+            .field(&self.node)
+            .finish()
+    }
+}