rust_assistant/
search.rs

1//! The `search` module.
2//!
3//! Focused on providing search functionalities within crates. This module might contain
4//! implementations for searching through crate contents, such as source code files, documentation,
5//! and other relevant data. It could include various search algorithms and data structures optimized
6//! for quick and efficient search operations, like `SearchIndex`.
7//!
8use fnv::FnvHashMap;
9use serde::{Deserialize, Serialize};
10use std::num::NonZeroUsize;
11use std::path::Path;
12use std::sync::Arc;
13use syn::spanned::Spanned;
14use syn::{Attribute, ItemEnum, ItemFn, ItemImpl, ItemMacro, ItemStruct, ItemTrait};
15
16use crate::{Item, ItemQuery, ItemType};
17
18/// A mutable search index containing categorized items for searching within a crate.
19///
20/// This struct stores various items like structs, enums, traits, etc., in categorized hash maps,
21/// facilitating efficient search operations.
22///
23#[derive(Debug, Default, Serialize, Deserialize)]
24pub struct SearchIndexMut {
25    pub structs: FnvHashMap<String, Vec<Item>>,
26    pub enums: FnvHashMap<String, Vec<Item>>,
27    pub traits: FnvHashMap<String, Vec<Item>>,
28    pub impl_types: FnvHashMap<String, Vec<Item>>,
29    pub impl_trait_for_types: FnvHashMap<String, Vec<Item>>,
30    pub macros: FnvHashMap<String, Vec<Item>>,
31    pub attribute_macros: FnvHashMap<String, Vec<Item>>,
32    pub functions: FnvHashMap<String, Vec<Item>>,
33    pub type_aliases: FnvHashMap<String, Vec<Item>>,
34}
35
36impl SearchIndexMut {
37    /// Searches for items within the index based on the provided query.
38    ///
39    pub fn search(&self, query: &ItemQuery) -> Vec<Item> {
40        let ItemQuery { type_, query, path } = query;
41        let query = query.to_lowercase();
42        let path = path.as_ref().map(|p| p.as_path());
43        match type_ {
44            ItemType::All => {
45                let mut all = Vec::new();
46                all.extend(filter_items(&query, &self.structs, path));
47                all.extend(filter_items(&query, &self.enums, path));
48                all.extend(filter_items(&query, &self.traits, path));
49                all.extend(filter_items(&query, &self.impl_types, path));
50                all.extend(filter_items(&query, &self.impl_trait_for_types, path));
51                all.extend(filter_items(&query, &self.macros, path));
52                all.extend(filter_items(&query, &self.attribute_macros, path));
53                all.extend(filter_items(&query, &self.functions, path));
54                all.extend(filter_items(&query, &self.type_aliases, path));
55                all
56            }
57            ItemType::Struct => filter_items(&query, &self.structs, path),
58            ItemType::Enum => filter_items(&query, &self.enums, path),
59            ItemType::Trait => filter_items(&query, &self.traits, path),
60            ItemType::ImplType => filter_items(&query, &self.impl_types, path),
61            ItemType::ImplTraitForType => filter_items(&query, &self.impl_trait_for_types, path),
62            ItemType::Macro => filter_items(&query, &self.macros, path),
63            ItemType::AttributeMacro => filter_items(&query, &self.attribute_macros, path),
64            ItemType::Function => filter_items(&query, &self.functions, path),
65            ItemType::TypeAlias => filter_items(&query, &self.type_aliases, path),
66        }
67    }
68}
69
70/// Filters items from a hashmap based on a query and optional path.
71///
72fn filter_items(
73    query: &str,
74    items: &FnvHashMap<String, Vec<Item>>,
75    path: Option<&Path>,
76) -> Vec<Item> {
77    let flatten = items
78        .iter()
79        .filter(|(name, _)| name.contains(&query))
80        .map(|(_, item)| item)
81        .flatten();
82    match path {
83        None => flatten.cloned().collect::<Vec<Item>>(),
84        Some(path) => flatten
85            .filter(|item| item.file.starts_with(path))
86            .cloned()
87            .collect::<Vec<Item>>(),
88    }
89}
90
91/// Shared immutable search index, used for efficient read access across multiple threads.
92pub type SearchIndex = Arc<SearchIndexMut>;
93
94impl SearchIndexMut {
95    /// Freezes the mutable search index into an immutable one.
96    ///
97    pub fn freeze(self) -> SearchIndex {
98        Arc::new(self)
99    }
100}
101
102/// A builder for constructing a `SearchIndex`.
103///
104/// This struct facilitates the creation and population of a `SearchIndexMut`
105/// by parsing Rust source files and adding items to the index.
106#[derive(Debug, Default)]
107pub struct SearchIndexBuilder {
108    index: SearchIndexMut,
109}
110
111impl SearchIndexBuilder {
112    /// Updates the search index with items parsed from a Rust source file.
113    ///
114    pub fn update<P: AsRef<Path>>(&mut self, file: P, content: &str) -> bool {
115        let mut visitor = IndexVisitor::new(&mut self.index, file);
116        if let Ok(ast) = syn::parse_file(content) {
117            syn::visit::visit_file(&mut visitor, &ast);
118            true
119        } else {
120            false
121        }
122    }
123
124    /// Finalizes the construction of the `SearchIndex`.
125    ///
126    pub fn finish(self) -> SearchIndex {
127        self.index.freeze()
128    }
129}
130
131/// A visitor struct for traversing and indexing Rust syntax trees.
132///
133/// This struct is used in conjunction with `syn::visit::Visit` to extract items from Rust source files
134/// and add them to a `SearchIndexMut`.
135pub struct IndexVisitor<'i> {
136    index: &'i mut SearchIndexMut,
137    current_file: Arc<Path>,
138}
139
140impl<'i> IndexVisitor<'i> {
141    /// Creates a new `IndexVisitor`.
142    ///
143    pub fn new<P: AsRef<Path>>(index: &'i mut SearchIndexMut, current_file: P) -> Self {
144        IndexVisitor {
145            index,
146            current_file: Arc::from(current_file.as_ref()),
147        }
148    }
149
150    fn create_item(
151        &self,
152        name: String,
153        type_: ItemType,
154        item_span: proc_macro2::Span,
155        attrs: &[Attribute],
156    ) -> Item {
157        // 获取项的 span
158        let mut start_line = item_span.start().line;
159        let end_line = item_span.end().line;
160
161        // 检查并调整起始行号以包含文档注释
162        for attr in attrs {
163            if attr.path().is_ident("doc") {
164                let attr_span = attr.span();
165                start_line = start_line.min(attr_span.start().line);
166            }
167        }
168
169        let start_line = NonZeroUsize::new(start_line).unwrap_or(NonZeroUsize::MIN);
170        let end_line = NonZeroUsize::new(end_line).unwrap_or(NonZeroUsize::MAX);
171
172        Item {
173            name,
174            type_,
175            file: self.current_file.clone(),
176            line_range: start_line..=end_line,
177        }
178    }
179}
180
181impl<'i, 'ast> syn::visit::Visit<'ast> for IndexVisitor<'i> {
182    fn visit_item_enum(&mut self, i: &'ast ItemEnum) {
183        let name = i.ident.to_string();
184        let item = self.create_item(name, ItemType::Enum, i.span(), &i.attrs);
185        self.index
186            .enums
187            .entry(item.name.to_lowercase())
188            .or_default()
189            .push(item);
190    }
191
192    fn visit_item_fn(&mut self, i: &'ast ItemFn) {
193        if is_attribute_macro(&i.attrs) {
194            let name = i.sig.ident.to_string();
195            let item = self.create_item(name, ItemType::AttributeMacro, i.span(), &i.attrs);
196            self.index
197                .attribute_macros
198                .entry(item.name.to_lowercase())
199                .or_default()
200                .push(item);
201        } else {
202            let name = i.sig.ident.to_string();
203            let item = self.create_item(name, ItemType::Function, i.span(), &i.attrs);
204            self.index
205                .functions
206                .entry(item.name.to_lowercase())
207                .or_default()
208                .push(item);
209        }
210    }
211
212    fn visit_item_impl(&mut self, i: &'ast ItemImpl) {
213        let self_ty = &i.self_ty;
214
215        match &i.trait_ {
216            Some((_, path, _)) => {
217                // impl Trait for Type
218                let impl_name = format!(
219                    "impl {} for {}",
220                    quote::quote! { #path },
221                    quote::quote! { #self_ty }
222                );
223                let item =
224                    self.create_item(impl_name, ItemType::ImplTraitForType, i.span(), &i.attrs);
225                self.index
226                    .impl_trait_for_types
227                    .entry(item.name.to_lowercase())
228                    .or_default()
229                    .push(item);
230            }
231            None => {
232                // impl Type
233                let impl_name = format!("impl {}", quote::quote! { #self_ty });
234                let item = self.create_item(impl_name, ItemType::ImplType, i.span(), &i.attrs);
235                self.index
236                    .impl_types
237                    .entry(item.name.to_lowercase())
238                    .or_default()
239                    .push(item);
240            }
241        };
242    }
243
244    fn visit_item_macro(&mut self, i: &'ast ItemMacro) {
245        if let Some(ident) = &i.ident {
246            let name = ident.to_string();
247            let item = self.create_item(name, ItemType::Macro, i.span(), &i.attrs);
248            self.index
249                .macros
250                .entry(item.name.to_lowercase())
251                .or_default()
252                .push(item);
253        }
254    }
255
256    fn visit_item_struct(&mut self, i: &'ast ItemStruct) {
257        let name = i.ident.to_string();
258        let item = self.create_item(name, ItemType::Struct, i.span(), &i.attrs);
259        self.index
260            .structs
261            .entry(item.name.to_lowercase())
262            .or_default()
263            .push(item);
264    }
265
266    fn visit_item_trait(&mut self, i: &'ast ItemTrait) {
267        let name = i.ident.to_string();
268        let item = self.create_item(name, ItemType::Trait, i.span(), &i.attrs);
269        self.index
270            .traits
271            .entry(item.name.to_lowercase())
272            .or_default()
273            .push(item);
274    }
275
276    fn visit_item_type(&mut self, i: &'ast syn::ItemType) {
277        let name = i.ident.to_string();
278        let item = self.create_item(name, ItemType::TypeAlias, i.span(), &i.attrs);
279        self.index
280            .type_aliases
281            .entry(item.name.to_lowercase())
282            .or_default()
283            .push(item);
284    }
285}
286
287fn is_attribute_macro(attrs: &[Attribute]) -> bool {
288    attrs.iter().any(|attr| {
289        // check proc_macro_attribute
290        attr.path().is_ident("proc_macro_attribute")
291    })
292}
rust_assistant/search.rs

rust_assistant/
search.rs
