Relax array(dim=) checks for reshape

Author: t-kalinowski

R/r2f-constructors.R

@@ -89,13 +89,20 @@ r2f_handlers[["array"]] <- function(args, scope = NULL, ...) {
     # R semantics: `array()` flattens its input (dropping dim) then reshapes.
     # We implement this as `reshape()`; recycling is not supported here.
     dim_vec <- r2f(args$dim, scope, ...)
-    if (is.null(dim_vec@value) || dim_vec@value@mode != "integer") {
-      stop("array(dim=) must be an integer vector", call. = FALSE)
+    if (
+      is.null(dim_vec@value) ||
+        !(dim_vec@value@mode %in% c("integer", "double"))
+    ) {
+      stop("array(dim=) must be an integer/numeric value", call. = FALSE)
     }
-    if (dim_vec@value@rank != 1L) {
-      stop("array(dim=) must be a 1-d integer vector", call. = FALSE)
+    shape <- if (passes_as_scalar(dim_vec@value)) {
+      glue("[int({dim_vec})]")
+    } else if (dim_vec@value@rank == 1L) {
+      glue("int({dim_vec})")
+    } else {
+      stop("array(dim=) must be a scalar or 1-d vector", call. = FALSE)
     }
-    out <- Fortran(glue("reshape({out}, int({dim_vec}))"), out@value)
+    out <- Fortran(glue("reshape({out}, {shape})"), out@value)
   }
 
   out@value <- Variable(

tests/testthat/test-array-reshape.R

@@ -8,3 +8,28 @@ test_that("array() supports reshaping non-scalar data", {
   x <- array(sample(1:10, 24, replace = TRUE), dim = c(2L, 3L, 4L))
   expect_quick_identical(fn, list(x))
 })
+
+test_that("array() reshape accepts numeric dim vectors", {
+  fn <- function(x) {
+    declare(type(x = integer(2L, 3L, 4L)))
+    array(as.double(x), dim = c(2, 3, 4))
+  }
+
+  set.seed(1)
+  x <- array(sample(1:10, 24, replace = TRUE), dim = c(2L, 3L, 4L))
+  expect_quick_identical(fn, list(x))
+})
+
+test_that("array() reshape accepts scalar dims", {
+  fn <- function(x) {
+    declare(type(x = integer(24L)))
+    # Rank-1 arrays carry a `dim` attribute in base R, but quickr treats them as
+    # plain vectors; wrap in `c()` so both sides compare identically while still
+    # exercising the `array(dim=scalar)` lowering.
+    c(array(as.double(x), dim = 24))
+  }
+
+  set.seed(1)
+  x <- sample(1:10, 24, replace = TRUE)
+  expect_quick_identical(fn, list(x))
+})