Make Definition a salsa-ingredient (#12151)

This commit is contained in:
Micha Reiser 2024-07-04 08:46:08 +02:00 committed by GitHub
parent e6e09ea93a
commit 3ce8b9fcae
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 499 additions and 707 deletions

View file

@ -3,52 +3,64 @@ use std::sync::Arc;
use rustc_hash::FxHashMap;
use ruff_db::parsed::ParsedModule;
use ruff_db::vfs::VfsFile;
use ruff_index::IndexVec;
use ruff_python_ast as ast;
use ruff_python_ast::name::Name;
use ruff_python_ast::visitor::{walk_expr, walk_stmt, Visitor};
use crate::node_key::NodeKey;
use crate::semantic_index::ast_ids::{AstId, AstIdsBuilder, ScopedClassId, ScopedFunctionId};
use crate::semantic_index::definition::{Definition, ImportDefinition, ImportFromDefinition};
use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey;
use crate::semantic_index::ast_ids::AstIdsBuilder;
use crate::semantic_index::definition::{Definition, DefinitionNodeKey, DefinitionNodeRef};
use crate::semantic_index::symbol::{
FileScopeId, Scope, ScopeKind, ScopedSymbolId, SymbolFlags, SymbolTableBuilder,
FileScopeId, NodeWithScopeKey, NodeWithScopeRef, Scope, ScopeId, ScopedSymbolId, SymbolFlags,
SymbolTableBuilder,
};
use crate::semantic_index::{NodeWithScopeId, NodeWithScopeKey, SemanticIndex};
use crate::semantic_index::SemanticIndex;
use crate::Db;
pub(super) struct SemanticIndexBuilder<'a> {
pub(super) struct SemanticIndexBuilder<'db, 'ast> {
// Builder state
module: &'a ParsedModule,
db: &'db dyn Db,
file: VfsFile,
module: &'db ParsedModule,
scope_stack: Vec<FileScopeId>,
/// the definition whose target(s) we are currently walking
current_definition: Option<Definition>,
/// the target we're currently inferring
current_target: Option<CurrentTarget<'ast>>,
// Semantic Index fields
scopes: IndexVec<FileScopeId, Scope>,
symbol_tables: IndexVec<FileScopeId, SymbolTableBuilder>,
scope_ids_by_scope: IndexVec<FileScopeId, ScopeId<'db>>,
symbol_tables: IndexVec<FileScopeId, SymbolTableBuilder<'db>>,
ast_ids: IndexVec<FileScopeId, AstIdsBuilder>,
scopes_by_expression: FxHashMap<NodeKey, FileScopeId>,
scopes_by_definition: FxHashMap<NodeWithScopeKey, FileScopeId>,
scopes_by_node: FxHashMap<NodeWithScopeKey, FileScopeId>,
scopes_by_expression: FxHashMap<ExpressionNodeKey, FileScopeId>,
definitions_by_node: FxHashMap<DefinitionNodeKey, Definition<'db>>,
}
impl<'a> SemanticIndexBuilder<'a> {
pub(super) fn new(parsed: &'a ParsedModule) -> Self {
impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast>
where
'db: 'ast,
{
pub(super) fn new(db: &'db dyn Db, file: VfsFile, parsed: &'db ParsedModule) -> Self {
let mut builder = Self {
db,
file,
module: parsed,
scope_stack: Vec::new(),
current_definition: None,
current_target: None,
scopes: IndexVec::new(),
symbol_tables: IndexVec::new(),
ast_ids: IndexVec::new(),
scope_ids_by_scope: IndexVec::new(),
scopes_by_expression: FxHashMap::default(),
scopes_by_definition: FxHashMap::default(),
scopes_by_node: FxHashMap::default(),
definitions_by_node: FxHashMap::default(),
};
builder.push_scope_with_parent(
&NodeWithScope::new(parsed.syntax(), NodeWithScopeId::Module),
None,
);
builder.push_scope_with_parent(NodeWithScopeRef::Module, None);
builder
}
@ -60,29 +72,40 @@ impl<'a> SemanticIndexBuilder<'a> {
.expect("Always to have a root scope")
}
fn push_scope(&mut self, node: &NodeWithScope) {
fn push_scope(&mut self, node: NodeWithScopeRef<'ast>) {
let parent = self.current_scope();
self.push_scope_with_parent(node, Some(parent));
}
fn push_scope_with_parent(&mut self, node: &NodeWithScope, parent: Option<FileScopeId>) {
fn push_scope_with_parent(
&mut self,
node: NodeWithScopeRef<'ast>,
parent: Option<FileScopeId>,
) {
let children_start = self.scopes.next_index() + 1;
let scope = Scope {
node: node.id(),
parent,
kind: node.scope_kind(),
descendents: children_start..children_start,
};
let scope_id = self.scopes.push(scope);
let file_scope_id = self.scopes.push(scope);
self.symbol_tables.push(SymbolTableBuilder::new());
let ast_id_scope = self.ast_ids.push(AstIdsBuilder::new());
debug_assert_eq!(ast_id_scope, scope_id);
#[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())
});
self.scope_stack.push(scope_id);
self.scopes_by_definition.insert(node.key(), scope_id);
self.scope_ids_by_scope.push(scope_id);
self.scopes_by_node.insert(node.node_key(), file_scope_id);
debug_assert_eq!(ast_id_scope, file_scope_id);
self.scope_stack.push(file_scope_id);
}
fn pop_scope(&mut self) -> FileScopeId {
@ -93,7 +116,7 @@ impl<'a> SemanticIndexBuilder<'a> {
id
}
fn current_symbol_table(&mut self) -> &mut SymbolTableBuilder {
fn current_symbol_table(&mut self) -> &mut SymbolTableBuilder<'db> {
let scope_id = self.current_scope();
&mut self.symbol_tables[scope_id]
}
@ -105,33 +128,64 @@ impl<'a> SemanticIndexBuilder<'a> {
fn add_or_update_symbol(&mut self, name: Name, flags: SymbolFlags) -> ScopedSymbolId {
let symbol_table = self.current_symbol_table();
symbol_table.add_or_update_symbol(name, flags, None)
symbol_table.add_or_update_symbol(name, flags)
}
fn add_definition(
&mut self,
definition_node: impl Into<DefinitionNodeRef<'ast>>,
symbol_id: ScopedSymbolId,
) -> Definition<'db> {
let definition_node = definition_node.into();
let definition = Definition::new(
self.db,
self.file,
self.current_scope(),
symbol_id,
#[allow(unsafe_code)]
unsafe {
definition_node.into_owned(self.module.clone())
},
);
self.definitions_by_node
.insert(definition_node.key(), definition);
definition
}
fn add_or_update_symbol_with_definition(
&mut self,
name: Name,
definition: Definition,
) -> ScopedSymbolId {
definition: impl Into<DefinitionNodeRef<'ast>>,
) -> (ScopedSymbolId, Definition<'db>) {
let symbol_table = self.current_symbol_table();
symbol_table.add_or_update_symbol(name, SymbolFlags::IS_DEFINED, Some(definition))
let id = symbol_table.add_or_update_symbol(name, SymbolFlags::IS_DEFINED);
let definition = self.add_definition(definition, id);
self.current_symbol_table().add_definition(id, definition);
(id, definition)
}
fn with_type_params(
&mut self,
with_params: &WithTypeParams,
with_params: &WithTypeParams<'ast>,
nested: impl FnOnce(&mut Self) -> FileScopeId,
) -> FileScopeId {
let type_params = with_params.type_parameters();
if let Some(type_params) = type_params {
let type_params_id = match with_params {
WithTypeParams::ClassDef { id, .. } => NodeWithScopeId::ClassTypeParams(*id),
WithTypeParams::FunctionDef { id, .. } => NodeWithScopeId::FunctionTypeParams(*id),
let with_scope = match with_params {
WithTypeParams::ClassDef { node, .. } => {
NodeWithScopeRef::ClassTypeParameters(node)
}
WithTypeParams::FunctionDef { node, .. } => {
NodeWithScopeRef::FunctionTypeParameters(node)
}
};
self.push_scope(&NodeWithScope::new(type_params, type_params_id));
self.push_scope(with_scope);
for type_param in &type_params.type_params {
let name = match type_param {
ast::TypeParam::TypeVar(ast::TypeParamTypeVar { name, .. }) => name,
@ -151,7 +205,7 @@ impl<'a> SemanticIndexBuilder<'a> {
nested_scope
}
pub(super) fn build(mut self) -> SemanticIndex {
pub(super) fn build(mut self) -> SemanticIndex<'db> {
let module = self.module;
self.visit_body(module.suite());
@ -159,7 +213,7 @@ impl<'a> SemanticIndexBuilder<'a> {
self.pop_scope();
assert!(self.scope_stack.is_empty());
assert!(self.current_definition.is_none());
assert!(self.current_target.is_none());
let mut symbol_tables: IndexVec<_, _> = self
.symbol_tables
@ -177,53 +231,48 @@ impl<'a> SemanticIndexBuilder<'a> {
ast_ids.shrink_to_fit();
symbol_tables.shrink_to_fit();
self.scopes_by_expression.shrink_to_fit();
self.definitions_by_node.shrink_to_fit();
self.scope_ids_by_scope.shrink_to_fit();
self.scopes_by_node.shrink_to_fit();
SemanticIndex {
symbol_tables,
scopes: self.scopes,
scopes_by_definition: self.scopes_by_definition,
definitions_by_node: self.definitions_by_node,
scope_ids_by_scope: self.scope_ids_by_scope,
ast_ids,
scopes_by_expression: self.scopes_by_expression,
scopes_by_node: self.scopes_by_node,
}
}
}
impl Visitor<'_> for SemanticIndexBuilder<'_> {
fn visit_stmt(&mut self, stmt: &ast::Stmt) {
let module = self.module;
#[allow(unsafe_code)]
let statement_id = unsafe {
// SAFETY: The builder only visits nodes that are part of `module`. This guarantees that
// the current statement must be a child of `module`.
self.current_ast_ids().record_statement(stmt, module)
};
impl<'db, 'ast> Visitor<'ast> for SemanticIndexBuilder<'db, 'ast>
where
'db: 'ast,
{
fn visit_stmt(&mut self, stmt: &'ast ast::Stmt) {
match stmt {
ast::Stmt::FunctionDef(function_def) => {
for decorator in &function_def.decorator_list {
self.visit_decorator(decorator);
}
let name = &function_def.name.id;
let function_id = ScopedFunctionId(statement_id);
let definition = Definition::FunctionDef(function_id);
let scope = self.current_scope();
self.add_or_update_symbol_with_definition(name.clone(), definition);
self.add_or_update_symbol_with_definition(
function_def.name.id.clone(),
function_def,
);
self.with_type_params(
&WithTypeParams::FunctionDef {
node: function_def,
id: AstId::new(scope, function_id),
},
&WithTypeParams::FunctionDef { node: function_def },
|builder| {
builder.visit_parameters(&function_def.parameters);
for expr in &function_def.returns {
builder.visit_annotation(expr);
}
builder.push_scope(&NodeWithScope::new(
function_def,
NodeWithScopeId::Function(AstId::new(scope, function_id)),
));
builder.push_scope(NodeWithScopeRef::Function(function_def));
builder.visit_body(&function_def.body);
builder.pop_scope()
},
@ -234,46 +283,28 @@ impl Visitor<'_> for SemanticIndexBuilder<'_> {
self.visit_decorator(decorator);
}
let name = &class.name.id;
let class_id = ScopedClassId(statement_id);
let definition = Definition::ClassDef(class_id);
let scope = self.current_scope();
self.add_or_update_symbol_with_definition(class.name.id.clone(), class);
self.add_or_update_symbol_with_definition(name.clone(), definition);
self.with_type_params(&WithTypeParams::ClassDef { node: class }, |builder| {
if let Some(arguments) = &class.arguments {
builder.visit_arguments(arguments);
}
self.with_type_params(
&WithTypeParams::ClassDef {
node: class,
id: AstId::new(scope, class_id),
},
|builder| {
if let Some(arguments) = &class.arguments {
builder.visit_arguments(arguments);
}
builder.push_scope(NodeWithScopeRef::Class(class));
builder.visit_body(&class.body);
builder.push_scope(&NodeWithScope::new(
class,
NodeWithScopeId::Class(AstId::new(scope, class_id)),
));
builder.visit_body(&class.body);
builder.pop_scope()
},
);
builder.pop_scope()
});
}
ast::Stmt::Import(ast::StmtImport { names, .. }) => {
for (i, alias) in names.iter().enumerate() {
for alias in names {
let symbol_name = if let Some(asname) = &alias.asname {
asname.id.clone()
} else {
Name::new(alias.name.id.split('.').next().unwrap())
};
let def = Definition::Import(ImportDefinition {
import_id: statement_id,
alias: u32::try_from(i).unwrap(),
});
self.add_or_update_symbol_with_definition(symbol_name, def);
self.add_or_update_symbol_with_definition(symbol_name, alias);
}
}
ast::Stmt::ImportFrom(ast::StmtImportFrom {
@ -282,27 +313,24 @@ impl Visitor<'_> for SemanticIndexBuilder<'_> {
level: _,
..
}) => {
for (i, alias) in names.iter().enumerate() {
for alias in names {
let symbol_name = if let Some(asname) = &alias.asname {
&asname.id
} else {
&alias.name.id
};
let def = Definition::ImportFrom(ImportFromDefinition {
import_id: statement_id,
name: u32::try_from(i).unwrap(),
});
self.add_or_update_symbol_with_definition(symbol_name.clone(), def);
self.add_or_update_symbol_with_definition(symbol_name.clone(), alias);
}
}
ast::Stmt::Assign(node) => {
debug_assert!(self.current_definition.is_none());
debug_assert!(self.current_target.is_none());
self.visit_expr(&node.value);
self.current_definition = Some(Definition::Assignment(statement_id));
for target in &node.targets {
self.current_target = Some(CurrentTarget::Expr(target));
self.visit_expr(target);
}
self.current_definition = None;
self.current_target = None;
}
_ => {
walk_stmt(self, stmt);
@ -310,17 +338,10 @@ impl Visitor<'_> for SemanticIndexBuilder<'_> {
}
}
fn visit_expr(&mut self, expr: &'_ ast::Expr) {
let module = self.module;
#[allow(unsafe_code)]
let expression_id = unsafe {
// SAFETY: The builder only visits nodes that are part of `module`. This guarantees that
// the current expression must be a child of `module`.
self.current_ast_ids().record_expression(expr, module)
};
fn visit_expr(&mut self, expr: &'ast ast::Expr) {
self.scopes_by_expression
.insert(NodeKey::from_node(expr), self.current_scope());
.insert(expr.into(), self.current_scope());
self.current_ast_ids().record_expression(expr);
match expr {
ast::Expr::Name(ast::ExprName { id, ctx, .. }) => {
@ -330,9 +351,9 @@ impl Visitor<'_> for SemanticIndexBuilder<'_> {
ast::ExprContext::Del => SymbolFlags::IS_DEFINED,
ast::ExprContext::Invalid => SymbolFlags::empty(),
};
match self.current_definition {
Some(definition) if flags.contains(SymbolFlags::IS_DEFINED) => {
self.add_or_update_symbol_with_definition(id.clone(), definition);
match self.current_target {
Some(target) if flags.contains(SymbolFlags::IS_DEFINED) => {
self.add_or_update_symbol_with_definition(id.clone(), target);
}
_ => {
self.add_or_update_symbol(id.clone(), flags);
@ -342,11 +363,11 @@ impl Visitor<'_> for SemanticIndexBuilder<'_> {
walk_expr(self, expr);
}
ast::Expr::Named(node) => {
debug_assert!(self.current_definition.is_none());
self.current_definition = Some(Definition::NamedExpr(expression_id));
debug_assert!(self.current_target.is_none());
self.current_target = Some(CurrentTarget::ExprNamed(node));
// TODO walrus in comprehensions is implicitly nonlocal
self.visit_expr(&node.target);
self.current_definition = None;
self.current_target = None;
self.visit_expr(&node.value);
}
ast::Expr::If(ast::ExprIf {
@ -382,19 +403,13 @@ impl Visitor<'_> for SemanticIndexBuilder<'_> {
}
}
enum WithTypeParams<'a> {
ClassDef {
node: &'a ast::StmtClassDef,
id: AstId<ScopedClassId>,
},
FunctionDef {
node: &'a ast::StmtFunctionDef,
id: AstId<ScopedFunctionId>,
},
enum WithTypeParams<'node> {
ClassDef { node: &'node ast::StmtClassDef },
FunctionDef { node: &'node ast::StmtFunctionDef },
}
impl<'a> WithTypeParams<'a> {
fn type_parameters(&self) -> Option<&'a ast::TypeParams> {
impl<'node> WithTypeParams<'node> {
fn type_parameters(&self) -> Option<&'node ast::TypeParams> {
match self {
WithTypeParams::ClassDef { node, .. } => node.type_params.as_deref(),
WithTypeParams::FunctionDef { node, .. } => node.type_params.as_deref(),
@ -402,35 +417,17 @@ impl<'a> WithTypeParams<'a> {
}
}
struct NodeWithScope {
id: NodeWithScopeId,
key: NodeWithScopeKey,
#[derive(Copy, Clone, Debug)]
enum CurrentTarget<'a> {
Expr(&'a ast::Expr),
ExprNamed(&'a ast::ExprNamed),
}
impl NodeWithScope {
fn new(node: impl Into<NodeWithScopeKey>, id: NodeWithScopeId) -> Self {
Self {
id,
key: node.into(),
}
}
fn id(&self) -> NodeWithScopeId {
self.id
}
fn key(&self) -> NodeWithScopeKey {
self.key
}
fn scope_kind(&self) -> ScopeKind {
match self.id {
NodeWithScopeId::Module => ScopeKind::Module,
NodeWithScopeId::Class(_) => ScopeKind::Class,
NodeWithScopeId::Function(_) => ScopeKind::Function,
NodeWithScopeId::ClassTypeParams(_) | NodeWithScopeId::FunctionTypeParams(_) => {
ScopeKind::Annotation
}
impl<'a> From<CurrentTarget<'a>> for DefinitionNodeRef<'a> {
fn from(val: CurrentTarget<'a>) -> Self {
match val {
CurrentTarget::Expr(expression) => DefinitionNodeRef::Target(expression),
CurrentTarget::ExprNamed(named) => DefinitionNodeRef::NamedExpression(named),
}
}
}