use std::collections::{HashMap, HashSet};
use syn::{Item, UseTree};
#[derive(Clone, Debug, Default)]
pub struct FileScope {
pub path: String,
pub name_to_path: HashMap<String, String>,
pub glob_imports: Vec<String>,
pub local_items: HashSet<String>,
}
impl FileScope {
pub fn from_syn_file(path: impl Into<String>, syn_file: &syn::File) -> Self {
let mut scope = Self {
path: path.into(),
..Default::default()
};
scope.visit_items(&syn_file.items);
scope
}
pub fn from_file_content(
path: impl Into<String>,
content: &str,
) -> Result<Self, syn::Error> {
let syn_file = syn::parse_file(content)?;
Ok(Self::from_syn_file(path, &syn_file))
}
pub fn resolve_name(&self, name: &str) -> Option<String> {
if let Some(path) = self.name_to_path.get(name) {
return Some(path.clone());
}
if self.local_items.contains(name) {
return Some(format!("self::{}", name));
}
None
}
pub fn resolve_name_candidates(&self, name: &str) -> Vec<String> {
if let Some(path) = self.name_to_path.get(name) {
return vec![path.clone()];
}
if self.local_items.contains(name) {
return vec![format!("self::{}", name)];
}
self.glob_imports
.iter()
.map(|m| format!("{}::{}", m, name))
.collect()
}
fn visit_items(&mut self, items: &[Item]) {
for item in items {
match item {
Item::Use(item_use) => {
self.collect_use_tree(&item_use.tree, "");
}
Item::Struct(s) => {
self.local_items.insert(s.ident.to_string());
}
Item::Enum(e) => {
self.local_items.insert(e.ident.to_string());
}
Item::Fn(f) => {
self.local_items.insert(f.sig.ident.to_string());
}
Item::Trait(t) => {
self.local_items.insert(t.ident.to_string());
}
Item::Type(ty) => {
self.local_items.insert(ty.ident.to_string());
}
Item::Const(c) => {
self.local_items.insert(c.ident.to_string());
}
Item::Static(s) => {
self.local_items.insert(s.ident.to_string());
}
Item::Mod(m) => {
self.local_items.insert(m.ident.to_string());
if let Some((_, ref sub_items)) = m.content {
self.visit_items(sub_items);
}
}
Item::Impl(impl_block) => {
if let syn::Type::Path(type_path) = &*impl_block.self_ty {
if let Some(last) = type_path.path.segments.last() {
self.local_items.insert(last.ident.to_string());
}
}
}
_ => {}
}
}
}
fn collect_use_tree(&mut self, tree: &UseTree, prefix: &str) {
match tree {
UseTree::Path(use_path) => {
let ident = use_path.ident.to_string();
let new_prefix = if prefix.is_empty() {
ident
} else {
format!("{}::{}", prefix, ident)
};
self.collect_use_tree(&use_path.tree, &new_prefix);
}
UseTree::Name(use_name) => {
let name = use_name.ident.to_string();
let full_path = if prefix.is_empty() {
name.clone()
} else {
format!("{}::{}", prefix, name)
};
self.name_to_path.insert(name, full_path);
}
UseTree::Rename(use_rename) => {
let original = use_rename.ident.to_string();
let alias = use_rename.rename.to_string();
let full_path = if prefix.is_empty() {
original
} else {
format!("{}::{}", prefix, original)
};
self.name_to_path.insert(alias, full_path);
}
UseTree::Glob(_) => {
if !prefix.is_empty() {
self.glob_imports.push(prefix.to_string());
}
}
UseTree::Group(use_group) => {
for sub_tree in &use_group.items {
self.collect_use_tree(sub_tree, prefix);
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn build(content: &str) -> FileScope {
FileScope::from_file_content("test.rs", content).unwrap()
}
#[test]
fn test_simple_use() {
let scope = build("use crate::state::State;");
assert_eq!(
scope.resolve_name("State"),
Some("crate::state::State".to_string())
);
}
#[test]
fn test_grouped_use() {
let scope = build("use foo::{Bar, Baz};");
assert_eq!(scope.resolve_name("Bar"), Some("foo::Bar".to_string()));
assert_eq!(scope.resolve_name("Baz"), Some("foo::Baz".to_string()));
}
#[test]
fn test_rename_use() {
let scope = build("use foo::Bar as Qux;");
assert_eq!(scope.resolve_name("Qux"), Some("foo::Bar".to_string()));
assert_eq!(scope.resolve_name("Bar"), None);
}
#[test]
fn test_glob_use() {
let scope = build("use anchor_lang::prelude::*;");
assert_eq!(scope.glob_imports, vec!["anchor_lang::prelude".to_string()]);
let candidates = scope.resolve_name_candidates("Pubkey");
assert_eq!(candidates, vec!["anchor_lang::prelude::Pubkey".to_string()]);
}
#[test]
fn test_nested_grouped() {
let scope = build(
"use foo::{
bar::{Baz, Qux},
quux::Corge as Grault,
};",
);
assert_eq!(scope.resolve_name("Baz"), Some("foo::bar::Baz".to_string()));
assert_eq!(scope.resolve_name("Qux"), Some("foo::bar::Qux".to_string()));
assert_eq!(
scope.resolve_name("Grault"),
Some("foo::quux::Corge".to_string())
);
}
#[test]
fn test_local_struct() {
let scope = build("pub struct Initialize<'info> { pub state: String }");
assert!(scope.local_items.contains("Initialize"));
assert_eq!(
scope.resolve_name("Initialize"),
Some("self::Initialize".to_string())
);
}
#[test]
fn test_inline_module() {
let scope = build(
r#"
pub mod inner {
use crate::foo::Bar;
}
"#,
);
assert_eq!(scope.resolve_name("Bar"), Some("crate::foo::Bar".to_string()));
assert!(scope.local_items.contains("inner"));
}
#[test]
fn test_anchor_program_pattern() {
let scope = build(
r#"
use crate::state::State;
use anchor_lang::prelude::*;
#[program]
pub mod marinade_finance {
use super::*;
pub fn initialize(ctx: Context<Initialize>) -> Result<()> { Ok(()) }
}
"#,
);
assert_eq!(
scope.resolve_name("State"),
Some("crate::state::State".to_string())
);
assert!(scope.glob_imports.contains(&"anchor_lang::prelude".to_string()));
}
}