Classify function assignment names as function semantic tokens

R/semantic.R

@@ -69,6 +69,67 @@ empty_semantic_data <- function() {
     )
 }
 
+#' Parse-data ids of symbols assigned `function()` or `\()`
+#'
+#' Same cases as `scope_completion_functs_xpath` in `completion.R`.
+#' @noRd
+function_assignment_symbol_ids <- function(data) {
+    fun_expr_ids <- unique(data$parent[data$token %in% c("FUNCTION", "'\\\\'")])
+    if (!length(fun_expr_ids)) {
+        return(integer())
+    }
+
+    assign_rows <- which(
+        data$token %in% c("LEFT_ASSIGN", "EQ_ASSIGN", "RIGHT_ASSIGN")
+    )
+    lhs_ids <- integer()
+    for (row in assign_rows) {
+        children <- data[data$parent == data$parent[[row]], , drop = FALSE]
+        assign_pos <- match(data$id[[row]], children$id)
+        if (is.na(assign_pos)) {
+            next
+        }
+        before_idx <- seq_len(assign_pos - 1L)
+        after_idx <- if (assign_pos < nrow(children)) {
+            seq.int(assign_pos + 1L, nrow(children))
+        } else {
+            integer()
+        }
+        if (data$token[[row]] == "RIGHT_ASSIGN") {
+            rhs <- children[before_idx, , drop = FALSE]
+            lhs <- children[after_idx, , drop = FALSE]
+        } else {
+            lhs <- children[before_idx, , drop = FALSE]
+            rhs <- children[after_idx, , drop = FALSE]
+        }
+        if (!any(rhs$id %in% fun_expr_ids)) {
+            next
+        }
+        for (expr_id in lhs$id[lhs$token %in% c("expr", "expr_or_assign_or_help")]) {
+            child_rows <- which(data$parent == expr_id)
+            terminal_rows <- child_rows[data$terminal[child_rows]]
+            if (length(terminal_rows) == 1L &&
+                    data$token[[terminal_rows]] == "SYMBOL") {
+                lhs_ids <- c(lhs_ids, data$id[[terminal_rows]])
+            }
+        }
+    }
+    unique(lhs_ids)
+}
+
+#' Whether an XML SYMBOL is assigned `function()` or `\()`
+#' @noRd
+xml_symbol_is_function_assignment <- function(node) {
+    !inherits(
+        xml_find_first(node, paste(
+            "./parent::expr[count(*)=1]/following-sibling::*[self::LEFT_ASSIGN or self::EQ_ASSIGN][following-sibling::expr/*[self::FUNCTION or self::OP-LAMBDA]]",
+            "./parent::expr[count(*)=1]/preceding-sibling::RIGHT_ASSIGN[preceding-sibling::expr/*[self::FUNCTION or self::OP-LAMBDA]]",
+            sep = "|"
+        )),
+        "xml_missing"
+    )
+}
+
 #' Build a compact semantic token index from R parse data
 #'
 #' This runs in the parse worker. Keeping ordinary integer vectors here avoids
@@ -121,6 +182,16 @@ semantic_parse_data <- function(data, content) {
     )
     modifiers <- as.integer(modifiers)
 
+    fun_lhs_ids <- function_assignment_symbol_ids(data)
+    if (length(fun_lhs_ids)) {
+        fun_lhs <- data$id[single_rows] %in% fun_lhs_ids
+        types[fun_lhs] <- SemanticTokenTypes[["function"]]
+        modifiers[fun_lhs] <- bitwOr(
+            modifiers[fun_lhs],
+            bitwShiftL(1L, SemanticTokenModifiers$declaration)
+        )
+    }
+
     non_ascii_lines <- nchar(content, type = "bytes") !=
         nchar(content, type = "chars")
     convert <- which(non_ascii_lines[data$line1[single_rows]])
@@ -395,7 +466,10 @@ extract_semantic_tokens <- function(uri, workspace, document, range = NULL) {
         modifiers <- 0L  # Start with no modifiers
 
         # Determine modifiers based on context
-        if (token_name == "SYMBOL_FUNCTION_CALL") {
+        if (token_name == "SYMBOL" && xml_symbol_is_function_assignment(token_node)) {
+            token_type <- SemanticTokenTypes[["function"]]
+            modifiers <- bitwOr(modifiers, 2^SemanticTokenModifiers$declaration)
+        } else if (token_name == "SYMBOL_FUNCTION_CALL") {
             # Function calls might be declared elsewhere
         } else if (token_name == "SYMBOL_FORMALS") {
             # Parameters are declarations

tests/testthat/test-semantic-tokens.R

@@ -180,6 +180,59 @@ test_that("Semantic parse data handles UTF-16 and multiline tokens", {
     )
 })
 
+test_that("Function assignment names are function declarations", {
+    content <- c(
+        "fn <- function(x) x",
+        "gn = function(x) x",
+        "hn <- \\(x) x",
+        "value <- 1",
+        "nested <- foo(function(x) x)",
+        "fn(value)"
+    )
+    parsed <- parse(text = content, keep.source = TRUE)
+    semantic <- semantic_parse_data(
+        utils::getParseData(parsed, includeText = TRUE),
+        content
+    )
+
+    token_at <- function(line, name) {
+        which(
+            semantic$lines == line &
+                semantic$cols == 0L &
+                semantic$lengths == nchar(name)
+        )
+    }
+
+    expect_equal(
+        semantic$types[token_at(0L, "fn")],
+        SemanticTokenTypes[["function"]]
+    )
+    expect_equal(
+        semantic$modifiers[token_at(0L, "fn")],
+        bitwShiftL(1L, SemanticTokenModifiers$declaration)
+    )
+    expect_equal(
+        semantic$types[token_at(1L, "gn")],
+        SemanticTokenTypes[["function"]]
+    )
+    expect_equal(
+        semantic$types[token_at(2L, "hn")],
+        SemanticTokenTypes[["function"]]
+    )
+    expect_equal(
+        semantic$types[token_at(3L, "value")],
+        SemanticTokenTypes$variable
+    )
+    expect_equal(
+        semantic$types[token_at(4L, "nested")],
+        SemanticTokenTypes$variable
+    )
+    expect_equal(
+        semantic$types[token_at(5L, "fn")],
+        SemanticTokenTypes[["function"]]
+    )
+})
+
 test_that("Semantic ranges select overlapping tokens and re-encode them", {
     fixture <- provider_fixture(c("alpha <- 1", "beta <- alpha", "gamma <- 3"))
     data <- fixture$document$parse_data$semantic_data
@@ -319,6 +372,12 @@ test_that("Legacy XML semantic extraction handles ranges and declarations", {
 
     tokens <- extract_semantic_tokens(uri, workspace, document)
     expect_gt(length(tokens), 0L)
+    fn_declaration <- Filter(function(token) {
+        token$line == 0L && token$col == 0L && token$length == nchar("fn")
+    }, tokens)
+    expect_length(fn_declaration, 1L)
+    expect_equal(fn_declaration[[1L]]$tokenType, SemanticTokenTypes[["function"]])
+    expect_true(fn_declaration[[1L]]$tokenModifiers != 0L)
     parameter <- Filter(function(token) {
         token$tokenType == SemanticTokenTypes$parameter
     }, tokens)