Skip to content

Commit c4e86c0

Browse files
t-kalinowskit-kalinowski
authored
Merge pull request #92 from t-kalinowski/add-any-and-all-handlers
Fix any()/all() reductions with masked subsets and scalar semantics
2 parents 2da9d92 + 1d1340e commit c4e86c0

2 files changed

Lines changed: 336 additions & 1 deletion

File tree

‎R/r2f-reductions.R‎

Lines changed: 154 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
# r2f-reductions.R
2-
# Handlers for reduction operations: max, min, sum, prod, which.max, which.min
2+
# Handlers for reduction operations:
3+
# - numeric: max, min, sum, prod
4+
# - logical: any, all
5+
# - index: which.max, which.min
36

47
# --- Handlers ---
58

@@ -68,6 +71,156 @@ register_r2f_handler(
6871
}
6972
)
7073

74+
register_r2f_handler(
75+
c("any", "all"),
76+
function(
77+
args,
78+
scope,
79+
...
80+
) {
81+
# For now, we only support the most common `any(x)` / `all(x)` shape.
82+
# We intentionally do not support named arguments like `na.rm`.
83+
arg_names <- names(args) %||% character()
84+
if (length(arg_names) && any(nzchar(arg_names))) {
85+
stop(
86+
"any()/all() do not support named arguments (e.g. `na.rm`)",
87+
call. = FALSE
88+
)
89+
}
90+
91+
call_name <- last(list(...)$calls)
92+
intrinsic <- switch(
93+
call_name,
94+
any = "any",
95+
all = "all",
96+
stop("internal error: unexpected call: ", call_name, call. = FALSE)
97+
)
98+
99+
# Match R's base semantics: any() == FALSE, all() == TRUE.
100+
if (length(args) == 0L) {
101+
lit <- if (identical(call_name, "any")) ".false." else ".true."
102+
return(Fortran(lit, Variable("logical")))
103+
}
104+
105+
reduce_arg <- function(arg) {
106+
mask_hoist <- create_mask_hoist()
107+
x <- r2f(arg, scope, ..., hoist_mask = mask_hoist$try_set)
108+
if (mask_hoist$has_conflict()) {
109+
stop(
110+
"reduction expressions only support a single logical mask",
111+
call. = FALSE
112+
)
113+
}
114+
115+
if (!identical(x@value@mode, "logical")) {
116+
stop("any()/all() only implemented for logical", call. = FALSE)
117+
}
118+
119+
hoisted_mask <- mask_hoist$get_hoisted()
120+
121+
# Scalar logical: any(x) == x, all(x) == x
122+
if (x@value@is_scalar) {
123+
if (is.null(hoisted_mask)) {
124+
# `c(FALSE)` lowers to a 1-element Fortran array constructor
125+
# (`[.false.]`) but any()/all() must still return scalars.
126+
x_code <- trimws(as.character(x))
127+
if (startsWith(x_code, "[")) {
128+
return(Fortran(glue("{intrinsic}({x})"), Variable("logical")))
129+
}
130+
return(x)
131+
}
132+
133+
# For scalar `x`, `x[mask]` is empty iff `!any(mask)`.
134+
#
135+
# Note: `logical(1)` masks are represented as rank-1 (dims = list(1L))
136+
# but pass as scalars in the ABI and must *not* be wrapped in `any()` /
137+
# `all()` (compilers reject `any()` / `all()` on scalar arguments).
138+
#
139+
# Conversely, literal masks like `c(FALSE)` compile to array constructors
140+
# (e.g. `[ .false. ]`) and must be reduced to a scalar condition.
141+
mask_code <- trimws(as.character(hoisted_mask))
142+
is_array_ctor <- startsWith(mask_code, "[")
143+
mask_is_scalar <-
144+
!is.null(hoisted_mask@value) &&
145+
passes_as_scalar(hoisted_mask@value) &&
146+
!is_array_ctor
147+
148+
mask_len1 <-
149+
!is.null(hoisted_mask@value) &&
150+
identical(hoisted_mask@value@dims, list(1L))
151+
152+
if (!mask_is_scalar && !mask_len1) {
153+
stop(
154+
"any()/all(): scalar masked subsets only support scalar or length-1 masks",
155+
call. = FALSE
156+
)
157+
}
158+
159+
mask_scalar <- if (mask_is_scalar) {
160+
glue("{hoisted_mask}")
161+
} else {
162+
glue("any({hoisted_mask})")
163+
}
164+
165+
# When `[` hoists a scalar mask (x[mask] -> x with a hoisted mask),
166+
# we must preserve empty-selection semantics:
167+
# - any(logical(0)) == FALSE
168+
# - all(logical(0)) == TRUE
169+
identity <- if (identical(call_name, "any")) ".false." else ".true."
170+
x_code <- trimws(as.character(x))
171+
x_scalar <- if (startsWith(x_code, "[")) {
172+
glue("{intrinsic}({x})")
173+
} else {
174+
glue("{x}")
175+
}
176+
return(Fortran(
177+
glue("merge({x_scalar}, {identity}, {mask_scalar})"),
178+
Variable("logical", x@value@dims)
179+
))
180+
}
181+
182+
x_expr <- if (is.null(hoisted_mask)) {
183+
glue("{x}")
184+
} else {
185+
# Avoid `pack()` temporaries. For a mask-selected subset:
186+
# - any(x[mask]) is equivalent to any(x .and. mask)
187+
# - all(x[mask]) is equivalent to all((.not. mask) .or. x)
188+
# Both preserve empty-selection semantics.
189+
#
190+
# Note: A length-1 mask constructor like `c(TRUE)` compiles to a rank-1
191+
# array constructor (`[ .true. ]`). In R, this is recycled as a scalar
192+
# mask, so we must scalarize it to keep elementwise ops conformable.
193+
mask_code <- trimws(as.character(hoisted_mask))
194+
mask_is_array_ctor <- startsWith(mask_code, "[")
195+
mask_ctor_len1 <-
196+
mask_is_array_ctor &&
197+
!is.null(hoisted_mask@value) &&
198+
identical(hoisted_mask@value@dims, list(1L))
199+
mask_expr <- if (mask_ctor_len1) {
200+
glue("any({hoisted_mask})")
201+
} else {
202+
glue("{hoisted_mask}")
203+
}
204+
if (identical(call_name, "any")) {
205+
glue("(({x}) .and. ({mask_expr}))")
206+
} else {
207+
glue("((.not. ({mask_expr})) .or. ({x}))")
208+
}
209+
}
210+
211+
Fortran(glue("{intrinsic}({x_expr})"), Variable("logical"))
212+
}
213+
214+
if (length(args) == 1L) {
215+
return(reduce_arg(args[[1L]]))
216+
}
217+
218+
args <- lapply(args, reduce_arg)
219+
op <- if (identical(call_name, "any")) ".or." else ".and."
220+
Fortran(glue("({str_flatten(args, glue(' {op} '))})"), Variable("logical"))
221+
}
222+
)
223+
71224

