ast: Preserve @nospecialize/@specialize in remove_macrocalls (#347)

The `remove_macrocalls` function converts macrocall nodes to block nodes to
enable LSP features to work with local variables inside macrocalls. However,
this transformation caused `@nospecialize` and `@specialize` macros in
function argument lists to become invalid block expressions, which prevented
JuliaLowering from generating correct lowered trees.

This commit adds special handling for these macros: instead of transforming
them, we keep them intact. JuliaLowering.jl provides new macro style
definitions for `@nospecialize` and `@specialize`, so they don't need to be
removed in the first place.

This fix enables document highlight and rename to work correctly for
function parameters annotated with these macros.

Author: aviatesk

CHANGELOG.md

@@ -36,6 +36,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
   opened simultaneously. Added global lock to `activate_do` to serialize
   environment switching operations. This fixes spurious "Failed to identify
   package environment" warnings.
+- Fixed document highlight and rename not working for function parameters
+  annotated with `@nospecialize` or `@specialize`.
 
 ### Internal
 

src/utils/ast.jl

@@ -31,8 +31,30 @@ function without_kinds(st::JL.SyntaxTree, kinds::Tuple{Vararg{JS.Kind}})
         _without_kinds(st, kinds)[1])::JL.SyntaxTree
 end
 
+function is_nospecialize_or_specialize_macrocall(st::JL.SyntaxTree)
+    JS.kind(st) === JS.K"macrocall" || return false
+    JS.numchildren(st) >= 1 || return false
+    macro_name_node = st[1]
+    JS.kind(macro_name_node) === JS.K"macro_name" || return false
+    JS.numchildren(macro_name_node) >= 1 || return false
+    macro_name = macro_name_node[1]
+    JS.kind(macro_name) === JS.K"Identifier" || return false
+    hasproperty(macro_name, :name_val) || return false
+    return macro_name.name_val == "nospecialize" || macro_name.name_val == "specialize"
+end
+
 function _remove_macrocalls(st::JL.SyntaxTree)
     if JS.kind(st) === JS.K"macrocall"
+        if is_nospecialize_or_specialize_macrocall(st)
+            # Special case `@nospecialize`/`@specialize`:
+            # These macros are sometimes used in method definition argument lists, but
+            # if we apply the `_remove_macrocalls` transformation directly, it would
+            # result in a `:block` expression being inserted into the argument list,
+            # preventing generation of a correct lowered tree.
+            # Furthermore, JuliaLowering.jl provides new macro style definitions for
+            # these macros, so there's no need to remove them in the first place.
+            return st, false
+        end
         new_children = JL.SyntaxList(JL.syntax_graph(st))
         for i = 2:JS.numchildren(st)
             push!(new_children, _remove_macrocalls(st[i])[1])

test/test_document_highlight.jl

@@ -119,6 +119,33 @@ end
             end
         end
 
+        @testset "highlight with @nospecialize" begin
+            code = """
+            function func(@nospecialize(│xxx│), yyy)
+                zzz = │xxx│, yyy
+                zzz, yyy
+            end
+            """
+            clean_code, positions = JETLS.get_text_and_positions(code)
+            @test length(positions) == 4
+            fi = JETLS.FileInfo(#=version=#0, clean_code, @__FILE__)
+            @test issorted(positions; by = x -> JETLS.xy_to_offset(fi, x))
+            for pos in positions
+                highlights = JETLS.document_highlights(fi, pos)
+                @test length(highlights) == 2
+                @test any(highlights) do highlight
+                    highlight.range.start == positions[1] &&
+                    highlight.range.var"end" == positions[2] &&
+                    highlight.kind == DocumentHighlightKind.Write
+                end
+                @test any(highlights) do highlight
+                    highlight.range.start == positions[3] &&
+                    highlight.range.var"end" == positions[4] &&
+                    highlight.kind == DocumentHighlightKind.Read
+                end
+            end
+        end
+
         let code = """
             let │xxx│, │yyy│ = :yyy
                 │xxx│ = :xxx