Fix any()/all() reductions with masked subsets and scalar semantics

R/r2f-reductions.R

@@ -1,5 +1,8 @@
 # r2f-reductions.R
-# Handlers for reduction operations: max, min, sum, prod, which.max, which.min
+# Handlers for reduction operations:
+# - numeric: max, min, sum, prod
+# - logical: any, all
+# - index: which.max, which.min
 
 # --- Handlers ---
 
@@ -58,6 +61,156 @@ register_r2f_handler(
   }
 )
 
+register_r2f_handler(
+  c("any", "all"),
+  function(
+    args,
+    scope,
+    ...
+  ) {
+    # For now, we only support the most common `any(x)` / `all(x)` shape.
+    # We intentionally do not support named arguments like `na.rm`.
+    arg_names <- names(args) %||% character()
+    if (length(arg_names) && any(nzchar(arg_names))) {
+      stop(
+        "any()/all() do not support named arguments (e.g. `na.rm`)",
+        call. = FALSE
+      )
+    }
+
+    call_name <- last(list(...)$calls)
+    intrinsic <- switch(
+      call_name,
+      any = "any",
+      all = "all",
+      stop("internal error: unexpected call: ", call_name, call. = FALSE)
+    )
+
+    # Match R's base semantics: any() == FALSE, all() == TRUE.
+    if (length(args) == 0L) {
+      lit <- if (identical(call_name, "any")) ".false." else ".true."
+      return(Fortran(lit, Variable("logical")))
+    }
+
+    reduce_arg <- function(arg) {
+      mask_hoist <- create_mask_hoist()
+      x <- r2f(arg, scope, ..., hoist_mask = mask_hoist$try_set)
+      if (mask_hoist$has_conflict()) {
+        stop(
+          "reduction expressions only support a single logical mask",
+          call. = FALSE
+        )
+      }
+
+      if (!identical(x@value@mode, "logical")) {
+        stop("any()/all() only implemented for logical", call. = FALSE)
+      }
+
+      hoisted_mask <- mask_hoist$get_hoisted()
+
+      # Scalar logical: any(x) == x, all(x) == x
+      if (x@value@is_scalar) {
+        if (is.null(hoisted_mask)) {
+          # `c(FALSE)` lowers to a 1-element Fortran array constructor
+          # (`[.false.]`) but any()/all() must still return scalars.
+          x_code <- trimws(as.character(x))
+          if (startsWith(x_code, "[")) {
+            return(Fortran(glue("{intrinsic}({x})"), Variable("logical")))
+          }
+          return(x)
+        }
+
+        # For scalar `x`, `x[mask]` is empty iff `!any(mask)`.
+        #
+        # Note: `logical(1)` masks are represented as rank-1 (dims = list(1L))
+        # but pass as scalars in the ABI and must *not* be wrapped in `any()` /
+        # `all()` (compilers reject `any()` / `all()` on scalar arguments).
+        #
+        # Conversely, literal masks like `c(FALSE)` compile to array constructors
+        # (e.g. `[ .false. ]`) and must be reduced to a scalar condition.
+        mask_code <- trimws(as.character(hoisted_mask))
+        is_array_ctor <- startsWith(mask_code, "[")
+        mask_is_scalar <-
+          !is.null(hoisted_mask@value) &&
+          passes_as_scalar(hoisted_mask@value) &&
+          !is_array_ctor
+
+        mask_len1 <-
+          !is.null(hoisted_mask@value) &&
+          identical(hoisted_mask@value@dims, list(1L))
+
+        if (!mask_is_scalar && !mask_len1) {
+          stop(
+            "any()/all(): scalar masked subsets only support scalar or length-1 masks",
+            call. = FALSE
+          )
+        }
+
+        mask_scalar <- if (mask_is_scalar) {
+          glue("{hoisted_mask}")
+        } else {
+          glue("any({hoisted_mask})")
+        }
+
+        # When `[` hoists a scalar mask (x[mask] -> x with a hoisted mask),
+        # we must preserve empty-selection semantics:
+        # - any(logical(0)) == FALSE
+        # - all(logical(0)) == TRUE
+        identity <- if (identical(call_name, "any")) ".false." else ".true."
+        x_code <- trimws(as.character(x))
+        x_scalar <- if (startsWith(x_code, "[")) {
+          glue("{intrinsic}({x})")
+        } else {
+          glue("{x}")
+        }
+        return(Fortran(
+          glue("merge({x_scalar}, {identity}, {mask_scalar})"),
+          Variable("logical", x@value@dims)
+        ))
+      }
+
+      x_expr <- if (is.null(hoisted_mask)) {
+        glue("{x}")
+      } else {
+        # Avoid `pack()` temporaries. For a mask-selected subset:
+        # - any(x[mask]) is equivalent to any(x .and. mask)
+        # - all(x[mask]) is equivalent to all((.not. mask) .or. x)
+        # Both preserve empty-selection semantics.
+        #
+        # Note: A length-1 mask constructor like `c(TRUE)` compiles to a rank-1
+        # array constructor (`[ .true. ]`). In R, this is recycled as a scalar
+        # mask, so we must scalarize it to keep elementwise ops conformable.
+        mask_code <- trimws(as.character(hoisted_mask))
+        mask_is_array_ctor <- startsWith(mask_code, "[")
+        mask_ctor_len1 <-
+          mask_is_array_ctor &&
+          !is.null(hoisted_mask@value) &&
+          identical(hoisted_mask@value@dims, list(1L))
+        mask_expr <- if (mask_ctor_len1) {
+          glue("any({hoisted_mask})")
+        } else {
+          glue("{hoisted_mask}")
+        }
+        if (identical(call_name, "any")) {
+          glue("(({x}) .and. ({mask_expr}))")
+        } else {
+          glue("((.not. ({mask_expr})) .or. ({x}))")
+        }
+      }
+
+      Fortran(glue("{intrinsic}({x_expr})"), Variable("logical"))
+    }
+
+    if (length(args) == 1L) {
+      return(reduce_arg(args[[1L]]))
+    }
+
+    args <- lapply(args, reduce_arg)
+    op <- if (identical(call_name, "any")) ".or." else ".and."
+    Fortran(glue("({str_flatten(args, glue(' {op} '))})"), Variable("logical"))
+  }
+)
+
 
 r2f_handlers[["which.max"]] <- r2f_handlers[["which.min"]] <-
   function(args, scope = NULL, ...) {

tests/testthat/test-subset-reduction.R

@@ -73,6 +73,188 @@ test_that("reduction intrinsics cover scalar, multi-arg, and mask cases", {
   expect_quick_equal(prod_two, list(c(1, 2, 3), c(4, 5, 6)))
 })
 
+test_that("any/all reduction intrinsics cover scalar, multi-arg, and mask cases", {
+  any_basic <- function(x) {
+    declare(type(x = logical(NA)))
+    any(x)
+  }
+  fsub <- r2f(any_basic)
+  expect_match(as.character(fsub), "any\\(")
+  expect_quick_identical(any_basic, list(c(FALSE, FALSE)), list(c(FALSE, TRUE)))
+
+  all_basic <- function(x) {
+    declare(type(x = logical(NA)))
+    all(x)
+  }
+  fsub <- r2f(all_basic)
+  expect_match(as.character(fsub), "all\\(")
+  expect_quick_identical(all_basic, list(c(TRUE, TRUE)), list(c(TRUE, FALSE)))
+
+  any_two <- function(a, b) {
+    declare(type(a = logical(NA)), type(b = logical(NA)))
+    any(a, b)
+  }
+  fsub <- r2f(any_two)
+  expect_match(as.character(fsub), "\\.or\\.")
+  expect_quick_identical(
+    any_two,
+    list(a = c(FALSE, FALSE), b = c(FALSE, FALSE)),
+    list(a = c(FALSE, FALSE), b = c(TRUE, FALSE))
+  )
+
+  all_two <- function(a, b) {
+    declare(type(a = logical(NA)), type(b = logical(NA)))
+    all(a, b)
+  }
+  fsub <- r2f(all_two)
+  expect_match(as.character(fsub), "\\.and\\.")
+  expect_quick_identical(
+    all_two,
+    list(a = c(TRUE, TRUE), b = c(TRUE, TRUE)),
+    list(a = c(TRUE, TRUE), b = c(FALSE, TRUE))
+  )
+
+  # Scalar masked subset: preserve empty-selection semantics.
+  any_scalar_masked_empty <- function(x) {
+    declare(type(x = logical(1)))
+    any(x[c(FALSE)])
+  }
+  expect_quick_identical(any_scalar_masked_empty, list(TRUE), list(FALSE))
+
+  all_scalar_masked_empty <- function(x) {
+    declare(type(x = logical(1)))
+    all(x[c(FALSE)])
+  }
+  expect_quick_identical(all_scalar_masked_empty, list(TRUE), list(FALSE))
+
+  # Scalar masked subset with a scalar mask variable: should compile and
+  # preserve empty-selection semantics.
+  any_scalar_masked_var <- function(x, m) {
+    declare(type(x = logical(1)), type(m = logical(1)))
+    any(x[m])
+  }
+  expect_quick_identical(
+    any_scalar_masked_var,
+    list(x = TRUE, m = FALSE),
+    list(x = FALSE, m = TRUE)
+  )
+
+  all_scalar_masked_var <- function(x, m) {
+    declare(type(x = logical(1)), type(m = logical(1)))
+    all(x[m])
+  }
+  expect_quick_identical(
+    all_scalar_masked_var,
+    list(x = FALSE, m = FALSE),
+    list(x = FALSE, m = TRUE)
+  )
+
+  # Scalar `x` with a longer logical mask diverges from simple empty/non-empty
+  # semantics in R because selecting out-of-range positions yields NAs. We don't
+  # implement that behavior here, so fail fast instead of silently emitting
+  # wrong code.
+  any_scalar_masked_long_mask <- function(x) {
+    declare(type(x = logical(1)))
+    any(x[c(FALSE, TRUE)])
+  }
+  expect_error(r2f(any_scalar_masked_long_mask), "scalar masked subsets")
+
+  all_scalar_masked_long_mask <- function(x) {
+    declare(type(x = logical(1)))
+    all(x[c(FALSE, TRUE)])
+  }
+  expect_error(r2f(all_scalar_masked_long_mask), "scalar masked subsets")
+
+  # 1-element vector expressions like c(FALSE) compile to Fortran array
+  # constructors (`[.false.]`) but any()/all() must still return scalars.
+  any_array_ctor_len1 <- function() {
+    any(c(FALSE))
+  }
+  expect_quick_identical(any_array_ctor_len1, list())
+
+  all_array_ctor_len1 <- function() {
+    all(c(TRUE))
+  }
+  expect_quick_identical(all_array_ctor_len1, list())
+
+  # Masked reduction over a 1-element array constructor: `[` can hoist the mask
+  # while leaving `x` as a rank-1 array constructor (`[.true.]`), but the
+  # reduction result must still be scalar.
+  any_array_ctor_len1_masked <- function() {
+    any(c(TRUE)[c(TRUE)])
+  }
+  expect_quick_identical(any_array_ctor_len1_masked, list())
+
+  all_array_ctor_len1_masked_empty <- function() {
+    all(c(TRUE)[c(FALSE)])
+  }
+  expect_quick_identical(all_array_ctor_len1_masked_empty, list())
+
+  # Length-1 mask constructors (c(TRUE)/c(FALSE)) are hoisted as rank-1 array
+  # constructors but should behave like scalars via recycling in R. The lowering
+  # must scalarize them so elementwise `.and.`/`.or.` remain conformable.
+  any_mask_ctor_len1 <- function(x) {
+    declare(type(x = double(NA)))
+    pred <- x > 1
+    any(pred[c(TRUE)])
+  }
+  fsub <- r2f(any_mask_ctor_len1)
+  expect_false(grepl("\\.and\\. \\(\\[", as.character(fsub)))
+  expect_quick_identical(
+    any_mask_ctor_len1,
+    list(c(-2, -1)),
+    list(c(0.2, 2))
+  )
+
+  all_mask_ctor_len1_empty <- function(x) {
+    declare(type(x = double(NA)))
+    pred <- x > 1
+    all(pred[c(FALSE)])
+  }
+  fsub <- r2f(all_mask_ctor_len1_empty)
+  expect_false(grepl("\\.not\\. \\(\\[", as.character(fsub)))
+  expect_false(grepl("\\.or\\. \\(\\[", as.character(fsub)))
+  expect_quick_identical(
+    all_mask_ctor_len1_empty,
+    list(c(-2, -1)),
+    list(c(0.2, 2))
+  )
+
+  # Masked subset: preserve empty-selection semantics without pack() temporaries.
+  any_masked <- function(x) {
+    declare(type(x = double(NA)))
+    pred <- x > 1
+    any(pred[x > 0])
+  }
+  fsub <- r2f(any_masked)
+  expect_false(grepl("pack\\(", as.character(fsub)))
+  expect_match(as.character(fsub), "\\.and\\.")
+  expect_match(as.character(fsub), "any\\(")
+  expect_quick_identical(
+    any_masked,
+    list(c(-2, -1)), # empty selection -> any(logical(0)) == FALSE
+    list(c(0.2, 0.3)), # selection all FALSE
+    list(c(0.5, 2)) # selection contains TRUE
+  )
+
+  all_masked <- function(x) {
+    declare(type(x = double(NA)))
+    pred <- x > 1
+    all(pred[x > 0])
+  }
+  fsub <- r2f(all_masked)
+  expect_false(grepl("pack\\(", as.character(fsub)))
+  expect_match(as.character(fsub), "\\.not\\.")
+  expect_match(as.character(fsub), "\\.or\\.")
+  expect_match(as.character(fsub), "all\\(")
+  expect_quick_identical(
+    all_masked,
+    list(c(-2, -1)), # empty selection -> all(logical(0)) == TRUE
+    list(c(0.2, 0.3)), # selection all FALSE -> FALSE
+    list(c(0.5, 2)) # selection has a FALSE -> FALSE
+  )
+})
+
 
 test_that("1x1 subsetting keeps dims and C bridge builds", {
   fn <- function(m) {