Remove AST-node dependency from FunctionType and ClassType (#14087)

Co-authored-by: Alex Waygood <[email protected]>

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -121,9 +121,11 @@ impl<'db> SemanticIndexBuilder<'db> {
     fn push_scope_with_parent(&mut self, node: NodeWithScopeRef, parent: Option<FileScopeId>) {
         let children_start = self.scopes.next_index() + 1;
 
+        #[allow(unsafe_code)]
         let scope = Scope {
             parent,
-            kind: node.scope_kind(),
+            // SAFETY: `node` is guaranteed to be a child of `self.module`
+            node: unsafe { node.to_kind(self.module.clone()) },
             descendents: children_start..children_start,
         };
         self.try_node_context_stack_manager.enter_nested_scope();
@@ -133,15 +135,7 @@ impl<'db> SemanticIndexBuilder<'db> {
         self.use_def_maps.push(UseDefMapBuilder::new());
         let ast_id_scope = self.ast_ids.push(AstIdsBuilder::new());
 
-        #[allow(unsafe_code)]
-        // SAFETY: `node` is guaranteed to be a child of `self.module`
-        let scope_id = ScopeId::new(
-            self.db,
-            self.file,
-            file_scope_id,
-            unsafe { node.to_kind(self.module.clone()) },
-            countme::Count::default(),
-        );
+        let scope_id = ScopeId::new(self.db, self.file, file_scope_id, countme::Count::default());
 
         self.scope_ids_by_scope.push(scope_id);
         self.scopes_by_node.insert(node.node_key(), file_scope_id);

crates/red_knot_python_semantic/src/semantic_index/definition.rs

@@ -9,6 +9,19 @@ use crate::semantic_index::symbol::{FileScopeId, ScopeId, ScopedSymbolId};
 use crate::unpack::Unpack;
 use crate::Db;
 
+/// A definition of a symbol.
+///
+/// ## Module-local type
+/// This type should not be used as part of any cross-module API because
+/// it holds a reference to the AST node. Range-offset changes
+/// then propagate through all usages, and deserialization requires
+/// reparsing the entire module.
+///
+/// E.g. don't use this type in:
+///
+/// * a return type of a cross-module query
+/// * a field of a type that is a return type of a cross-module query
+/// * an argument of a cross-module query
 #[salsa::tracked]
 pub struct Definition<'db> {
     /// The file in which the definition occurs.

crates/red_knot_python_semantic/src/semantic_index/expression.rs

@@ -8,6 +8,18 @@ use salsa;
 /// An independently type-inferable expression.
 ///
 /// Includes constraint expressions (e.g. if tests) and the RHS of an unpacking assignment.
+///
+/// ## Module-local type
+/// This type should not be used as part of any cross-module API because
+/// it holds a reference to the AST node. Range-offset changes
+/// then propagate through all usages, and deserialization requires
+/// reparsing the entire module.
+///
+/// E.g. don't use this type in:
+///
+/// * a return type of a cross-module query
+/// * a field of a type that is a return type of a cross-module query
+/// * an argument of a cross-module query
 #[salsa::tracked]
 pub(crate) struct Expression<'db> {
     /// The file in which the expression occurs.

crates/red_knot_python_semantic/src/semantic_index/symbol.rs

@@ -103,14 +103,10 @@ pub struct ScopedSymbolId;
 pub struct ScopeId<'db> {
     #[id]
     pub file: File,
+
     #[id]
     pub file_scope_id: FileScopeId,
 
-    /// The node that introduces this scope.
-    #[no_eq]
-    #[return_ref]
-    pub node: NodeWithScopeKind,
-
     #[no_eq]
     count: countme::Count<ScopeId<'static>>,
 }
@@ -131,6 +127,14 @@ impl<'db> ScopeId<'db> {
         )
     }
 
+    pub(crate) fn node(self, db: &dyn Db) -> &NodeWithScopeKind {
+        self.scope(db).node()
+    }
+
+    pub(crate) fn scope(self, db: &dyn Db) -> &Scope {
+        semantic_index(db, self.file(db)).scope(self.file_scope_id(db))
+    }
+
     #[cfg(test)]
     pub(crate) fn name(self, db: &'db dyn Db) -> &'db str {
         match self.node(db) {
@@ -169,10 +173,10 @@ impl FileScopeId {
     }
 }
 
-#[derive(Debug, Eq, PartialEq)]
+#[derive(Debug)]
 pub struct Scope {
     pub(super) parent: Option<FileScopeId>,
-    pub(super) kind: ScopeKind,
+    pub(super) node: NodeWithScopeKind,
     pub(super) descendents: Range<FileScopeId>,
 }
 
@@ -181,8 +185,12 @@ impl Scope {
         self.parent
     }
 
+    pub fn node(&self) -> &NodeWithScopeKind {
+        &self.node
+    }
+
     pub fn kind(&self) -> ScopeKind {
-        self.kind
+        self.node().scope_kind()
     }
 }
 
@@ -376,21 +384,6 @@ impl NodeWithScopeRef<'_> {
         }
     }
 
-    pub(super) fn scope_kind(self) -> ScopeKind {
-        match self {
-            NodeWithScopeRef::Module => ScopeKind::Module,
-            NodeWithScopeRef::Class(_) => ScopeKind::Class,
-            NodeWithScopeRef::Function(_) => ScopeKind::Function,
-            NodeWithScopeRef::Lambda(_) => ScopeKind::Function,
-            NodeWithScopeRef::FunctionTypeParameters(_)
-            | NodeWithScopeRef::ClassTypeParameters(_) => ScopeKind::Annotation,
-            NodeWithScopeRef::ListComprehension(_)
-            | NodeWithScopeRef::SetComprehension(_)
-            | NodeWithScopeRef::DictComprehension(_)
-            | NodeWithScopeRef::GeneratorExpression(_) => ScopeKind::Comprehension,
-        }
-    }
-
     pub(crate) fn node_key(self) -> NodeWithScopeKey {
         match self {
             NodeWithScopeRef::Module => NodeWithScopeKey::Module,
@@ -438,6 +431,36 @@ pub enum NodeWithScopeKind {
     GeneratorExpression(AstNodeRef<ast::ExprGenerator>),
 }
 
+impl NodeWithScopeKind {
+    pub(super) const fn scope_kind(&self) -> ScopeKind {
+        match self {
+            Self::Module => ScopeKind::Module,
+            Self::Class(_) => ScopeKind::Class,
+            Self::Function(_) => ScopeKind::Function,
+            Self::Lambda(_) => ScopeKind::Function,
+            Self::FunctionTypeParameters(_) | Self::ClassTypeParameters(_) => ScopeKind::Annotation,
+            Self::ListComprehension(_)
+            | Self::SetComprehension(_)
+            | Self::DictComprehension(_)
+            | Self::GeneratorExpression(_) => ScopeKind::Comprehension,
+        }
+    }
+
+    pub fn expect_class(&self) -> &ast::StmtClassDef {
+        match self {
+            Self::Class(class) => class.node(),
+            _ => panic!("expected class"),
+        }
+    }
+
+    pub fn expect_function(&self) -> &ast::StmtFunctionDef {
+        match self {
+            Self::Function(function) => function.node(),
+            _ => panic!("expected function"),
+        }
+    }
+}
+
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
 pub(crate) enum NodeWithScopeKey {
     Module,