Skip to main content

oak_rust/lsp/
mod.rs

1#![doc = include_str!("readme.md")]
2#[cfg(feature = "oak-highlight")]
3pub mod highlighter;
4
5#[cfg(feature = "lsp")]
6use {
7    futures::{Future, FutureExt},
8    oak_hover::{Hover, HoverProvider},
9    oak_lsp::service::LanguageService,
10    oak_vfs::Vfs,
11};
12#[cfg(feature = "oak-pretty-print")]
13pub mod formatter;
14use crate::{RustLanguage, parser::RustElementType};
15use core::range::Range;
16use dashmap::DashMap;
17use oak_core::{
18    GreenNode, Source,
19    language::{ElementType, TokenType},
20    parser::{ParseCache, Parser, session::ParseSession},
21    tree::RedNode,
22};
23/// Hover provider implementation for Rust.
24#[cfg(feature = "lsp")]
25pub struct RustHoverProvider;
26#[cfg(feature = "lsp")]
27impl HoverProvider<RustLanguage> for RustHoverProvider {
28    fn hover(&self, node: &RedNode<RustLanguage>, _range: Range<usize>) -> Option<Hover> {
29        let kind = node.green.kind;
30        // Provide context-aware hover information
31        let contents = match kind {
32            RustElementType::Function => "### Rust Function\nDefines a callable block of code.",
33            RustElementType::StructItem => "### Rust Struct\nDefines a custom data type.",
34            RustElementType::ModuleItem => "### Rust Module\nOrganizes code into namespaces.",
35            RustElementType::Trait => "### Rust Trait\nDefines a shared behavior.",
36            _ => return None,
37        };
38        Some(Hover { contents: contents.to_string(), range: Some(node.span()) })
39    }
40}
41/// Language service implementation for Rust.
42#[cfg(feature = "lsp")]
43pub struct RustLanguageService<V: Vfs> {
44    vfs: V,
45    workspace: oak_lsp::workspace::WorkspaceManager,
46    hover_provider: RustHoverProvider,
47    sessions: DashMap<String, Box<ParseSession<RustLanguage>>>,
48}
49impl<V: Vfs> RustLanguageService<V> {
50    /// Creates a new `RustLanguageService`.
51    pub fn new(vfs: V) -> Self {
52        Self { vfs, workspace: oak_lsp::workspace::WorkspaceManager::default(), hover_provider: RustHoverProvider, sessions: DashMap::new() }
53    }
54    /// Collects definitions of a name starting from the given node.
55    ///
56    /// This method performs a recursive search through the syntax tree to find nodes
57    /// that represent definitions (e.g., functions, structs, traits) and checks if
58    /// their identifier matches the provided name.
59    ///
60    /// # Arguments
61    ///
62    /// * `node` - The root node to start the search from.
63    /// * `name` - The name of the symbol to look for.
64    /// * `source` - The source text of the file.
65    /// * `uri` - The URI of the file being searched.
66    /// * `definitions` - A mutable vector to store the found locations.
67    fn collect_definitions<S: Source + ?Sized>(&self, node: &RedNode<RustLanguage>, name: &str, source: &S, uri: &str, definitions: &mut Vec<oak_lsp::LocationRange>) {
68        use oak_core::{
69            language::{ElementRole, UniversalElementRole},
70            tree::RedTree,
71        };
72        let role = node.green.kind.role();
73        if role.universal() == UniversalElementRole::Definition {
74            for child in node.children() {
75                if let RedTree::Leaf(leaf) = child {
76                    if leaf.kind.is_universal(oak_core::language::UniversalTokenRole::Name) {
77                        let leaf_name = source.get_text_in(leaf.span.clone());
78                        if leaf_name.as_ref() == name {
79                            definitions.push(oak_lsp::LocationRange { uri: uri.to_string().into(), range: leaf.span });
80                            return;
81                        }
82                    }
83                }
84            }
85        }
86        for child in node.children() {
87            if let RedTree::Node(child_node) = child {
88                self.collect_definitions(&child_node, name, source, uri, definitions);
89            }
90        }
91    }
92}
93impl<V: Vfs + Send + Sync + 'static + oak_vfs::WritableVfs> LanguageService for RustLanguageService<V> {
94    type Lang = RustLanguage;
95    type Vfs = V;
96    fn vfs(&self) -> &Self::Vfs {
97        &self.vfs
98    }
99    fn workspace(&self) -> &oak_lsp::workspace::WorkspaceManager {
100        &self.workspace
101    }
102    fn get_root(&self, uri: &str) -> impl Future<Output = Option<RedNode<'_, RustLanguage>>> + Send + '_ {
103        let uri = uri.to_string();
104        async move {
105            let source = self.vfs().get_source(&uri)?;
106            let mut session = self.sessions.entry(uri.clone()).or_insert_with(|| Box::new(ParseSession::<RustLanguage>::default()));
107            let language = RustLanguage::default();
108            let parser = crate::parser::RustParser::new(&language);
109            let session_guard = session.value_mut();
110            let session_ptr: *mut ParseSession<RustLanguage> = session_guard.as_mut();
111            let output = parser.parse(&source, &[], session_guard.as_mut());
112            // Commit the generation so it can be retrieved via last_root
113            let root_green = output.result.ok()?;
114            unsafe { (*session_ptr).commit_generation(root_green) }
115            // Safety: The root is stored in the ParseSession, which is inside a Box in the DashMap.
116            // Since the DashMap is owned by self, and the Box provides stable addressing,
117            // we can safely extend the lifetime to the lifetime of self.
118            let root_ptr = unsafe {
119                let ptr = (*session_ptr).last_root()?;
120                std::mem::transmute::<&GreenNode<RustLanguage>, &GreenNode<RustLanguage>>(ptr)
121            };
122            Some(RedNode::new(root_ptr, 0))
123        }
124    }
125    fn definition<'a>(&'a self, uri: &'a str, range: Range<usize>) -> impl Future<Output = Vec<oak_lsp::LocationRange>> + Send + 'a {
126        let uri = uri.to_string();
127        async move {
128            let root = self.get_root(&uri).await?;
129            let source = self.vfs().get_source(&uri)?;
130            let leaf = root.leaf_at_offset(range.start)?;
131            if !leaf.kind.is_universal(oak_core::language::UniversalTokenRole::Name) {
132                return None;
133            }
134            let name = source.get_text_in(leaf.span.clone());
135            let name = name.to_string();
136            // Search for definitions in all files in the workspace
137            let mut all_definitions = Vec::new();
138            let files = self.list_all_files(&uri).await;
139            for file_uri in files {
140                if let Some(file_root) = self.get_root(&file_uri).await {
141                    if let Some(file_source) = self.vfs().get_source(&file_uri) {
142                        self.collect_definitions(&file_root, &name, &file_source, &file_uri, &mut all_definitions);
143                    }
144                }
145            }
146            Some(all_definitions)
147        }
148        .then(|opt| async { opt.unwrap_or_default() })
149    }
150    fn references<'a>(&'a self, uri: &'a str, range: Range<usize>) -> impl Future<Output = Vec<oak_lsp::LocationRange>> + Send + 'a {
151        let uri = uri.to_string();
152        async move {
153            let root = self.get_root(&uri).await?;
154            let source = self.vfs().get_source(&uri)?;
155            let leaf = root.leaf_at_offset(range.start)?;
156            if !leaf.kind.is_universal(oak_core::language::UniversalTokenRole::Name) {
157                return None;
158            }
159            let name = source.get_text_in(leaf.span.clone());
160            let name = name.to_string();
161            // Search in all files in the workspace
162            // Note: In a real LSP, we would use an index for performance
163            let mut all_refs = Vec::new();
164            // Get all workspace folders or just list files from root
165            // For this example, we'll just search the current file and any other files we can find
166            let files = self.list_all_files(&uri).await; // This is a bit hacky as it uses current file as root
167            for file_uri in files {
168                if let Some(file_root) = self.get_root(&file_uri).await {
169                    if let Some(file_source) = self.vfs().get_source(&file_uri) {
170                        let source_ref: &dyn oak_core::Source = &file_source;
171                        // For SimpleReferenceFinder::find, we need the text source content as a string slice
172                        // But SimpleReferenceFinder might expect different arguments based on the error
173                        // "expected `&str`, found `&<V as Vfs>::Source`"
174                        // This implies SimpleReferenceFinder::find expects `&str` for source.
175                        // We need to get the full text from the source.
176                        let source_len = source_ref.length();
177                        let full_text = source_ref.get_text_in(oak_core::Range { start: 0, end: source_len });
178                        let full_text_str = full_text.as_ref();
179                        let refs = oak_navigation::SimpleReferenceFinder::find(&file_root, &name, full_text_str, file_uri.clone());
180                        all_refs.extend(refs.into_iter().map(|l| oak_lsp::LocationRange { uri: l.uri, range: l.range }));
181                    }
182                }
183            }
184            Some(all_refs)
185        }
186        .then(|opt| async { opt.unwrap_or_default() })
187    }
188    fn rename<'a>(&'a self, uri: &'a str, range: Range<usize>, new_name: String) -> impl Future<Output = Option<oak_lsp::WorkspaceEdit>> + Send + 'a {
189        let uri = uri.to_string();
190        async move {
191            let refs = self.references(&uri, range).await;
192            if refs.is_empty() {
193                return None;
194            }
195            let mut changes = std::collections::HashMap::new();
196            for r in refs {
197                changes.entry(r.uri.to_string()).or_insert_with(Vec::new).push(oak_lsp::TextEdit { range: r.range, new_text: new_name.clone() });
198            }
199            Some(oak_lsp::WorkspaceEdit { changes })
200        }
201    }
202    fn hover(&self, uri: &str, range: Range<usize>) -> impl Future<Output = Option<oak_lsp::Hover>> + Send + '_ {
203        let uri = uri.to_string();
204        async move {
205            self.with_root(&uri, |root| {
206                // In a real implementation, you would find the specific node at offset
207                // For this example, we just check the root or simple children
208                self.hover_provider.hover(&root, range).map(|h| oak_lsp::Hover { contents: h.contents, range: h.range })
209            })
210            .await
211            .flatten()
212        }
213    }
214}