refactor(service): merge logic of finding block references

g-plane · Dec 23, 2024 · 80d9e67 · 80d9e67
1 parent 45ad65d
commit 80d9e67
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 96 deletions.
diff --git a/crates/service/src/binder.rs b/crates/service/src/binder.rs
@@ -595,6 +595,31 @@ impl SymbolTable {
             .iter()
             .filter(move |symbol| symbol.kind == kind && symbol.region == key)
     }
+
+    pub fn find_block_references<'a>(
+        &'a self,
+        def_key: &'a SymbolItemKey,
+        with_decl: bool,
+    ) -> impl Iterator<Item = &'a SymbolItem> {
+        if with_decl {
+            self.symbols.iter().find(|symbol| symbol.key == *def_key)
+        } else {
+            None
+        }
+        .into_iter()
+        .chain(
+            self.blocks
+                .iter()
+                .filter(|block| {
+                    block.def_key == *def_key && block.ref_idx.is_defined_by(&block.def_idx)
+                })
+                .filter_map(|block| {
+                    self.symbols
+                        .iter()
+                        .find(|symbol| symbol.key == block.ref_key)
+                }),
+        )
+    }
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]

diff --git a/crates/service/src/features/document_highlight.rs b/crates/service/src/features/document_highlight.rs
@@ -1,5 +1,5 @@
 use crate::{
-    binder::{SymbolItem, SymbolItemKey, SymbolItemKind, SymbolTable, SymbolTablesCtx},
+    binder::{SymbolItem, SymbolItemKind, SymbolTablesCtx},
     files::FilesCtx,
     helpers, LanguageService,
 };
@@ -164,20 +164,24 @@ impl LanguageService {
                                 .collect(),
                         )
                     }
-                    SymbolItemKind::BlockDef => Some(create_block_highlights(
-                        &current_symbol.key,
-                        &symbol_table,
-                        &line_index,
-                        &root,
-                    )),
+                    SymbolItemKind::BlockDef => Some(
+                        symbol_table
+                            .find_block_references(&current_symbol.key, true)
+                            .filter_map(|symbol| {
+                                create_symbol_highlight(symbol, &root, &line_index)
+                            })
+                            .collect(),
+                    ),
                     SymbolItemKind::BlockRef => {
                         let def_key = symbol_table.find_block_def(&key)?;
-                        Some(create_block_highlights(
-                            def_key,
-                            &symbol_table,
-                            &line_index,
-                            &root,
-                        ))
+                        Some(
+                            symbol_table
+                                .find_block_references(def_key, true)
+                                .filter_map(|symbol| {
+                                    create_symbol_highlight(symbol, &root, &line_index)
+                                })
+                                .collect(),
+                        )
                     }
                 }
             }
@@ -209,39 +213,6 @@ fn create_symbol_highlight(
         })
 }
 
-fn create_block_highlights(
-    def_key: &SymbolItemKey,
-    symbol_table: &SymbolTable,
-    line_index: &LineIndex,
-    root: &SyntaxNode,
-) -> Vec<DocumentHighlight> {
-    let mut highlights = Vec::with_capacity(1);
-    if let Some(highlight) = symbol_table
-        .symbols
-        .iter()
-        .find(|symbol| symbol.key == *def_key)
-        .and_then(|def_symbol| create_symbol_highlight(def_symbol, root, line_index))
-    {
-        highlights.push(highlight);
-    }
-    highlights.extend(
-        symbol_table
-            .blocks
-            .iter()
-            .filter(|block| {
-                block.def_key == *def_key && block.ref_idx.is_defined_by(&block.def_idx)
-            })
-            .map(|block| DocumentHighlight {
-                range: helpers::rowan_range_to_lsp_range(
-                    line_index,
-                    block.ref_key.ptr.text_range(),
-                ),
-                kind: Some(DocumentHighlightKind::READ),
-            }),
-    );
-    highlights
-}
-
 fn get_highlight_kind_of_symbol(
     symbol: &SymbolItem,
     root: &SyntaxNode,

diff --git a/crates/service/src/features/references.rs b/crates/service/src/features/references.rs
@@ -1,6 +1,6 @@
 use super::find_meaningful_token;
 use crate::{
-    binder::{SymbolItem, SymbolItemKey, SymbolItemKind, SymbolTable, SymbolTablesCtx},
+    binder::{SymbolItem, SymbolItemKind, SymbolTablesCtx},
     files::FilesCtx,
     helpers, LanguageService,
 };
@@ -141,62 +141,27 @@ impl LanguageService {
                         .collect(),
                 )
             }
-            SymbolItemKind::BlockDef => Some(get_block_refs(
-                &current_symbol.key,
-                &params,
-                &symbol_table,
-                &line_index,
-                &root,
-            )),
+            SymbolItemKind::BlockDef => Some(
+                symbol_table
+                    .find_block_references(&current_symbol.key, params.context.include_declaration)
+                    .map(|symbol| create_location_by_symbol(&params, &line_index, symbol, &root))
+                    .collect(),
+            ),
             SymbolItemKind::BlockRef => {
                 let def_key = symbol_table.find_block_def(&key)?;
-                Some(get_block_refs(
-                    def_key,
-                    &params,
-                    &symbol_table,
-                    &line_index,
-                    &root,
-                ))
+                Some(
+                    symbol_table
+                        .find_block_references(def_key, params.context.include_declaration)
+                        .map(|symbol| {
+                            create_location_by_symbol(&params, &line_index, symbol, &root)
+                        })
+                        .collect(),
+                )
             }
         }
     }
 }
 
-fn get_block_refs(
-    def_key: &SymbolItemKey,
-    params: &ReferenceParams,
-    symbol_table: &SymbolTable,
-    line_index: &LineIndex,
-    root: &SyntaxNode,
-) -> Vec<Location> {
-    let mut locations = Vec::with_capacity(1);
-    if params.context.include_declaration {
-        if let Some(symbol) = symbol_table
-            .symbols
-            .iter()
-            .find(|symbol| symbol.key == *def_key)
-        {
-            locations.push(create_location_by_symbol(params, line_index, symbol, root));
-        }
-    }
-    locations.extend(
-        symbol_table
-            .blocks
-            .iter()
-            .filter(|block| {
-                block.def_key == *def_key && block.ref_idx.is_defined_by(&block.def_idx)
-            })
-            .filter_map(|block| {
-                symbol_table
-                    .symbols
-                    .iter()
-                    .find(|symbol| symbol.key == block.ref_key)
-            })
-            .map(|symbol| create_location_by_symbol(params, line_index, symbol, root)),
-    );
-    locations
-}
-
 fn create_location_by_symbol(
     params: &ReferenceParams,
     line_index: &LineIndex,