Fix member completion for assigned and collected Polars queries

NEWS.md

@@ -4,6 +4,8 @@
   list/environment factories, declarative R6 classes and package method
   registries. Package metadata is prepared in a background worker; completion
   never calls document expressions, constructors, methods or active getters.
+- Preserve member completion for assigned and reassigned Polars query chains,
+  including DataFrame results after `collect()`, within the inference budget.
 - Use the same static member inference for signature help and hover, including
   chained methods and named argument documentation. Member requests wait for the
   current parse after edits, fixing stale completions on the first `$` trigger.

R/member-completion.R

@@ -260,6 +260,9 @@ member_context_index <- function(workspace, uri, document, at, parsed = NULL) {
         seen <- c(seen, pending)
         for (name in pending) {
             history <- document$parse_data$member_data$bindings[[name]]
+            # A dangling $ can make the following assignment parse as a
+            # member write. Later bindings must not hide the receiver's package.
+            history <- Filter(function(item) member_before(item$end, at), history)
             for (item in utils::tail(history, 1L)) {
                 member_walk(item$expr, function(node) {
                     if (is.symbol(node)) referenced <<- c(referenced, as.character(node))

R/member-extraction.R

@@ -557,6 +557,53 @@ member_package_index <- function(input) {
     }
     index$roots <- index$package_roots[intersect(names(index$package_roots), index$exports)]
     for (name in names(index$roots)) index$namespace_roots[paste(input$package, name, sep = "::")] <- list(index$roots[[name]])
+    member_prepare_returns(index)
     index$cache <- new.env(parent = emptyenv())
     index
 }
+
+# Binding many defaults can dominate a request even when the result class does
+# not depend on any argument. Prepare those guarantees in the package worker.
+member_prepare_returns <- function(index) {
+    index$method_results <- list()
+    backed <- names(index$properties)[vapply(index$properties, function(fields) {
+        any(vapply(fields, function(value) {
+            length(value$type) == 1L &&
+                any(index$members[[value$type]] %in% index$native_factories, na.rm = TRUE)
+        }, logical(1L)))
+    }, logical(1L))]
+    backed <- intersect(backed, names(index$classes))
+    if (!length(backed)) return(invisible(NULL))
+    prepare <- function(key, type = NULL) {
+        fn <- member_definition(index$definitions, key)
+        # This is a preparation-cost heuristic, not a return-type rule. Small
+        # functions keep using ordinary argument-sensitive request inference.
+        if (!member_head(fn, "function") || length(fn[[2L]]) < 10L) return()
+        env <- index$package_roots
+        for (name in names(fn[[2L]])) env[name] <- list(member_value())
+        if (!is.null(type)) {
+            receiver <- member_lookup(index$method_receivers, paste(type, key, sep = "|"))
+            if (is.null(receiver)) receiver <- "self"
+            # Use only class guarantees, never a representative instance's fields.
+            env[receiver] <- list(member_value(type = type))
+        }
+        budget <- new.env(parent = emptyenv())
+        budget$remaining <- 10000L
+        budget$exhausted <- budget$transient <- FALSE
+        result <- member_infer(fn[[3L]], index, env, budget = budget,
+            context = list(key = key, formals = names(fn[[2L]]), actuals = list()))
+        if (!budget$exhausted && length(result$type) == 1L && result$type %in% backed) {
+            id <- if (is.null(type)) key else paste(type, key, sep = "|")
+            # Native-backed declarative classes expose their members in the
+            # index. Argument-dependent fields must not enter a universal summary.
+            index$method_results[id] <- list(member_value(type = result$type, classes = result$classes))
+        }
+    }
+    for (key in names(index$definitions)) prepare(key)
+    for (type in intersect(names(index$members), names(index$classes))) {
+        if (any(index$members[[type]] %in% index$native_factories, na.rm = TRUE)) next
+        for (key in unique(unname(index$members[[type]]))) if (!is.na(key)) prepare(key, type)
+    }
+    index$return_definitions <- digest::digest(index$definitions, algo = "xxhash64")
+    invisible(NULL)
+}

R/member-inference.R

@@ -60,9 +60,13 @@ member_strings <- function(x) {
 }
 
 member_lookup <- function(x, key) {
-    if (is.null(x) || is.null(key) || length(key) != 1L || is.na(key) || !key %in% names(x)) {
+    if (is.null(x) || !is.character(key) || length(key) != 1L || is.na(key) || !nzchar(key)) {
         return(NULL)
     }
+    # Lists and environments return NULL for absent names. Avoid scanning
+    # large namespace maps twice for each syntax node.
+    if (is.list(x) || is.environment(x)) return(x[[key]])
+    if (!key %in% names(x)) return(NULL)
     x[[key]]
 }
 
@@ -255,9 +259,18 @@ member_infer <- function(
     # start timing after R's first-call JIT compilation, before traversing ASTs.
     if (!is.null(budget$time_limit) && is.null(budget$deadline)) {
         budget$deadline <- proc.time()[[3L]] + budget$time_limit
-    }
-    if (depth > 64L || budget$remaining < 0L ||
-        (!is.null(budget$deadline) && proc.time()[[3L]] > budget$deadline)) {
+        budget$next_time_check <- budget$remaining - 128L
+    }
+    # Reading the clock for every AST node can consume most of the request's
+    # time limit itself. Check periodically; node and depth bounds still apply
+    # on every visit, and an externally supplied deadline is checked first.
+    timed_out <- FALSE
+    if (!is.null(budget$deadline) && (is.null(budget$next_time_check) ||
+                budget$remaining <= budget$next_time_check)) {
+        budget$next_time_check <- budget$remaining - 128L
+        timed_out <- proc.time()[[3L]] > budget$deadline
+    }
+    if (isTRUE(budget$exhausted) || depth > 64L || budget$remaining < 0L || timed_out) {
         budget$exhausted <- TRUE
         return(member_value(reason = "budget"))
     }
@@ -273,8 +286,8 @@ member_infer <- function(
         )))
         !name %in% names(bindings) &&
             !shadowed &&
-            ((!name %in% names(index$definitions) && name %in% member_base_intrinsics) ||
-                (name %in% index$intrinsics && (is.null(index$document_bindings) || !is.null(context$key))))
+            ((is.null(member_lookup(index$definitions, name)) && name %in% member_base_intrinsics) ||
+                    (name %in% index$intrinsics && (is.null(index$document_bindings) || !is.null(context$key))))
     }
     package_env <- if (!is.null(index$package_roots)) index$package_roots else index$roots
     join_env <- function(a, b) {
@@ -404,7 +417,7 @@ member_infer <- function(
         state <- flow(body, env)
         if (state$falls) member_join(state$returns, state$value) else state$returns
     }
-    bind_arguments <- function(fn, actuals, env) {
+    bind_arguments <- function(fn, actuals, env, defaults = TRUE) {
         formals <- as.list(fn[[2L]])
         keys <- names(formals)
         dots <- match("...", keys)
@@ -441,7 +454,9 @@ member_infer <- function(
         }
         if (!is.na(dots)) env["..."] <- list(member_value(type = "list", elements = actuals[!matched]))
         for (name in setdiff(keys, c(used, "..."))) {
-            env[name] <- list(if (identical(formals[[name]], quote(expr = ))) {
+            env[name] <- list(if (!defaults) {
+                unknown
+            } else if (identical(formals[[name]], quote(expr = ))) {
                 member_value(type = ".missing")
             } else {
                 infer(formals[[name]], env)
@@ -641,8 +656,8 @@ member_infer <- function(
         if (identical(key, "self") && !is.null(receiver)) {
             return(member_value(type = receiver))
         }
-        if (key %in% names(index$definitions) && (is.null(index$document_bindings) ||
-            !is.null(context$key))) {
+        if (!is.null(member_lookup(index$definitions, key)) && (is.null(index$document_bindings) ||
+                    !is.null(context$key))) {
             return(member_value(function_key = key))
         }
         # Imported roots must be supplied by the document/package lexical
@@ -809,6 +824,26 @@ member_infer <- function(
         return(member_value(type = ".never"))
     }
     args <- as.list(expr)[-1L]
+    callee <- NULL
+    if (!intrinsic(head)) {
+        callee <- infer(expr[[1L]])
+        if (!is.null(callee$function_key) && is.null(callee$metadata)) {
+            self <- callee$receiver_value
+            key <- if (is.null(self)) callee$function_key else paste(self$type, callee$function_key, sep = "|")
+            summary <- member_lookup(index$method_results, key)
+            if (is.null(summary)) summary <- member_lookup(index$method_results, callee$function_key)
+            if (!is.null(summary) && !any(vapply(args, identical, logical(1L), as.name("...")))) {
+                fn <- member_definition(index$definitions, callee$function_key)
+                if (member_head(fn, "function")) {
+                    actuals <- lapply(args, function(arg) unknown)
+                    if (is.null(bind_arguments(fn, actuals, list(), defaults = FALSE))) {
+                        return(member_value(reason = "argument_matching"))
+                    }
+                    return(summary)
+                }
+            }
+        }
+    }
     if (head == "missing" && intrinsic(head) && length(args) == 1L &&
         identical(args[[1L]], as.name("..."))) {
         dots <- member_lookup(bindings, "...")
@@ -1066,7 +1101,7 @@ member_infer <- function(
         }
         return(unknown)
     }
-    callee <- infer(expr[[1L]])
+    if (is.null(callee)) callee <- infer(expr[[1L]])
     if (identical(callee$metadata, "S7") && callee$function_key %in% member_s7_intrinsics) {
         return(member_s7_call(expr, index, bindings, budget, callee$function_key))
     }

R/member-metadata.R

@@ -363,6 +363,10 @@ member_index_thaw <- function(snapshot) {
         return(NULL)
     }
     index <- list2env(snapshot, parent = emptyenv())
+    if (!is.null(index$return_definitions) &&
+            !identical(index$return_definitions, digest::digest(index$definitions, algo = "xxhash64"))) {
+        index$method_results <- list()
+    }
     index$cache <- new.env(parent = emptyenv())
     index
 }

tests/testthat/test-completion.R

@@ -1784,3 +1784,59 @@ test_that("The first dollar trigger after an edit uses current Polars members",
     expect_false(any(c("fileext", "infer_schema_files", "row.names") %in% labels))
     expect_true(all(vapply(result$items, function(item) identical(item$data$type, "member"), logical(1L))))
 })
+
+test_that("The first dollar trigger completes assigned and collected Polars queries", {
+    skip_on_cran()
+    skip_if_not_installed("polars")
+    client <- language_client()
+    temp_file <- withr::local_tempfile(fileext = ".R")
+    uri <- path_to_uri(temp_file)
+    lines <- c(
+        "library(polars)", "",
+        "csv_file <- tempfile(fileext = \".csv\")",
+        "write.csv(iris, csv_file, row.names = FALSE)", "",
+        "q <- pl$scan_csv(csv_file, infer_schema_files = 10)", "",
+        "q1 <- q$filter(pl$col(\"Sepal.Length\") > 5)", "q1", "",
+        "q1 <- q$filter(pl$col(\"Sepal.Length\") > 5)$group_by(\"Species\")$agg(pl$all()$median())$collect()",
+        "q1 # collected", "",
+        "q2 <- q1$group_by(\"Species\")$agg(pl$all()$sum())", "q2 # collected"
+    )
+    did_open(client, temp_file, text = paste(lines, collapse = "\n"))
+    # Prepare package metadata using only q; q1 and q2 must complete on their
+    # first request, without warming their method summaries or retrying.
+    deadline <- Sys.time() + 15
+    repeat {
+        ready <- respond_completion(client, temp_file, c(7L, 8L), retry = FALSE)
+        if (any(vapply(ready$items, function(item) identical(item$data$type, "member"), logical(1L)))) break
+        if (Sys.time() > deadline) break
+        Sys.sleep(0.1)
+    }
+    expect_true(any(vapply(ready$items, function(item) identical(item$data$type, "member"), logical(1L))))
+    notify(client, "workspace/didChangeConfiguration", list(settings = list(parse_delay = 0.5)))
+    for (row in c(8L, 11L, 14L)) {
+        end <- list(line = row, character = 2L)
+        notify(client, "textDocument/didChange", list(
+            textDocument = list(uri = uri, version = row),
+            contentChanges = list(list(range = list(start = end, end = end), text = "$"))
+        ))
+        result <- respond(client, "textDocument/completion", list(
+            textDocument = list(uri = uri),
+            position = list(line = row, character = end$character + 1L),
+            context = list(triggerKind = 2L, triggerCharacter = "$")
+        ), retry = FALSE)
+        labels <- vapply(result$items, `[[`, character(1L), "label")
+        expect_true(all(c("filter", "group_by") %in% labels))
+        expect_identical("collect" %in% labels, row == 8L)
+        expect_identical("lazy" %in% labels, row != 8L)
+        expect_true(all(vapply(result$items, function(item) identical(item$data$type, "member"), logical(1L))))
+        expect_false(any(c("fileext", "infer_schema_files", "row.names") %in% labels))
+        # Restore a complete statement before editing the next receiver.
+        notify(client, "textDocument/didChange", list(
+            textDocument = list(uri = uri, version = row + 1L),
+            contentChanges = list(list(
+                range = list(start = end, end = list(line = row, character = end$character + 1L)),
+                text = ""
+            ))
+        ))
+    }
+})