rust_docs_mcp/docs/
query.rs

1use anyhow::{Context, Result};
2use rmcp::schemars;
3use rustdoc_types::{Crate, Id, Item, ItemEnum};
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6
7/// Query interface for rustdoc JSON data
8#[derive(Debug)]
9pub struct DocQuery {
10    crate_data: Crate,
11}
12
13/// Simplified item information for API responses
14#[derive(Debug, Serialize, Deserialize, JsonSchema)]
15pub struct ItemInfo {
16    pub id: String,
17    pub name: String,
18    pub kind: String,
19    pub path: Vec<String>,
20    pub docs: Option<String>,
21    pub visibility: String,
22}
23
24/// Source location information
25#[derive(Debug, Serialize, Deserialize, JsonSchema)]
26pub struct SourceLocation {
27    pub filename: String,
28    pub line_start: usize,
29    pub column_start: usize,
30    pub line_end: usize,
31    pub column_end: usize,
32}
33
34/// Source code information for an item
35#[derive(Debug, Serialize, Deserialize, JsonSchema)]
36pub struct SourceInfo {
37    pub location: SourceLocation,
38    pub code: String,
39    pub context_lines: Option<usize>,
40}
41
42/// Detailed item information including signatures
43#[derive(Debug, Serialize, Deserialize, JsonSchema)]
44pub struct DetailedItem {
45    pub info: ItemInfo,
46    pub signature: Option<String>,
47    pub generics: Option<serde_json::Value>,
48    pub fields: Option<Vec<ItemInfo>>,
49    pub variants: Option<Vec<ItemInfo>>,
50    pub methods: Option<Vec<ItemInfo>>,
51    pub source_location: Option<SourceLocation>,
52}
53
54impl DocQuery {
55    /// Create a new query interface for a crate's documentation
56    pub fn new(crate_data: Crate) -> Self {
57        Self { crate_data }
58    }
59
60    /// List all items in the crate, optionally filtered by kind
61    pub fn list_items(&self, kind_filter: Option<&str>) -> Vec<ItemInfo> {
62        let mut items = Vec::new();
63
64        for (id, item) in &self.crate_data.index {
65            if let Some(filter) = &kind_filter
66                && self.get_item_kind_string(&item.inner) != *filter
67            {
68                continue;
69            }
70
71            if let Some(info) = self.item_to_info(id, item) {
72                items.push(info);
73            }
74        }
75
76        // Sort by path and name for consistent output
77        items.sort_by(|a, b| a.path.cmp(&b.path).then_with(|| a.name.cmp(&b.name)));
78        items
79    }
80
81    /// Search for items by name pattern
82    pub fn search_items(&self, pattern: &str) -> Vec<ItemInfo> {
83        let pattern_lower = pattern.to_lowercase();
84        let mut items = Vec::new();
85
86        for (id, item) in &self.crate_data.index {
87            // First check if item has a direct name
88            let item_name = if let Some(name) = &item.name {
89                Some(name.clone())
90            } else if let Some(path_summary) = self.crate_data.paths.get(id) {
91                // Fall back to using the last component of the path
92                path_summary.path.last().cloned()
93            } else {
94                None
95            };
96
97            if let Some(name) = item_name
98                && name.to_lowercase().contains(&pattern_lower)
99                && let Some(info) = self.item_to_info(id, item)
100            {
101                items.push(info);
102            }
103        }
104
105        items.sort_by(|a, b| {
106            // Sort by relevance (exact match first, then prefix match, then contains)
107            let a_exact = a.name.to_lowercase() == pattern_lower;
108            let b_exact = b.name.to_lowercase() == pattern_lower;
109            let a_prefix = a.name.to_lowercase().starts_with(&pattern_lower);
110            let b_prefix = b.name.to_lowercase().starts_with(&pattern_lower);
111
112            b_exact
113                .cmp(&a_exact)
114                .then_with(|| b_prefix.cmp(&a_prefix))
115                .then_with(|| a.name.len().cmp(&b.name.len()))
116                .then_with(|| a.name.cmp(&b.name))
117        });
118
119        items
120    }
121
122    /// Get detailed information about a specific item by ID
123    pub fn get_item_details(&self, item_id: u32) -> Result<DetailedItem> {
124        let id = Id(item_id);
125        let item = self.crate_data.index.get(&id).context("Item not found")?;
126
127        let info = self
128            .item_to_info(&id, item)
129            .context("Failed to convert item to info")?;
130
131        let mut details = DetailedItem {
132            info,
133            signature: self.get_item_signature(item),
134            generics: None,
135            fields: None,
136            variants: None,
137            methods: None,
138            source_location: self.get_item_source_location(item),
139        };
140
141        // Add type-specific information
142        match &item.inner {
143            ItemEnum::Struct(s) => {
144                details.generics = serde_json::to_value(&s.generics).ok();
145                details.fields = Some(self.get_struct_fields(s));
146            }
147            ItemEnum::Enum(e) => {
148                details.generics = serde_json::to_value(&e.generics).ok();
149                details.variants = Some(self.get_enum_variants(e));
150            }
151            ItemEnum::Trait(t) => {
152                details.generics = serde_json::to_value(&t.generics).ok();
153                details.methods = Some(self.get_trait_items(&t.items));
154            }
155            ItemEnum::Impl(i) => {
156                details.generics = serde_json::to_value(&i.generics).ok();
157                details.methods = Some(self.get_impl_items(&i.items));
158            }
159            ItemEnum::Function(f) => {
160                details.generics = serde_json::to_value(&f.generics).ok();
161            }
162            _ => {}
163        }
164
165        Ok(details)
166    }
167
168    /// Get documentation for a specific item
169    pub fn get_item_docs(&self, item_id: u32) -> Result<Option<String>> {
170        let id = Id(item_id);
171        let item = self.crate_data.index.get(&id).context("Item not found")?;
172
173        Ok(item.docs.clone())
174    }
175
176    /// Helper to convert an Item to ItemInfo
177    fn item_to_info(&self, id: &Id, item: &Item) -> Option<ItemInfo> {
178        // Get name from item or from paths
179        let name = if let Some(name) = &item.name {
180            name.clone()
181        } else if let Some(path_summary) = self.crate_data.paths.get(id) {
182            path_summary.path.last()?.clone()
183        } else {
184            return None;
185        };
186
187        let kind = self.get_item_kind_string(&item.inner);
188        let path = self.get_item_path(id);
189        let visibility = self.get_visibility_string(&item.visibility);
190
191        Some(ItemInfo {
192            id: id.0.to_string(),
193            name,
194            kind,
195            path,
196            docs: item.docs.clone(),
197            visibility,
198        })
199    }
200
201    /// Get the kind of an item as a string
202    fn get_item_kind_string(&self, inner: &ItemEnum) -> String {
203        use ItemEnum::*;
204        match inner {
205            Module(_) => "module",
206            Struct(_) => "struct",
207            Enum(_) => "enum",
208            Function(_) => "function",
209            Trait(_) => "trait",
210            Impl(_) => "impl",
211            TypeAlias(_) => "type_alias",
212            Constant { .. } => "constant",
213            Static(_) => "static",
214            Macro(_) => "macro",
215            ExternCrate { .. } => "extern_crate",
216            Use(_) => "use",
217            Union(_) => "union",
218            StructField(_) => "field",
219            Variant(_) => "variant",
220            TraitAlias(_) => "trait_alias",
221            ProcMacro(_) => "proc_macro",
222            Primitive(_) => "primitive",
223            AssocConst { .. } => "assoc_const",
224            AssocType { .. } => "assoc_type",
225            ExternType => "extern_type",
226        }
227        .to_string()
228    }
229
230    /// Get the full path of an item
231    fn get_item_path(&self, id: &Id) -> Vec<String> {
232        if let Some(summary) = self.crate_data.paths.get(id) {
233            summary.path.clone()
234        } else {
235            Vec::new()
236        }
237    }
238
239    /// Get visibility as a string
240    fn get_visibility_string(&self, vis: &rustdoc_types::Visibility) -> String {
241        use rustdoc_types::Visibility::*;
242        match vis {
243            Public => "public".to_string(),
244            Default => "default".to_string(),
245            Crate => "crate".to_string(),
246            Restricted { parent, .. } => format!("restricted({})", parent.0),
247        }
248    }
249
250    /// Get a signature representation for an item
251    fn get_item_signature(&self, item: &Item) -> Option<String> {
252        use ItemEnum::*;
253        match &item.inner {
254            Function(f) => {
255                let name = item.name.as_ref()?;
256                let generics = self.format_generics(&f.generics);
257                let params = self.format_fn_params(&f.sig.inputs);
258                let output = self.format_fn_output(&f.sig.output);
259                Some(format!("fn {name}{generics}{params}{output}"))
260            }
261            _ => None,
262        }
263    }
264
265    /// Format generic parameters
266    fn format_generics(&self, generics: &rustdoc_types::Generics) -> String {
267        // Simplified generic formatting
268        if generics.params.is_empty() {
269            String::new()
270        } else {
271            "<...>".to_string()
272        }
273    }
274
275    /// Format function parameters
276    fn format_fn_params(&self, params: &[(String, rustdoc_types::Type)]) -> String {
277        let param_strs: Vec<String> = params.iter().map(|(name, _)| name.clone()).collect();
278        format!("({})", param_strs.join(", "))
279    }
280
281    /// Format function output
282    fn format_fn_output(&self, output: &Option<rustdoc_types::Type>) -> String {
283        output
284            .as_ref()
285            .map(|_| " -> ...".to_string())
286            .unwrap_or_default()
287    }
288
289    /// Get struct fields as ItemInfo
290    fn get_struct_fields(&self, s: &rustdoc_types::Struct) -> Vec<ItemInfo> {
291        use rustdoc_types::StructKind;
292        match &s.kind {
293            StructKind::Unit => vec![],
294            StructKind::Tuple(fields) => fields
295                .iter()
296                .enumerate()
297                .filter_map(|(i, field_id)| {
298                    if let Some(field_id) = field_id {
299                        let item = self.crate_data.index.get(field_id)?;
300                        let mut info = self.item_to_info(field_id, item)?;
301                        if info.name.is_empty() {
302                            info.name = i.to_string();
303                        }
304                        Some(info)
305                    } else {
306                        Some(ItemInfo {
307                            id: String::new(),
308                            name: format!("(field {i} stripped)"),
309                            kind: "field".to_string(),
310                            path: Vec::new(),
311                            docs: None,
312                            visibility: "private".to_string(),
313                        })
314                    }
315                })
316                .collect(),
317            StructKind::Plain {
318                fields,
319                has_stripped_fields,
320            } => {
321                let mut field_infos: Vec<ItemInfo> = fields
322                    .iter()
323                    .filter_map(|field_id| {
324                        let item = self.crate_data.index.get(field_id)?;
325                        self.item_to_info(field_id, item)
326                    })
327                    .collect();
328
329                if *has_stripped_fields {
330                    field_infos.push(ItemInfo {
331                        id: String::new(),
332                        name: "(some fields stripped)".to_string(),
333                        kind: "note".to_string(),
334                        path: Vec::new(),
335                        docs: None,
336                        visibility: "private".to_string(),
337                    });
338                }
339
340                field_infos
341            }
342        }
343    }
344
345    /// Get enum variants as ItemInfo
346    fn get_enum_variants(&self, e: &rustdoc_types::Enum) -> Vec<ItemInfo> {
347        let mut variant_infos: Vec<ItemInfo> = e
348            .variants
349            .iter()
350            .filter_map(|variant_id| {
351                let item = self.crate_data.index.get(variant_id)?;
352                self.item_to_info(variant_id, item)
353            })
354            .collect();
355
356        if e.has_stripped_variants {
357            variant_infos.push(ItemInfo {
358                id: String::new(),
359                name: "(some variants stripped)".to_string(),
360                kind: "note".to_string(),
361                path: Vec::new(),
362                docs: None,
363                visibility: "private".to_string(),
364            });
365        }
366
367        variant_infos
368    }
369
370    /// Get trait items as ItemInfo
371    fn get_trait_items(&self, items: &[Id]) -> Vec<ItemInfo> {
372        items
373            .iter()
374            .filter_map(|item_id| {
375                let item = self.crate_data.index.get(item_id)?;
376                self.item_to_info(item_id, item)
377            })
378            .collect()
379    }
380
381    /// Get impl items as ItemInfo
382    fn get_impl_items(&self, items: &[Id]) -> Vec<ItemInfo> {
383        items
384            .iter()
385            .filter_map(|item_id| {
386                let item = self.crate_data.index.get(item_id)?;
387                self.item_to_info(item_id, item)
388            })
389            .collect()
390    }
391
392    /// Get source location information for an item
393    fn get_item_source_location(&self, item: &Item) -> Option<SourceLocation> {
394        let span = item.span.as_ref()?;
395        Some(SourceLocation {
396            filename: span.filename.to_string_lossy().to_string(),
397            line_start: span.begin.0,
398            column_start: span.begin.1,
399            line_end: span.end.0,
400            column_end: span.end.1,
401        })
402    }
403
404    /// Get source code for a specific item by ID
405    pub fn get_item_source(
406        &self,
407        item_id: u32,
408        base_path: &std::path::Path,
409        context_lines: usize,
410    ) -> Result<SourceInfo> {
411        let id = Id(item_id);
412        let item = self.crate_data.index.get(&id).context("Item not found")?;
413
414        let span = item.span.as_ref().context("Item has no source span")?;
415        let source_path = base_path.join(&span.filename);
416
417        if !source_path.exists() {
418            anyhow::bail!("Source file not found: {}", source_path.display());
419        }
420
421        let content = std::fs::read_to_string(&source_path)
422            .with_context(|| format!("Failed to read source file: {}", source_path.display()))?;
423
424        let lines: Vec<&str> = content.lines().collect();
425
426        // Calculate line range with context
427        let start_line = span.begin.0.saturating_sub(1).saturating_sub(context_lines);
428        let end_line = std::cmp::min(span.end.0 + context_lines, lines.len());
429
430        // Extract the relevant lines
431        let code_lines: Vec<String> = lines[start_line..end_line]
432            .iter()
433            .map(|line| line.to_string())
434            .collect();
435
436        Ok(SourceInfo {
437            location: SourceLocation {
438                filename: span.filename.to_string_lossy().to_string(),
439                line_start: span.begin.0,
440                column_start: span.begin.1,
441                line_end: span.end.0,
442                column_end: span.end.1,
443            },
444            code: code_lines.join("\n"),
445            context_lines: Some(context_lines),
446        })
447    }
448}
rust_docs_mcp/docs/query.rs

rust_docs_mcp/docs/
query.rs
