Skip to content

Commit 63d755e

Browse files
authored
Merge pull request #93 from t-kalinowski/fix-array-nonscalar-reshape
Lower array(data=<non-scalar>, dim=...) via reshape()
2 parents c4e86c0 + f60c0b8 commit 63d755e

3 files changed

Lines changed: 333 additions & 7 deletions

File tree

R/r2f-constructors.R

Lines changed: 191 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -63,9 +63,9 @@ r2f_handlers[["character"]] <- r2f_handlers[["raw"]] <-
6363
.r2f_handler_not_implemented_yet
6464

6565

66-
r2f_handlers[["matrix"]] <- function(args, scope = NULL, ...) {
66+
r2f_handlers[["matrix"]] <- function(args, scope = NULL, ..., hoist = NULL) {
6767
args$data %||% stop("matrix(data=) must be provided, cannot be NA")
68-
out <- r2f(args$data, scope, ...)
68+
out <- r2f(args$data, scope, ..., hoist = hoist)
6969
out@value <- Variable(
7070
mode = out@value@mode,
7171
dims = r2dims(list(args$nrow, args$ncol), scope)
@@ -75,7 +75,7 @@ r2f_handlers[["matrix"]] <- function(args, scope = NULL, ...) {
7575
# TODO: reshape() if !passes_as_scalar(out)
7676
}
7777

78-
r2f_handlers[["array"]] <- function(args, scope = NULL, ...) {
78+
r2f_handlers[["array"]] <- function(args, scope = NULL, ..., hoist = NULL) {
7979
args$data %||% stop("array(data=) must be provided, cannot be NA")
8080
if (is.null(args$dim)) {
8181
stop("array(dim=) must be provided, cannot be NA")
@@ -84,14 +84,199 @@ r2f_handlers[["array"]] <- function(args, scope = NULL, ...) {
8484
stop("array(dimnames=) not supported")
8585
}
8686

87-
out <- r2f(args$data, scope, ...)
87+
dim_to_dims <- function(dim_arg) {
88+
if (
89+
is.atomic(dim_arg) &&
90+
typeof(dim_arg) %in% c("integer", "double")
91+
) {
92+
if (!length(dim_arg) || anyNA(dim_arg)) {
93+
stop(
94+
"array(dim=) must be non-empty and must not contain NA",
95+
call. = FALSE
96+
)
97+
}
98+
dim_arg <- vapply(
99+
dim_arg,
100+
function(x) {
101+
if (!is_wholenumber(x)) {
102+
stop(
103+
"array(dim=) must be whole numbers, found: ",
104+
x,
105+
call. = FALSE
106+
)
107+
}
108+
as.integer(x)
109+
},
110+
integer(1L)
111+
)
112+
return(as.list(dim_arg))
113+
}
114+
115+
if (is.call(dim_arg) && is.symbol(dim_arg[[1L]])) {
116+
op <- as.character(dim_arg[[1L]])
117+
if (op == ":") {
118+
if (length(dim_arg) != 3L) {
119+
stop("bad dim sequence", call. = FALSE)
120+
}
121+
from <- dim_arg[[2L]]
122+
to <- dim_arg[[3L]]
123+
if (
124+
!(is.atomic(from) && length(from) == 1L && is_wholenumber(from)) ||
125+
!(is.atomic(to) && length(to) == 1L && is_wholenumber(to))
126+
) {
127+
stop(
128+
"array(dim=) only supports literal sequences like 2:4",
129+
call. = FALSE
130+
)
131+
}
132+
return(as.list(seq.int(as.integer(from), as.integer(to))))
133+
}
134+
}
135+
136+
if (is.symbol(dim_arg)) {
137+
var <- get0(as.character(dim_arg), scope)
138+
if (
139+
inherits(var, Variable) &&
140+
var@mode %in% c("integer", "double") &&
141+
var@rank == 1L &&
142+
(is.language(var@r) || is.atomic(var@r)) &&
143+
!identical(var@r, dim_arg)
144+
) {
145+
return(dim_to_dims(var@r))
146+
}
147+
}
148+
149+
r2dims(dim_arg, scope)
150+
}
151+
152+
out <- r2f(args$data, scope, ..., hoist = hoist)
153+
target_dims <- dim_to_dims(args$dim)
154+
if (!length(target_dims)) {
155+
stop("array(dim=) must not be empty", call. = FALSE)
156+
}
88157
if (!passes_as_scalar(out@value)) {
89-
stop("array(data=) must be a scalar for now")
158+
# R semantics: `array()` flattens its input (dropping dim) then reshapes.
159+
# We implement this as Fortran `reshape()`. Recycling (i.e. expanding a
160+
# shorter SOURCE to a larger target shape) is not supported.
161+
dims_f <- dims2f(target_dims, scope)
162+
scalar_target <- !nzchar(dims_f) && length(target_dims) == 1L
163+
if (scalar_target) {
164+
# `dim = 1` is scalar-like in quickr (rank-1 length-1 is declared scalar).
165+
# Avoid `reshape(..., [1])` (rank-1) and instead return the first element.
166+
if (is.null(hoist)) {
167+
stop("internal error: array() requires hoist context", call. = FALSE)
168+
}
169+
target_dims <- list(1L)
170+
tmp <- hoist$declare_tmp(mode = out@value@mode, dims = out@value@dims)
171+
hoist$emit(glue("{tmp@name} = {out}"))
172+
idxs <- rep("1", out@value@rank)
173+
out <- Fortran(
174+
glue("{tmp@name}({str_flatten_commas(idxs)})"),
175+
Variable(mode = out@value@mode, dims = list(1L))
176+
)
177+
} else {
178+
if (!nzchar(dims_f)) {
179+
dims_f <- "1"
180+
}
181+
if (grepl(":", dims_f, fixed = TRUE)) {
182+
stop("array(dim=) must be known", call. = FALSE)
183+
}
184+
shape <- glue("int([{dims_f}])")
185+
186+
data_r <- args$data
187+
is_fill_constructor <-
188+
is.call(data_r) &&
189+
is.symbol(data_r[[1L]]) &&
190+
as.character(data_r[[1L]]) %in%
191+
c(
192+
"logical",
193+
"integer",
194+
"double",
195+
"numeric"
196+
)
197+
198+
axis_terms <- vapply(
199+
target_dims,
200+
function(d) {
201+
axis <- dims2f(list(d), scope)
202+
if (!nzchar(axis)) {
203+
"1"
204+
} else {
205+
axis
206+
}
207+
},
208+
character(1L)
209+
)
210+
n_expr <- if (length(axis_terms) == 1L) {
211+
axis_terms[[1L]]
212+
} else {
213+
paste0("(", paste0("(", axis_terms, ")", collapse = " * "), ")")
214+
}
215+
216+
known_prod <- function(dims) {
217+
if (is.null(dims) || !length(dims)) {
218+
return(1)
219+
}
220+
vals <- vapply(
221+
dims,
222+
function(d) {
223+
if (
224+
is.atomic(d) &&
225+
length(d) == 1L &&
226+
!is.na(d) &&
227+
is_wholenumber(d)
228+
) {
229+
as.double(d)
230+
} else {
231+
NA_real_
232+
}
233+
},
234+
double(1L)
235+
)
236+
if (anyNA(vals)) {
237+
return(NA_real_)
238+
}
239+
prod(vals)
240+
}
241+
242+
source <- if (is_fill_constructor) {
243+
i <- scope@get_unique_var("integer")
244+
glue("[({out}, {i}=1, int({n_expr}))]")
245+
} else {
246+
n_target <- known_prod(target_dims)
247+
n_source <- known_prod(out@value@dims)
248+
if (!is.na(n_target) && !is.na(n_source) && n_target > n_source) {
249+
stop(
250+
"array() reshape does not support recycling: prod(dim)=",
251+
n_target,
252+
" > length(data)=",
253+
n_source,
254+
call. = FALSE
255+
)
256+
}
257+
if (!is.null(hoist)) {
258+
mark_scope_uses_errors(scope)
259+
err <- quickr_error_fortran_lines(
260+
"array() reshape does not support recycling (data shorter than prod(dim))",
261+
scope = scope
262+
)
263+
hoist$emit(glue("if (int({n_expr}) > size({out})) then"))
264+
hoist$emit(paste0(" ", err))
265+
hoist$emit("end if")
266+
}
267+
268+
# RESHAPE() requires `SOURCE` to be an array expression; array constructors
269+
# flatten array-valued expressions (which matches R's array() semantics).
270+
glue("[{out}]")
271+
}
272+
273+
out <- Fortran(glue("reshape({source}, {shape})"), out@value)
274+
}
90275
}
91276

92277
out@value <- Variable(
93278
mode = out@value@mode,
94-
dims = r2dims(args$dim, scope)
279+
dims = target_dims
95280
)
96281
out
97282
}
Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
test_that("array() supports reshaping non-scalar data", {
2+
fn <- function(x) {
3+
declare(type(x = integer(2L, 3L, 4L)))
4+
array(as.double(x), dim = c(2L, 3L, 4L))
5+
}
6+
7+
set.seed(1)
8+
x <- array(sample(1:10, 24, replace = TRUE), dim = c(2L, 3L, 4L))
9+
expect_quick_identical(fn, list(x))
10+
})
11+
12+
test_that("array() reshape accepts numeric dim vectors", {
13+
fn <- function(x) {
14+
declare(type(x = integer(2L, 3L, 4L)))
15+
array(as.double(x), dim = c(2, 3, 4))
16+
}
17+
18+
set.seed(1)
19+
x <- array(sample(1:10, 24, replace = TRUE), dim = c(2L, 3L, 4L))
20+
expect_quick_identical(fn, list(x))
21+
})
22+
23+
test_that("array() reshape accepts scalar dims", {
24+
fn <- function(x) {
25+
declare(type(x = integer(24L)))
26+
# Rank-1 arrays carry a `dim` attribute in base R, but quickr treats them as
27+
# plain vectors; wrap in `c()` so both sides compare identically while still
28+
# exercising the `array(dim=scalar)` lowering.
29+
c(array(as.double(x), dim = 24))
30+
}
31+
32+
set.seed(1)
33+
x <- sample(1:10, 24, replace = TRUE)
34+
expect_quick_identical(fn, list(x))
35+
})
36+
37+
test_that("array() reshape works when data is scalar-emitted (e.g. integer(n))", {
38+
fn <- function() {
39+
# `integer(3)` currently lowers to scalar `0` with a non-scalar value shape.
40+
# The array() reshape path must produce valid Fortran anyway.
41+
array(integer(3L), dim = c(1L, 3L))
42+
}
43+
44+
expect_quick_identical(fn, list())
45+
})
46+
47+
test_that("array() reshape accepts literal dim vectors in the AST", {
48+
dim_const <- c(2L, 3L, 4L)
49+
fn <- eval(bquote(function(x) {
50+
declare(type(x = integer(2L, 3L, 4L)))
51+
array(as.double(x), dim = .(dim_const))
52+
}))
53+
54+
set.seed(1)
55+
x <- array(sample(1:10, 24, replace = TRUE), dim = c(2L, 3L, 4L))
56+
expect_quick_identical(fn, list(x))
57+
})
58+
59+
test_that("array() reshape accepts dim as a literal sequence (2:4)", {
60+
fn <- function(x) {
61+
declare(type(x = integer(2L, 3L, 4L)))
62+
array(as.double(x), dim = 2:4)
63+
}
64+
65+
set.seed(1)
66+
x <- array(sample(1:10, 24, replace = TRUE), dim = c(2L, 3L, 4L))
67+
expect_quick_identical(fn, list(x))
68+
})
69+
70+
test_that("array() reshape accepts dim passed as a variable bound to a literal sequence", {
71+
fn <- function(x) {
72+
declare(type(x = integer(2L, 3L, 4L)))
73+
d <- 2:4
74+
array(as.double(x), dim = d)
75+
}
76+
77+
set.seed(1)
78+
x <- array(sample(1:10, 24, replace = TRUE), dim = c(2L, 3L, 4L))
79+
expect_quick_identical(fn, list(x))
80+
})
81+
82+
test_that("array() reshape fails early when recycling would be required", {
83+
fn <- function() {
84+
array(c(1L, 2L, 3L), dim = c(2L, 2L))
85+
}
86+
87+
# Fortran `reshape()` errors when the source is too short; quickr should stop
88+
# before generating uncompilable code.
89+
expect_error(quick(fn), "does not support recycling")
90+
})
91+
92+
test_that("array() fill reshape handles dim expressions that lower to comma-containing Fortran", {
93+
fn <- function(y, x) {
94+
declare(type(y = double(NA, NA)), type(x = double(nrow(y), ncol(y))))
95+
96+
# `dim(x)` uses the dims declared above, which dims2f() lowers to
97+
# `size(y, 1), size(y, 2)` (commas inside expressions). Codegen must not
98+
# split on commas in Fortran output.
99+
array(integer(nrow(y) * ncol(y)), dim = dim(x))
100+
}
101+
102+
set.seed(1)
103+
y <- matrix(runif(6), 2, 3)
104+
x <- y
105+
expect_quick_identical(fn, list(y, x))
106+
})
107+
108+
test_that("array() reshape supports dim = 1 for non-scalar data", {
109+
fn <- function(x) {
110+
declare(type(x = integer(2L, 3L, 4L)))
111+
# Rank-1 length-1 arrays are scalar-like in quickr; index the first element
112+
# to compare against base R without relying on `dim` attributes.
113+
array(as.double(x), dim = 1L)[1]
114+
}
115+
116+
set.seed(1)
117+
x <- array(sample(1:10, 24, replace = TRUE), dim = c(2L, 3L, 4L))
118+
expect_quick_identical(fn, list(x))
119+
})
120+
121+
test_that("array() forwards hoist when data needs hoisted temporaries", {
122+
fn <- function(x, y) {
123+
declare(type(x = double(2L, 2L)), type(y = double(2L, 2L)))
124+
array((x + y)[, 1], dim = c(2L, 1L))
125+
}
126+
127+
set.seed(1)
128+
x <- matrix(runif(4), 2, 2)
129+
y <- matrix(runif(4), 2, 2)
130+
expect_quick_identical(fn, list(x, y))
131+
})
132+
133+
test_that("array() rejects empty dim vectors (dim=c())", {
134+
fn <- function(x) {
135+
declare(type(x = double(2L, 2L)))
136+
array(as.double(x), dim = c())
137+
}
138+
139+
# Base R errors here ("'dims' cannot be of length 0"); quickr should fail
140+
# early too, rather than emitting rank-mismatched Fortran.
141+
expect_error(quick(fn), "dim")
142+
})

tests/testthat/test-r2f-r-attr.R

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,4 +13,3 @@ test_that("r2f() attaches `r` metadata for bind(c) logical symbols", {
1313
expect_true(inherits(a_var, Variable))
1414
expect_identical(a_var@r, quote(m))
1515
})
16-

0 commit comments

Comments
 (0)