72225
r2f_handlers[["which.max"]] <- r2f_handlers[["which.min"]] <-
73226
function(args, scope = NULL, ...) {

‎tests/testthat/test-subset-reduction.R‎

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,188 @@ test_that("reduction intrinsics cover scalar, multi-arg, and mask cases", {
7373
expect_quick_equal(prod_two, list(c(1, 2, 3), c(4, 5, 6)))
7474
})
7575

76+
test_that("any/all reduction intrinsics cover scalar, multi-arg, and mask cases", {
77+
any_basic <- function(x) {
78+
declare(type(x = logical(NA)))
79+
any(x)
80+
}
81+
fsub <- r2f(any_basic)
82+
expect_match(as.character(fsub), "any\\(")
83+
expect_quick_identical(any_basic, list(c(FALSE, FALSE)), list(c(FALSE, TRUE)))
84+
85+
all_basic <- function(x) {
86+
declare(type(x = logical(NA)))
87+
all(x)
88+
}
89+
fsub <- r2f(all_basic)
90+
expect_match(as.character(fsub), "all\\(")
91+
expect_quick_identical(all_basic, list(c(TRUE, TRUE)), list(c(TRUE, FALSE)))
92+
93+
any_two <- function(a, b) {
94+
declare(type(a = logical(NA)), type(b = logical(NA)))
95+
any(a, b)
96+
}
97+
fsub <- r2f(any_two)
98+
expect_match(as.character(fsub), "\\.or\\.")
99+
expect_quick_identical(
100+
any_two,
101+
list(a = c(FALSE, FALSE), b = c(FALSE, FALSE)),
102+
list(a = c(FALSE, FALSE), b = c(TRUE, FALSE))
103+
)
104+
105+
all_two <- function(a, b) {
106+
declare(type(a = logical(NA)), type(b = logical(NA)))
107+
all(a, b)
108+
}
109+
fsub <- r2f(all_two)
110+
expect_match(as.character(fsub), "\\.and\\.")
111+
expect_quick_identical(
112+
all_two,
113+
list(a = c(TRUE, TRUE), b = c(TRUE, TRUE)),
114+
list(a = c(TRUE, TRUE), b = c(FALSE, TRUE))
115+
)
116+
117+
# Scalar masked subset: preserve empty-selection semantics.
118+
any_scalar_masked_empty <- function(x) {
119+
declare(type(x = logical(1)))
120+
any(x[c(FALSE)])
121+
}
122+
expect_quick_identical(any_scalar_masked_empty, list(TRUE), list(FALSE))
123+
124+
all_scalar_masked_empty <- function(x) {
125+
declare(type(x = logical(1)))
126+
all(x[c(FALSE)])
127+
}
128+
expect_quick_identical(all_scalar_masked_empty, list(TRUE), list(FALSE))
129+
130+
# Scalar masked subset with a scalar mask variable: should compile and
131+
# preserve empty-selection semantics.
132+
any_scalar_masked_var <- function(x, m) {
133+
declare(type(x = logical(1)), type(m = logical(1)))
134+
any(x[m])
135+
}
136+
expect_quick_identical(
137+
any_scalar_masked_var,
138+
list(x = TRUE, m = FALSE),
139+
list(x = FALSE, m = TRUE)
140+
)
141+
142+
all_scalar_masked_var <- function(x, m) {
143+
declare(type(x = logical(1)), type(m = logical(1)))
144+
all(x[m])
145+
}
146+
expect_quick_identical(
147+
all_scalar_masked_var,
148+
list(x = FALSE, m = FALSE),
149+
list(x = FALSE, m = TRUE)
150+
)
151+
152+
# Scalar `x` with a longer logical mask diverges from simple empty/non-empty
153+
# semantics in R because selecting out-of-range positions yields NAs. We don't
154+
# implement that behavior here, so fail fast instead of silently emitting
155+
# wrong code.
156+
any_scalar_masked_long_mask <- function(x) {
157+
declare(type(x = logical(1)))
158+
any(x[c(FALSE, TRUE)])
159+
}
160+
expect_error(r2f(any_scalar_masked_long_mask), "scalar masked subsets")
161+
162+
all_scalar_masked_long_mask <- function(x) {
163+
declare(type(x = logical(1)))
164+
all(x[c(FALSE, TRUE)])
165+
}
166+
expect_error(r2f(all_scalar_masked_long_mask), "scalar masked subsets")
167+
168+
# 1-element vector expressions like c(FALSE) compile to Fortran array
169+
# constructors (`[.false.]`) but any()/all() must still return scalars.
170+
any_array_ctor_len1 <- function() {
171+
any(c(FALSE))
172+
}
173+
expect_quick_identical(any_array_ctor_len1, list())
174+
175+
all_array_ctor_len1 <- function() {
176+
all(c(TRUE))
177+
}
178+
expect_quick_identical(all_array_ctor_len1, list())
179+
180+
# Masked reduction over a 1-element array constructor: `[` can hoist the mask
181+
# while leaving `x` as a rank-1 array constructor (`[.true.]`), but the
182+
# reduction result must still be scalar.
183+
any_array_ctor_len1_masked <- function() {
184+
any(c(TRUE)[c(TRUE)])
185+
}
186+
expect_quick_identical(any_array_ctor_len1_masked, list())
187+
188+
all_array_ctor_len1_masked_empty <- function() {
189+
all(c(TRUE)[c(FALSE)])
190+
}
191+
expect_quick_identical(all_array_ctor_len1_masked_empty, list())
192+
193+
# Length-1 mask constructors (c(TRUE)/c(FALSE)) are hoisted as rank-1 array
194+
# constructors but should behave like scalars via recycling in R. The lowering
195+
# must scalarize them so elementwise `.and.`/`.or.` remain conformable.
196+
any_mask_ctor_len1 <- function(x) {
197+
declare(type(x = double(NA)))
198+
pred <- x > 1
199+
any(pred[c(TRUE)])
200+
}
201+
fsub <- r2f(any_mask_ctor_len1)
202+
expect_false(grepl("\\.and\\. \\(\\[", as.character(fsub)))
203+
expect_quick_identical(
204+
any_mask_ctor_len1,
205+
list(c(-2, -1)),
206+
list(c(0.2, 2))
207+
)
208+
209+
all_mask_ctor_len1_empty <- function(x) {
210+
declare(type(x = double(NA)))
211+
pred <- x > 1
212+
all(pred[c(FALSE)])
213+
}
214+
fsub <- r2f(all_mask_ctor_len1_empty)
215+
expect_false(grepl("\\.not\\. \\(\\[", as.character(fsub)))
216+
expect_false(grepl("\\.or\\. \\(\\[", as.character(fsub)))
217+
expect_quick_identical(
218+
all_mask_ctor_len1_empty,
219+
list(c(-2, -1)),
220+
list(c(0.2, 2))
221+
)
222+
223+
# Masked subset: preserve empty-selection semantics without pack() temporaries.
224+
any_masked <- function(x) {
225+
declare(type(x = double(NA)))
226+
pred <- x > 1
227+
any(pred[x > 0])
228+
}
229+
fsub <- r2f(any_masked)
230+
expect_false(grepl("pack\\(", as.character(fsub)))
231+
expect_match(as.character(fsub), "\\.and\\.")
232+
expect_match(as.character(fsub), "any\\(")
233+
expect_quick_identical(
234+
any_masked,
235+
list(c(-2, -1)), # empty selection -> any(logical(0)) == FALSE
236+
list(c(0.2, 0.3)), # selection all FALSE
237+
list(c(0.5, 2)) # selection contains TRUE
238+
)
239+
240+
all_masked <- function(x) {
241+
declare(type(x = double(NA)))
242+
pred <- x > 1
243+
all(pred[x > 0])
244+
}
245+
fsub <- r2f(all_masked)
246+
expect_false(grepl("pack\\(", as.character(fsub)))
247+
expect_match(as.character(fsub), "\\.not\\.")
248+
expect_match(as.character(fsub), "\\.or\\.")
249+
expect_match(as.character(fsub), "all\\(")
250+
expect_quick_identical(
251+
all_masked,
252+
list(c(-2, -1)), # empty selection -> all(logical(0)) == TRUE
253+
list(c(0.2, 0.3)), # selection all FALSE -> FALSE
254+
list(c(0.5, 2)) # selection has a FALSE -> FALSE
255+
)
256+
})
257+
76258

77259
test_that("1x1 subsetting keeps dims and C bridge builds", {
78260
fn <- function(m) {

0 commit comments

Comments
 (0)