feat: references for stages

alesbrelih · Apr 6, 2024 · 6670032 · 6670032
1 parent 32132c4
commit 6670032
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 14 deletions.
diff --git a/src/gitlab_ci_ls_parser/handlers.rs b/src/gitlab_ci_ls_parser/handlers.rs
@@ -294,7 +294,7 @@ impl LSPHandlers {
 
                 locations.append(&mut root);
             }
-            _ => {
+            parser::PositionType::None => {
                 error!("invalid position type for goto def");
                 return None;
             }
@@ -761,9 +761,9 @@ impl LSPHandlers {
             }
         }
 
-        let stages = self
-            .parser
-            .get_all_stages(params.text_document.uri.as_str(), content.as_str());
+        let stages =
+            self.parser
+                .get_all_stages(params.text_document.uri.as_str(), content.as_str(), None);
 
         let all_stages = self.stages.lock().unwrap();
         for stage in stages {
@@ -871,6 +871,15 @@ impl LSPHandlers {
                     }
                 }
             }
+            parser::PositionType::Stage => {
+                let word =
+                    parser_utils::ParserUtils::extract_word(line, position.character as usize);
+
+                for (uri, content) in store.iter() {
+                    let mut stages = self.parser.get_all_stages(uri, content.as_str(), word);
+                    references.append(&mut stages);
+                }
+            }
             _ => {}
         }
 

diff --git a/src/gitlab_ci_ls_parser/parser.rs b/src/gitlab_ci_ls_parser/parser.rs
@@ -21,7 +21,7 @@ pub trait Parser {
         content: &str,
         extend_name: Option<&str>,
     ) -> Vec<GitlabElement>;
-    fn get_all_stages(&self, uri: &str, content: &str) -> Vec<GitlabElement>;
+    fn get_all_stages(&self, uri: &str, content: &str, stage: Option<&str>) -> Vec<GitlabElement>;
     fn get_position_type(&self, content: &str, position: Position) -> PositionType;
     fn get_root_node(&self, uri: &str, content: &str, node_key: &str) -> Option<GitlabElement>;
     fn parse_contents(&self, uri: &Url, content: &str, _follow: bool) -> Option<ParseResults>;
@@ -197,8 +197,8 @@ impl Parser for ParserImpl {
         self.treesitter.get_all_extends(uri, content, extend_name)
     }
 
-    fn get_all_stages(&self, uri: &str, content: &str) -> Vec<GitlabElement> {
-        self.treesitter.get_all_stages(uri, content)
+    fn get_all_stages(&self, uri: &str, content: &str, stage: Option<&str>) -> Vec<GitlabElement> {
+        self.treesitter.get_all_stages(uri, content, stage)
     }
 
     fn get_position_type(&self, content: &str, position: Position) -> PositionType {

diff --git a/src/gitlab_ci_ls_parser/treesitter.rs b/src/gitlab_ci_ls_parser/treesitter.rs
@@ -14,7 +14,7 @@ pub trait Treesitter {
     fn get_all_root_nodes(&self, uri: &str, content: &str) -> Vec<GitlabElement>;
     fn get_root_variables(&self, uri: &str, content: &str) -> Vec<GitlabElement>;
     fn get_stage_definitions(&self, uri: &str, content: &str) -> Vec<GitlabElement>;
-    fn get_all_stages(&self, uri: &str, content: &str) -> Vec<GitlabElement>;
+    fn get_all_stages(&self, uri: &str, content: &str, stage: Option<&str>) -> Vec<GitlabElement>;
     fn get_all_extends(
         &self,
         uri: String,
@@ -245,7 +245,7 @@ impl Treesitter for TreesitterImpl {
         stages
     }
 
-    fn get_all_stages(&self, uri: &str, content: &str) -> Vec<GitlabElement> {
+    fn get_all_stages(&self, uri: &str, content: &str, stage: Option<&str>) -> Vec<GitlabElement> {
         let mut parser = tree_sitter::Parser::new();
         parser
             .set_language(tree_sitter_yaml::language())
@@ -254,7 +254,7 @@ impl Treesitter for TreesitterImpl {
         let tree = parser.parse(content, None).unwrap();
         let root_node = tree.root_node();
 
-        let query = Query::new(language(), &TreesitterQueries::get_all_stages()).unwrap();
+        let query = Query::new(language(), &TreesitterQueries::get_all_stages(stage)).unwrap();
         let mut cursor_qry = QueryCursor::new();
         let matches = cursor_qry.matches(&query, root_node, content.as_bytes());
 
@@ -885,7 +885,7 @@ job_two:
 
         let uri = "file://mocked";
         let treesitter = TreesitterImpl::new();
-        let all_stages = treesitter.get_all_stages(uri, cnt);
+        let all_stages = treesitter.get_all_stages(uri, cnt, None);
 
         assert_eq!(all_stages.len(), 2);
 

diff --git a/src/gitlab_ci_ls_parser/treesitter_queries.rs b/src/gitlab_ci_ls_parser/treesitter_queries.rs
@@ -99,8 +99,14 @@ impl TreesitterQueries {
         "#.to_string()
     }
 
-    pub fn get_all_stages() -> String {
-        r#"
+    pub fn get_all_stages(stage: Option<&str>) -> String {
+        let mut search = String::new();
+        if stage.is_some() {
+            search = format!("(#eq? @value \"{}\")", stage.unwrap());
+        }
+
+        format!(
+            r#"
         (
             block_mapping_pair
                 key: (
@@ -114,9 +120,10 @@ impl TreesitterQueries {
                     )
                 )
             (#eq? @key "stage")
+            {search}
         )
         "#
-        .to_string()
+        )
     }
 
     #[allow(clippy::too_many_lines)]
@@ -148,6 +155,16 @@ impl TreesitterQueries {
                     )
                 (#eq? @keystage "stage")
             )
+            (
+                block_mapping_pair
+                    key: (
+                        flow_node(
+                            plain_scalar(string_scalar) @keystage
+                        )
+                    )
+                    value: (block_node(block_sequence(block_sequence_item)@stage ))
+                (#eq? @keystage "stages")
+            )
         "#;
 
         let search_variables = r#"