diff --git a/crates/red_knot/src/db.rs b/crates/red_knot/src/db.rs index c4b5b67394..9cea81808c 100644 --- a/crates/red_knot/src/db.rs +++ b/crates/red_knot/src/db.rs @@ -26,10 +26,13 @@ pub trait SemanticDb: SourceDb { // queries fn resolve_module(&self, name: ModuleName) -> Option; + fn file_to_module(&self, file_id: FileId) -> Option; + + fn path_to_module(&self, path: &Path) -> Option; + fn symbol_table(&self, file_id: FileId) -> Arc; // mutations - fn path_to_module(&mut self, path: &Path) -> Option; fn add_module(&mut self, path: &Path) -> Option<(Module, Vec>)>; @@ -80,8 +83,8 @@ pub(crate) mod tests { use crate::files::{FileId, Files}; use crate::lint::{lint_syntax, Diagnostics}; use crate::module::{ - add_module, path_to_module, resolve_module, set_module_search_paths, Module, ModuleData, - ModuleName, ModuleSearchPath, + add_module, file_to_module, path_to_module, resolve_module, set_module_search_paths, + Module, ModuleData, ModuleName, ModuleSearchPath, }; use crate::parse::{parse, Parsed}; use crate::source::{source_text, Source}; @@ -148,14 +151,18 @@ pub(crate) mod tests { resolve_module(self, name) } - fn symbol_table(&self, file_id: FileId) -> Arc { - symbol_table(self, file_id) + fn file_to_module(&self, file_id: FileId) -> Option { + file_to_module(self, file_id) } - fn path_to_module(&mut self, path: &Path) -> Option { + fn path_to_module(&self, path: &Path) -> Option { path_to_module(self, path) } + fn symbol_table(&self, file_id: FileId) -> Arc { + symbol_table(self, file_id) + } + fn add_module(&mut self, path: &Path) -> Option<(Module, Vec>)> { add_module(self, path) } diff --git a/crates/red_knot/src/files.rs b/crates/red_knot/src/files.rs index defe4a2ee4..69b72281cf 100644 --- a/crates/red_knot/src/files.rs +++ b/crates/red_knot/src/files.rs @@ -23,7 +23,7 @@ pub struct Files { } impl Files { - #[tracing::instrument(level = "debug", skip(path))] + #[tracing::instrument(level = "debug", skip(self))] pub fn intern(&self, path: &Path) -> FileId { self.inner.write().intern(path) } @@ -32,7 +32,7 @@ impl Files { self.inner.read().try_get(path) } - #[tracing::instrument(level = "debug")] + #[tracing::instrument(level = "debug", skip(self))] pub fn path(&self, id: FileId) -> Arc { self.inner.read().path(id) } diff --git a/crates/red_knot/src/lib.rs b/crates/red_knot/src/lib.rs index 995a957bed..bf321f8e3b 100644 --- a/crates/red_knot/src/lib.rs +++ b/crates/red_knot/src/lib.rs @@ -1,4 +1,5 @@ use std::hash::BuildHasherDefault; +use std::ops::Deref; use std::path::{Path, PathBuf}; use rustc_hash::{FxHashSet, FxHasher}; @@ -81,3 +82,21 @@ impl Name { self.0.as_str() } } + +impl Deref for Name { + type Target = str; + + #[inline] + fn deref(&self) -> &Self::Target { + self.as_str() + } +} + +impl From for Name +where + T: Into, +{ + fn from(value: T) -> Self { + Self(value.into()) + } +} diff --git a/crates/red_knot/src/main.rs b/crates/red_knot/src/main.rs index 016cb5a3ea..9b1c0c0d03 100644 --- a/crates/red_knot/src/main.rs +++ b/crates/red_knot/src/main.rs @@ -151,7 +151,7 @@ impl MainLoop { .unwrap(); rayon::in_place_scope(|scope| { - let scheduler = RayonCheckScheduler { program, scope }; + let scheduler = RayonCheckScheduler::new(program, scope); let result = program.check(&scheduler, run_cancellation_token); match result { @@ -422,6 +422,7 @@ fn setup_tracing() { .with_indent_lines(true) .with_indent_amount(2) .with_bracketed_fields(true) + .with_thread_ids(true) .with_targets(true) .with_writer(|| Box::new(std::io::stderr())) .with_timer(Uptime::default()) diff --git a/crates/red_knot/src/module.rs b/crates/red_knot/src/module.rs index 62f48e5330..754b108db8 100644 --- a/crates/red_knot/src/module.rs +++ b/crates/red_knot/src/module.rs @@ -4,6 +4,7 @@ use std::sync::atomic::AtomicU32; use std::sync::Arc; use dashmap::mapref::entry::Entry; +use smol_str::SmolStr; use crate::db::{HasJar, SemanticDb, SemanticJar}; use crate::files::FileId; @@ -31,6 +32,55 @@ impl Module { modules.modules.get(self).unwrap().path.clone() } + + pub fn kind(&self, db: &Db) -> ModuleKind + where + Db: HasJar, + { + let modules = &db.jar().module_resolver; + + modules.modules.get(self).unwrap().kind + } + + pub fn relative_name(&self, db: &Db, level: u32, module: Option<&str>) -> Option + where + Db: HasJar, + { + let name = self.name(db); + let kind = self.kind(db); + + let mut components = name.components().peekable(); + + if level > 0 { + let start = match kind { + // `.` resolves to the enclosing package + ModuleKind::Module => 0, + // `.` resolves to the current package + ModuleKind::Package => 1, + }; + + // Skip over the relative parts. + for _ in start..level { + components.next_back()?; + } + } + + let mut name = String::new(); + + for part in components.chain(module) { + if !name.is_empty() { + name.push('.'); + } + + name.push_str(part); + } + + if name.is_empty() { + None + } else { + Some(ModuleName(SmolStr::new(name))) + } + } } /// A module name, e.g. `foo.bar`. @@ -46,12 +96,7 @@ impl ModuleName { Self(smol_str::SmolStr::new(name)) } - pub fn relative(_dots: u32, name: &str, _to: &Path) -> Self { - // FIXME: Take `to` and `dots` into account. - Self(smol_str::SmolStr::new(name)) - } - - pub fn from_relative_path(path: &Path) -> Option { + fn from_relative_path(path: &Path) -> Option { let path = if path.ends_with("__init__.py") || path.ends_with("__init__.pyi") { path.parent()? } else { @@ -102,6 +147,14 @@ impl std::fmt::Display for ModuleName { } } +#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)] +pub enum ModuleKind { + Module, + + /// A python package (a `__init__.py` or `__init__.pyi` file) + Package, +} + /// A search path in which to search modules. /// Corresponds to a path in [`sys.path`](https://docs.python.org/3/library/sys_path_init.html) at runtime. /// @@ -151,10 +204,17 @@ pub enum ModuleSearchPathKind { StandardLibrary, } +impl ModuleSearchPathKind { + pub const fn is_first_party(self) -> bool { + matches!(self, Self::FirstParty) + } +} + #[derive(Debug, Eq, PartialEq)] pub struct ModuleData { name: ModuleName, path: ModulePath, + kind: ModuleKind, } ////////////////////////////////////////////////////// @@ -177,7 +237,7 @@ where match entry { Entry::Occupied(entry) => Some(*entry.get()), Entry::Vacant(entry) => { - let (root_path, absolute_path) = resolve_name(&name, &modules.search_paths)?; + let (root_path, absolute_path, kind) = resolve_name(&name, &modules.search_paths)?; let normalized = absolute_path.canonicalize().ok()?; let file_id = db.file_id(&normalized); @@ -191,7 +251,7 @@ where modules .modules - .insert(id, Arc::from(ModuleData { name, path })); + .insert(id, Arc::from(ModuleData { name, path, kind })); // A path can map to multiple modules because of symlinks: // ``` @@ -209,24 +269,11 @@ where } } -////////////////////////////////////////////////////// -// Mutations -////////////////////////////////////////////////////// - -/// Changes the module search paths to `search_paths`. -pub fn set_module_search_paths(db: &mut Db, search_paths: Vec) -where - Db: SemanticDb + HasJar, -{ - let jar = db.jar_mut(); - - jar.module_resolver = ModuleResolver::new(search_paths); -} - /// Resolves the module id for the file with the given id. /// /// Returns `None` if the file is not a module in `sys.path`. -pub fn file_to_module(db: &mut Db, file: FileId) -> Option +#[tracing::instrument(level = "debug", skip(db))] +pub fn file_to_module(db: &Db, file: FileId) -> Option where Db: SemanticDb + HasJar, { @@ -237,28 +284,24 @@ where /// Resolves the module id for the given path. /// /// Returns `None` if the path is not a module in `sys.path`. -// WARNING!: It's important that this method takes `&mut self`. Without, the implementation is prone to race conditions. -// Note: This won't work with salsa because `Path` is not an ingredient. -pub fn path_to_module(db: &mut Db, path: &Path) -> Option +#[tracing::instrument(level = "debug", skip(db))] +pub fn path_to_module(db: &Db, path: &Path) -> Option where Db: SemanticDb + HasJar, { - let jar = db.jar_mut(); - let modules = &mut jar.module_resolver; + let jar = db.jar(); + let modules = &jar.module_resolver; debug_assert!(path.is_absolute()); if let Some(existing) = modules.by_path.get(path) { return Some(*existing); } - let root_path = modules - .search_paths - .iter() - .find(|root| path.starts_with(root.path()))? - .clone(); + let (root_path, relative_path) = modules.search_paths.iter().find_map(|root| { + let relative_path = path.strip_prefix(root.path()).ok()?; + Some((root.clone(), relative_path)) + })?; - // SAFETY: `strip_prefix` is guaranteed to succeed because we search the root that is a prefix of the path. - let relative_path = path.strip_prefix(root_path.path()).unwrap(); let module_name = ModuleName::from_relative_path(relative_path)?; // Resolve the module name to see if Python would resolve the name to the same path. @@ -266,7 +309,6 @@ where // root paths, but that the module corresponding to the past path is in a lower priority path, // in which case we ignore it. let module_id = resolve_module(db, module_name)?; - // Note: Guaranteed to be race-free because we're holding a mutable reference of `self` here. let module_path = module_id.path(db); if module_path.root() == &root_path { @@ -293,6 +335,20 @@ where } } +////////////////////////////////////////////////////// +// Mutations +////////////////////////////////////////////////////// + +/// Changes the module search paths to `search_paths`. +pub fn set_module_search_paths(db: &mut Db, search_paths: Vec) +where + Db: SemanticDb + HasJar, +{ + let jar = db.jar_mut(); + + jar.module_resolver = ModuleResolver::new(search_paths); +} + /// Adds a module to the resolver. /// /// Returns `None` if the path doesn't resolve to a module. @@ -444,7 +500,7 @@ impl ModulePath { fn resolve_name( name: &ModuleName, search_paths: &[ModuleSearchPath], -) -> Option<(ModuleSearchPath, PathBuf)> { +) -> Option<(ModuleSearchPath, PathBuf, ModuleKind)> { for search_path in search_paths { let mut components = name.components(); let module_name = components.next_back()?; @@ -456,21 +512,24 @@ fn resolve_name( package_path.push(module_name); // Must be a `__init__.pyi` or `__init__.py` or it isn't a package. - if package_path.is_dir() { + let kind = if package_path.is_dir() { package_path.push("__init__"); - } + ModuleKind::Package + } else { + ModuleKind::Module + }; // TODO Implement full https://peps.python.org/pep-0561/#type-checker-module-resolution-order resolution let stub = package_path.with_extension("pyi"); if stub.is_file() { - return Some((search_path.clone(), stub)); + return Some((search_path.clone(), stub, kind)); } let module = package_path.with_extension("py"); if module.is_file() { - return Some((search_path.clone(), module)); + return Some((search_path.clone(), module, kind)); } // For regular packages, don't search the next search path. All files of that @@ -578,7 +637,7 @@ impl PackageKind { mod tests { use crate::db::tests::TestDb; use crate::db::{SemanticDb, SourceDb}; - use crate::module::{ModuleName, ModuleSearchPath, ModuleSearchPathKind}; + use crate::module::{ModuleKind, ModuleName, ModuleSearchPath, ModuleSearchPathKind}; struct TestCase { temp_dir: tempfile::TempDir, @@ -619,7 +678,7 @@ mod tests { #[test] fn first_party_module() -> std::io::Result<()> { let TestCase { - mut db, + db, src, temp_dir: _temp_dir, .. @@ -634,6 +693,7 @@ mod tests { assert_eq!(ModuleName::new("foo"), foo_module.name(&db)); assert_eq!(&src, foo_module.path(&db).root()); + assert_eq!(ModuleKind::Module, foo_module.kind(&db)); assert_eq!(&foo_path, &*db.file_path(foo_module.path(&db).file())); assert_eq!(Some(foo_module), db.path_to_module(&foo_path)); @@ -645,7 +705,7 @@ mod tests { fn resolve_package() -> std::io::Result<()> { let TestCase { src, - mut db, + db, temp_dir: _temp_dir, .. } = create_resolver()?; @@ -672,7 +732,7 @@ mod tests { #[test] fn package_priority_over_module() -> std::io::Result<()> { let TestCase { - mut db, + db, temp_dir: _temp_dir, src, .. @@ -690,6 +750,7 @@ mod tests { assert_eq!(&src, foo_module.path(&db).root()); assert_eq!(&foo_init, &*db.file_path(foo_module.path(&db).file())); + assert_eq!(ModuleKind::Package, foo_module.kind(&db)); assert_eq!(Some(foo_module), db.path_to_module(&foo_init)); assert_eq!(None, db.path_to_module(&foo_py)); @@ -700,7 +761,7 @@ mod tests { #[test] fn typing_stub_over_module() -> std::io::Result<()> { let TestCase { - mut db, + db, src, temp_dir: _temp_dir, .. @@ -725,7 +786,7 @@ mod tests { #[test] fn sub_packages() -> std::io::Result<()> { let TestCase { - mut db, + db, src, temp_dir: _temp_dir, .. @@ -753,7 +814,7 @@ mod tests { #[test] fn namespace_package() -> std::io::Result<()> { let TestCase { - mut db, + db, temp_dir: _, src, site_packages, @@ -803,7 +864,7 @@ mod tests { #[test] fn regular_package_in_namespace_package() -> std::io::Result<()> { let TestCase { - mut db, + db, temp_dir: _, src, site_packages, @@ -850,7 +911,7 @@ mod tests { #[test] fn module_search_path_priority() -> std::io::Result<()> { let TestCase { - mut db, + db, src, site_packages, temp_dir: _temp_dir, @@ -877,7 +938,7 @@ mod tests { #[cfg(target_family = "unix")] fn symlink() -> std::io::Result<()> { let TestCase { - mut db, + db, src, temp_dir: _temp_dir, .. @@ -908,4 +969,66 @@ mod tests { Ok(()) } + + #[test] + fn relative_name() -> std::io::Result<()> { + let TestCase { + src, + db, + temp_dir: _temp_dir, + .. + } = create_resolver()?; + + let foo_dir = src.path().join("foo"); + let foo_path = foo_dir.join("__init__.py"); + let bar_path = foo_dir.join("bar.py"); + + std::fs::create_dir(&foo_dir)?; + std::fs::write(foo_path, "from .bar import test")?; + std::fs::write(bar_path, "test = 'Hello world'")?; + + let foo_module = db.resolve_module(ModuleName::new("foo")).unwrap(); + let bar_module = db.resolve_module(ModuleName::new("foo.bar")).unwrap(); + + // `from . import bar` in `foo/__init__.py` resolves to `foo` + assert_eq!( + Some(ModuleName::new("foo")), + foo_module.relative_name(&db, 1, None) + ); + + // `from baz import bar` in `foo/__init__.py` should resolve to `foo/baz.py` + assert_eq!( + Some(ModuleName::new("foo.baz")), + foo_module.relative_name(&db, 0, Some("baz")) + ); + + // from .bar import test in `foo/__init__.py` should resolve to `foo/bar.py` + assert_eq!( + Some(ModuleName::new("foo.bar")), + foo_module.relative_name(&db, 1, Some("bar")) + ); + + // from .. import test in `foo/__init__.py` resolves to `` which is not a module + assert_eq!(None, foo_module.relative_name(&db, 2, None)); + + // `from . import test` in `foo/bar.py` resolves to `foo` + assert_eq!( + Some(ModuleName::new("foo")), + bar_module.relative_name(&db, 1, None) + ); + + // `from baz import test` in `foo/bar.py` resolves to `foo.bar.baz` + assert_eq!( + Some(ModuleName::new("foo.bar.baz")), + bar_module.relative_name(&db, 0, Some("baz")) + ); + + // `from .baz import test` in `foo/bar.py` resolves to `foo.baz`. + assert_eq!( + Some(ModuleName::new("foo.baz")), + bar_module.relative_name(&db, 1, Some("baz")) + ); + + Ok(()) + } } diff --git a/crates/red_knot/src/program/check.rs b/crates/red_knot/src/program/check.rs index 1d56caa5bf..0154a17d61 100644 --- a/crates/red_knot/src/program/check.rs +++ b/crates/red_knot/src/program/check.rs @@ -1,5 +1,5 @@ use crate::cancellation::CancellationToken; -use crate::db::SourceDb; +use crate::db::{SemanticDb, SourceDb}; use crate::files::FileId; use crate::lint::Diagnostics; use crate::program::Program; @@ -41,7 +41,32 @@ impl Program { ) -> Result { context.cancelled_ok()?; - // TODO schedule the dependencies. + let symbol_table = self.symbol_table(file); + let dependencies = symbol_table.dependencies(); + + if !dependencies.is_empty() { + let module = self.file_to_module(file); + + // TODO scheduling all dependencies here is wasteful if we don't infer any types on them + // but I think that's unlikely, so it is okay? + // Anyway, we need to figure out a way to retrieve the dependencies of a module + // from the persistent cache. So maybe it should be a separate query after all. + for dependency in dependencies { + let dependency_name = dependency.module_name(self, module); + + if let Some(dependency_name) = dependency_name { + // TODO We may want to have a different check functions for non-first-party + // files because we only need to index them and not check them. + // Supporting non-first-party code also requires supporting typing stubs. + if let Some(dependency) = self.resolve_module(dependency_name) { + if dependency.path(self).root().kind().is_first_party() { + context.schedule_check_file(dependency.path(self).file()); + } + } + } + } + } + let mut diagnostics = Vec::new(); if self.workspace().is_file_open(file) { @@ -72,8 +97,8 @@ pub trait CheckScheduler { /// Scheduler that runs checks on a rayon thread pool. pub struct RayonCheckScheduler<'program, 'scope_ref, 'scope> { - pub program: &'program Program, - pub scope: &'scope_ref rayon::Scope<'scope>, + program: &'program Program, + scope: &'scope_ref rayon::Scope<'scope>, } impl<'program, 'scope_ref, 'scope> RayonCheckScheduler<'program, 'scope_ref, 'scope> { diff --git a/crates/red_knot/src/program/mod.rs b/crates/red_knot/src/program/mod.rs index f0dac73014..0070d53086 100644 --- a/crates/red_knot/src/program/mod.rs +++ b/crates/red_knot/src/program/mod.rs @@ -7,8 +7,8 @@ use crate::db::{Db, HasJar, SemanticDb, SemanticJar, SourceDb, SourceJar}; use crate::files::{FileId, Files}; use crate::lint::{lint_syntax, Diagnostics, LintSyntaxStorage}; use crate::module::{ - add_module, path_to_module, resolve_module, set_module_search_paths, Module, ModuleData, - ModuleName, ModuleResolver, ModuleSearchPath, + add_module, file_to_module, path_to_module, resolve_module, set_module_search_paths, Module, + ModuleData, ModuleName, ModuleResolver, ModuleSearchPath, }; use crate::parse::{parse, Parsed, ParsedStorage}; use crate::source::{source_text, Source, SourceStorage}; @@ -99,14 +99,19 @@ impl SemanticDb for Program { resolve_module(self, name) } + fn file_to_module(&self, file_id: FileId) -> Option { + file_to_module(self, file_id) + } + + fn path_to_module(&self, path: &Path) -> Option { + path_to_module(self, path) + } + fn symbol_table(&self, file_id: FileId) -> Arc { symbol_table(self, file_id) } // Mutations - fn path_to_module(&mut self, path: &Path) -> Option { - path_to_module(self, path) - } fn add_module(&mut self, path: &Path) -> Option<(Module, Vec>)> { add_module(self, path) diff --git a/crates/red_knot/src/symbols.rs b/crates/red_knot/src/symbols.rs index 0277e51d3b..6f82522d36 100644 --- a/crates/red_knot/src/symbols.rs +++ b/crates/red_knot/src/symbols.rs @@ -16,6 +16,7 @@ use crate::ast_ids::TypedNodeKey; use crate::cache::KeyValueCache; use crate::db::{HasJar, SemanticDb, SemanticJar}; use crate::files::FileId; +use crate::module::{Module, ModuleName}; use crate::Name; #[allow(unreachable_pub)] @@ -110,22 +111,43 @@ pub(crate) enum Definition { #[derive(Debug)] pub(crate) struct ImportDefinition { - pub(crate) module: String, + pub(crate) module: Name, } #[derive(Debug)] pub(crate) struct ImportFromDefinition { - pub(crate) module: Option, - pub(crate) name: String, + pub(crate) module: Option, + pub(crate) name: Name, pub(crate) level: u32, } +#[derive(Debug, Clone)] +pub(crate) enum Dependency { + Module(Name), + Relative { level: u32, module: Option }, +} + +impl Dependency { + pub(crate) fn module_name(&self, db: &Db, relative_to: Option) -> Option + where + Db: SemanticDb + HasJar, + { + match self { + Dependency::Module(name) => Some(ModuleName::new(name.as_str())), + Dependency::Relative { level, module } => { + relative_to?.relative_name(db, *level, module.as_deref()) + } + } + } +} + /// Table of all symbols in all scopes for a module. #[derive(Debug)] pub struct SymbolTable { scopes_by_id: IndexVec, symbols_by_id: IndexVec, defs: FxHashMap>, + dependencies: Vec, } impl SymbolTable { @@ -144,6 +166,7 @@ impl SymbolTable { scopes_by_id: IndexVec::new(), symbols_by_id: IndexVec::new(), defs: FxHashMap::default(), + dependencies: Vec::new(), }; table.scopes_by_id.push(Scope { name: Name::new(""), @@ -154,6 +177,10 @@ impl SymbolTable { table } + pub(crate) fn dependencies(&self) -> &[Dependency] { + &self.dependencies + } + pub(crate) const fn root_scope_id() -> ScopeId { ScopeId::from_usize(0) } @@ -445,10 +472,14 @@ impl PreorderVisitor<'_> for SymbolTableBuilder { } else { alias.name.id.split('.').next().unwrap() }; + + let module = Name::new(&alias.name.id); + let def = Definition::Import(ImportDefinition { - module: alias.name.id.clone(), + module: module.clone(), }); self.add_symbol_with_def(symbol_name, def); + self.table.dependencies.push(Dependency::Module(module)); } } ast::Stmt::ImportFrom(ast::StmtImportFrom { @@ -457,6 +488,8 @@ impl PreorderVisitor<'_> for SymbolTableBuilder { level, .. }) => { + let module = module.as_ref().map(|m| Name::new(&m.id)); + for alias in names { let symbol_name = if let Some(asname) = &alias.asname { asname.id.as_str() @@ -464,12 +497,30 @@ impl PreorderVisitor<'_> for SymbolTableBuilder { alias.name.id.as_str() }; let def = Definition::ImportFrom(ImportFromDefinition { - module: module.as_ref().map(|m| m.id.clone()), - name: alias.name.id.clone(), + module: module.clone(), + name: Name::new(&alias.name.id), level: *level, }); self.add_symbol_with_def(symbol_name, def); } + + let dependency = if let Some(module) = module { + if *level == 0 { + Dependency::Module(module) + } else { + Dependency::Relative { + level: *level, + module: Some(module), + } + } + } else { + Dependency::Relative { + level: *level, + module, + } + }; + + self.table.dependencies.push(dependency); } _ => { ast::visitor::preorder::walk_stmt(self, stmt);