implemented file-local find-refs for functions

Author: Strum355

server/main/src/linemap.rs

@@ -0,0 +1,80 @@
+use rust_lsp::lsp_types::Position;
+
+pub struct LineMap {
+    positions: Vec<usize>,
+}
+
+impl LineMap {
+    pub fn new(source: &str) -> Self {
+        let mut positions = vec![0];
+        for (i, char) in source.char_indices() {
+            if char == '\n' {
+                positions.push(i + 1);
+            }
+        }
+
+        LineMap { positions }
+    }
+
+    pub fn offset_for_position(&self, position: Position) -> usize {
+        self.positions[position.line as usize] + (position.character as usize)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use rust_lsp::lsp_types::Position;
+
+    use crate::linemap::LineMap;
+
+    #[test]
+    #[logging_macro::log_scope]
+    fn test_linemap() {
+        struct Test {
+            string: &'static str,
+            pos: Position,
+            offset: usize,
+        }
+
+        let cases = vec![
+            Test {
+                string: "sample\ntext",
+                pos: Position { line: 1, character: 2 },
+                offset: 9,
+            },
+            Test {
+                string: "banana",
+                pos: Position { line: 0, character: 0 },
+                offset: 0,
+            },
+            Test {
+                string: "banana",
+                pos: Position { line: 0, character: 1 },
+                offset: 1,
+            },
+            Test {
+                string: "sample\ntext",
+                pos: Position { line: 1, character: 0 },
+                offset: 7,
+            },
+            Test {
+                string: "sample\n\ttext",
+                pos: Position { line: 1, character: 2 },
+                offset: 9,
+            },
+            Test {
+                string: "sample\r\ntext",
+                pos: Position { line: 1, character: 0 },
+                offset: 8,
+            },
+        ];
+
+        for case in cases {
+            let linemap = LineMap::new(case.string);
+
+            let offset = linemap.offset_for_position(case.pos);
+
+            assert_eq!(offset, case.offset, "{:?}", case.string);
+        }
+    }
+}

server/main/src/main.rs

@@ -47,6 +47,7 @@ mod consts;
 mod dfs;
 mod diagnostics_parser;
 mod graph;
+mod linemap;
 mod lsp_ext;
 mod merge_views;
 mod navigation;
@@ -515,6 +516,7 @@ impl LanguageServerHandling for MinecraftShaderLanguageServer {
 
             let capabilities = ServerCapabilities {
                 definition_provider: Some(OneOf::Left(true)),
+                references_provider: Some(OneOf::Left(true)),
                 document_link_provider: Some(DocumentLinkOptions {
                     resolve_provider: None,
                     work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
@@ -689,31 +691,60 @@ impl LanguageServerHandling for MinecraftShaderLanguageServer {
 
     fn goto_definition(&mut self, params: TextDocumentPositionParams, completable: LSCompletable<Vec<Location>>) {
         logging::slog_with_trace_id(|| {
+            let path = PathBuf::from_url(params.text_document.uri);
+            if !path.starts_with(&self.root) {
+                return;
+            }
             let parser = &mut self.tree_sitter.borrow_mut();
-            let parser_ctx = match navigation::ParserContext::new(parser, &params.text_document.uri) {
+            let parser_ctx = match navigation::ParserContext::new(parser, &path) {
                 Ok(ctx) => ctx,
                 Err(e) => {
                     return completable.complete(Err(MethodError {
                         code: 42069,
-                        message: format!("error building parser context: {}", e.context(params.text_document.uri)),
+                        message: format!("error building parser context: error={}, path={:?}", e, path),
                         data: (),
                     }))
                 }
             };
 
-            match parser_ctx.find_definitions(&params.text_document.uri, params.position) {
-                Ok(locations) => completable.complete(Ok(locations)),
+            match parser_ctx.find_definitions(&path, params.position) {
+                Ok(locations) => completable.complete(Ok(locations.unwrap_or_default())),
                 Err(e) => completable.complete(Err(MethodError {
                     code: 42069,
-                    message: format!("error finding definitions: {}", e.context(params.text_document.uri)),
+                    message: format!("error finding definitions: error={}, path={:?}", e, path),
                     data: (),
                 })),
             }
         });
     }
 
-    fn references(&mut self, _: ReferenceParams, completable: LSCompletable<Vec<Location>>) {
-        completable.complete(Err(Self::error_not_available(())));
+    fn references(&mut self, params: ReferenceParams, completable: LSCompletable<Vec<Location>>) {
+        logging::slog_with_trace_id(|| {
+            let path = PathBuf::from_url(params.text_document_position.text_document.uri);
+            if !path.starts_with(&self.root) {
+                return;
+            }
+            let parser = &mut self.tree_sitter.borrow_mut();
+            let parser_ctx = match navigation::ParserContext::new(parser, &path) {
+                Ok(ctx) => ctx,
+                Err(e) => {
+                    return completable.complete(Err(MethodError {
+                        code: 42069,
+                        message: format!("error building parser context: error={}, path={:?}", e, path),
+                        data: (),
+                    }))
+                }
+            };
+
+            match parser_ctx.find_references(&path, params.text_document_position.position) {
+                Ok(locations) => completable.complete(Ok(locations.unwrap_or_default())),
+                Err(e) => completable.complete(Err(MethodError {
+                    code: 42069,
+                    message: format!("error finding definitions: error={}, path={:?}", e, path),
+                    data: (),
+                })),
+            }
+        });
     }
 
     fn document_highlight(&mut self, _: TextDocumentPositionParams, completable: LSCompletable<Vec<DocumentHighlight>>) {
@@ -743,7 +774,7 @@ impl LanguageServerHandling for MinecraftShaderLanguageServer {
     fn document_link(&mut self, params: DocumentLinkParams, completable: LSCompletable<Vec<DocumentLink>>) {
         logging::slog_with_trace_id(|| {
             // node for current document
-            let curr_doc = params.text_document.uri.to_file_path().unwrap();
+            let curr_doc = PathBuf::from_url(params.text_document.uri);
             let node = match self.graph.borrow_mut().find_node(&curr_doc) {
                 Some(n) => n,
                 None => {

server/main/src/navigation.rs

@@ -1,68 +1,89 @@
-use std::fs::read_to_string;
+use std::{fs::read_to_string, path::Path};
 
 use anyhow::Result;
 use rust_lsp::lsp_types::{Location, Position, Range};
-use slog_scope::{info, debug};
-use tree_sitter::{Node, Parser, Point, Tree, Query, QueryCursor};
+use slog_scope::{debug, info, trace};
+use tree_sitter::{Node, Parser, Point, Query, QueryCursor, Tree};
 use url::Url;
 
-macro_rules! find_function_str {
-    () => { 
+use crate::linemap::LineMap;
+
+macro_rules! find_function_def_str {
+    () => {
         r#"
             (
                 (function_declarator 
                     (identifier) @function)
-                    (#match? @function "{}")
+                (#match? @function "{}")
+            )
+        "#
+    };
+}
+
+macro_rules! find_function_refs_str {
+    () => {
+        r#"
+            (
+                (call_expression 
+                    (identifier) @call)
+                (#match? @call "{}")
             )
         "#
     };
 }
 pub struct ParserContext<'a> {
     source: String,
     tree: Tree,
+    linemap: LineMap,
     parser: &'a mut Parser,
 }
 
 impl<'a> ParserContext<'a> {
-    pub fn new(parser: &'a mut Parser, document_uri: &Url) -> Result<Self> {
-        let source = read_to_string(document_uri.path())?;
+    pub fn new(parser: &'a mut Parser, path: &Path) -> Result<Self> {
+        let source = read_to_string(path)?;
 
         let tree = parser.parse(&source, None).unwrap();
 
-        Ok(ParserContext { source, tree, parser })
+        let linemap = LineMap::new(&source);
+
+        Ok(ParserContext {
+            source,
+            tree,
+            linemap,
+            parser,
+        })
     }
 
-    pub fn find_definitions(&self, document_uri: &Url, point: Position) -> Result<Vec<Location>> {
+    pub fn find_definitions(&self, path: &Path, point: Position) -> Result<Option<Vec<Location>>> {
         let current_node = match self.find_node_at_point(point) {
             Some(node) => node,
-            None => return Ok(vec![]),
+            None => return Ok(None),
         };
 
         let parent = match current_node.parent() {
             Some(parent) => parent,
-            None => return Ok(vec![]),
+            None => return Ok(None),
         };
 
         let query = match (current_node.kind(), parent.kind()) {
             (_, "call_expression") => {
-                format!(find_function_str!(), current_node.utf8_text(self.source.as_bytes())?)
+                format!(find_function_def_str!(), current_node.utf8_text(self.source.as_bytes())?)
             }
-            _ => return Ok(vec![]),
+            _ => return Ok(None),
         };
 
         let ts_query = Query::new(tree_sitter_glsl::language(), query.as_str())?;
-
         let mut query_cursor = QueryCursor::new();
 
         let mut locations = vec![];
 
-        for m in query_cursor.matches(&ts_query, self.tree.root_node(), self.source.as_bytes()) {
+        for m in query_cursor.matches(&ts_query, self.root_node(), self.source.as_bytes()) {
             for capture in m.captures {
                 let start = capture.node.start_position();
                 let end = capture.node.end_position();
 
                 locations.push(Location {
-                    uri: document_uri.clone(),
+                    uri: Url::from_file_path(path).unwrap(),
                     range: Range {
                         start: Position {
                             line: start.row as u32,
@@ -79,32 +100,98 @@ impl<'a> ParserContext<'a> {
 
         info!("finished searching for definitions"; "definitions" => format!("{:?}", locations));
 
-        Ok(locations)
+        Ok(Some(locations))
     }
 
-    pub fn find_references(&self) -> Result<Vec<Location>> {
-        Ok(vec![])
+    pub fn find_references(&self, path: &Path, point: Position) -> Result<Option<Vec<Location>>> {
+        let current_node = match self.find_node_at_point(point) {
+            Some(node) => node,
+            None => return Ok(None),
+        };
+
+        let parent = match current_node.parent() {
+            Some(parent) => parent,
+            None => return Ok(None),
+        };
+
+        let query = match (current_node.kind(), parent.kind()) {
+            (_, "function_declarator") => {
+                format!(find_function_refs_str!(), current_node.utf8_text(self.source.as_bytes())?)
+            }
+            _ => return Ok(None),
+        };
+
+        let ts_query = Query::new(tree_sitter_glsl::language(), query.as_str())?;
+        let mut query_cursor = QueryCursor::new();
+
+        let mut locations = vec![];
+
+        for m in query_cursor.matches(&ts_query, self.root_node(), self.source.as_bytes()) {
+            for capture in m.captures {
+                let start = capture.node.start_position();
+                let end = capture.node.end_position();
+
+                locations.push(Location {
+                    uri: Url::from_file_path(path).unwrap(),
+                    range: Range {
+                        start: Position {
+                            line: start.row as u32,
+                            character: start.column as u32,
+                        },
+                        end: Position {
+                            line: end.row as u32,
+                            character: end.column as u32,
+                        },
+                    },
+                });
+            }
+        }
+
+        Ok(Some(locations))
     }
 
     fn root_node(&self) -> Node {
         self.tree.root_node()
     }
 
-    fn find_node_at_point(&self, point: Position) -> Option<Node> {
-        match self.root_node().named_descendant_for_point_range(
-            Point {
-                row: point.line as usize,
-                column: (point.character - 1) as usize,
-            },
-            Point {
-                row: point.line as usize,
-                column: point.character as usize,
-            },
-        ) {
+    fn find_node_at_point(&self, pos: Position) -> Option<Node> {
+        // if we're at the end of an ident, we need to look _back_ one char instead
+        // for tree-sitter to find the right node.
+        let look_behind = {
+            let offset = self.linemap.offset_for_position(pos);
+            let char_at = self.source.as_bytes()[offset];
+            trace!("looking for non-alpha for point adjustment";
+                "offset" => offset, 
+                "char" => char_at as char,
+                "point" => format!("{:?}", pos),
+                "look_behind" => !char_at.is_ascii_alphabetic());
+            !char_at.is_ascii_alphabetic()
+        };
+
+        let mut start = Point {
+            row: pos.line as usize,
+            column: pos.character as usize,
+        };
+        let mut end = Point {
+            row: pos.line as usize,
+            column: pos.character as usize,
+        };
+
+        if look_behind {
+            start.column -= 1;
+        } else {
+            end.column += 1;
+        }
+
+        match self.root_node().named_descendant_for_point_range(start, end) {
             Some(node) => {
-                debug!("found a node"; "node" => format!("{:?}", node));
+                debug!("found a node"; 
+                    "node" => format!("{:?}", node),
+                    "text" => node.utf8_text(self.source.as_bytes()).unwrap(),
+                    "start" => format!("{}", start),
+                    "end" => format!("{}", end));
                 Some(node)
-            },
+            }
             None => None,
         }
     }