Lower array(data=<non-scalar>, dim=...) via reshape()

R/r2f-constructors.R

@@ -63,9 +63,9 @@ r2f_handlers[["character"]] <- r2f_handlers[["raw"]] <-
   .r2f_handler_not_implemented_yet
 
 
-r2f_handlers[["matrix"]] <- function(args, scope = NULL, ...) {
+r2f_handlers[["matrix"]] <- function(args, scope = NULL, ..., hoist = NULL) {
   args$data %||% stop("matrix(data=) must be provided, cannot be NA")
-  out <- r2f(args$data, scope, ...)
+  out <- r2f(args$data, scope, ..., hoist = hoist)
   out@value <- Variable(
     mode = out@value@mode,
     dims = r2dims(list(args$nrow, args$ncol), scope)
@@ -75,7 +75,7 @@ r2f_handlers[["matrix"]] <- function(args, scope = NULL, ...) {
   # TODO: reshape() if !passes_as_scalar(out)
 }
 
-r2f_handlers[["array"]] <- function(args, scope = NULL, ...) {
+r2f_handlers[["array"]] <- function(args, scope = NULL, ..., hoist = NULL) {
   args$data %||% stop("array(data=) must be provided, cannot be NA")
   if (is.null(args$dim)) {
     stop("array(dim=) must be provided, cannot be NA")
@@ -84,14 +84,199 @@ r2f_handlers[["array"]] <- function(args, scope = NULL, ...) {
     stop("array(dimnames=) not supported")
   }
 
-  out <- r2f(args$data, scope, ...)
+  dim_to_dims <- function(dim_arg) {
+    if (
+      is.atomic(dim_arg) &&
+        typeof(dim_arg) %in% c("integer", "double")
+    ) {
+      if (!length(dim_arg) || anyNA(dim_arg)) {
+        stop(
+          "array(dim=) must be non-empty and must not contain NA",
+          call. = FALSE
+        )
+      }
+      dim_arg <- vapply(
+        dim_arg,
+        function(x) {
+          if (!is_wholenumber(x)) {
+            stop(
+              "array(dim=) must be whole numbers, found: ",
+              x,
+              call. = FALSE
+            )
+          }
+          as.integer(x)
+        },
+        integer(1L)
+      )
+      return(as.list(dim_arg))
+    }
+
+    if (is.call(dim_arg) && is.symbol(dim_arg[[1L]])) {
+      op <- as.character(dim_arg[[1L]])
+      if (op == ":") {
+        if (length(dim_arg) != 3L) {
+          stop("bad dim sequence", call. = FALSE)
+        }
+        from <- dim_arg[[2L]]
+        to <- dim_arg[[3L]]
+        if (
+          !(is.atomic(from) && length(from) == 1L && is_wholenumber(from)) ||
+            !(is.atomic(to) && length(to) == 1L && is_wholenumber(to))
+        ) {
+          stop(
+            "array(dim=) only supports literal sequences like 2:4",
+            call. = FALSE
+          )
+        }
+        return(as.list(seq.int(as.integer(from), as.integer(to))))
+      }
+    }
+
+    if (is.symbol(dim_arg)) {
+      var <- get0(as.character(dim_arg), scope)
+      if (
+        inherits(var, Variable) &&
+          var@mode %in% c("integer", "double") &&
+          var@rank == 1L &&
+          (is.language(var@r) || is.atomic(var@r)) &&
+          !identical(var@r, dim_arg)
+      ) {
+        return(dim_to_dims(var@r))
+      }
+    }
+
+    r2dims(dim_arg, scope)
+  }
+
+  out <- r2f(args$data, scope, ..., hoist = hoist)
+  target_dims <- dim_to_dims(args$dim)
+  if (!length(target_dims)) {
+    stop("array(dim=) must not be empty", call. = FALSE)
+  }
   if (!passes_as_scalar(out@value)) {
-    stop("array(data=) must be a scalar for now")
+    # R semantics: `array()` flattens its input (dropping dim) then reshapes.
+    # We implement this as Fortran `reshape()`. Recycling (i.e. expanding a
+    # shorter SOURCE to a larger target shape) is not supported.
+    dims_f <- dims2f(target_dims, scope)
+    scalar_target <- !nzchar(dims_f) && length(target_dims) == 1L
+    if (scalar_target) {
+      # `dim = 1` is scalar-like in quickr (rank-1 length-1 is declared scalar).
+      # Avoid `reshape(..., [1])` (rank-1) and instead return the first element.
+      if (is.null(hoist)) {
+        stop("internal error: array() requires hoist context", call. = FALSE)
+      }
+      target_dims <- list(1L)
+      tmp <- hoist$declare_tmp(mode = out@value@mode, dims = out@value@dims)
+      hoist$emit(glue("{tmp@name} = {out}"))
+      idxs <- rep("1", out@value@rank)
+      out <- Fortran(
+        glue("{tmp@name}({str_flatten_commas(idxs)})"),
+        Variable(mode = out@value@mode, dims = list(1L))
+      )
+    } else {
+      if (!nzchar(dims_f)) {
+        dims_f <- "1"
+      }
+      if (grepl(":", dims_f, fixed = TRUE)) {
+        stop("array(dim=) must be known", call. = FALSE)
+      }
+      shape <- glue("int([{dims_f}])")
+
+      data_r <- args$data
+      is_fill_constructor <-
+        is.call(data_r) &&
+        is.symbol(data_r[[1L]]) &&
+        as.character(data_r[[1L]]) %in%
+          c(
+            "logical",
+            "integer",
+            "double",
+            "numeric"
+          )
+
+      axis_terms <- vapply(
+        target_dims,
+        function(d) {
+          axis <- dims2f(list(d), scope)
+          if (!nzchar(axis)) {
+            "1"
+          } else {
+            axis
+          }
+        },
+        character(1L)
+      )
+      n_expr <- if (length(axis_terms) == 1L) {
+        axis_terms[[1L]]
+      } else {
+        paste0("(", paste0("(", axis_terms, ")", collapse = " * "), ")")
+      }
+
+      known_prod <- function(dims) {
+        if (is.null(dims) || !length(dims)) {
+          return(1)
+        }
+        vals <- vapply(
+          dims,
+          function(d) {
+            if (
+              is.atomic(d) &&
+                length(d) == 1L &&
+                !is.na(d) &&
+                is_wholenumber(d)
+            ) {
+              as.double(d)
+            } else {
+              NA_real_
+            }
+          },
+          double(1L)
+        )
+        if (anyNA(vals)) {
+          return(NA_real_)
+        }
+        prod(vals)
+      }
+
+      source <- if (is_fill_constructor) {
+        i <- scope@get_unique_var("integer")
+        glue("[({out}, {i}=1, int({n_expr}))]")
+      } else {
+        n_target <- known_prod(target_dims)
+        n_source <- known_prod(out@value@dims)
+        if (!is.na(n_target) && !is.na(n_source) && n_target > n_source) {
+          stop(
+            "array() reshape does not support recycling: prod(dim)=",
+            n_target,
+            " > length(data)=",
+            n_source,
+            call. = FALSE
+          )
+        }
+        if (!is.null(hoist)) {
+          mark_scope_uses_errors(scope)
+          err <- quickr_error_fortran_lines(
+            "array() reshape does not support recycling (data shorter than prod(dim))",
+            scope = scope
+          )
+          hoist$emit(glue("if (int({n_expr}) > size({out})) then"))
+          hoist$emit(paste0("  ", err))
+          hoist$emit("end if")
+        }
+
+        # RESHAPE() requires `SOURCE` to be an array expression; array constructors
+        # flatten array-valued expressions (which matches R's array() semantics).
+        glue("[{out}]")
+      }
+
+      out <- Fortran(glue("reshape({source}, {shape})"), out@value)
+    }
   }
 
   out@value <- Variable(
     mode = out@value@mode,
-    dims = r2dims(args$dim, scope)
+    dims = target_dims
   )
   out
 }

tests/testthat/test-array-reshape.R

@@ -0,0 +1,142 @@
+test_that("array() supports reshaping non-scalar data", {
+  fn <- function(x) {
+    declare(type(x = integer(2L, 3L, 4L)))
+    array(as.double(x), dim = c(2L, 3L, 4L))
+  }
+
+  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 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))
+})
+
+test_that("array() reshape works when data is scalar-emitted (e.g. integer(n))", {
+  fn <- function() {
+    # `integer(3)` currently lowers to scalar `0` with a non-scalar value shape.
+    # The array() reshape path must produce valid Fortran anyway.
+    array(integer(3L), dim = c(1L, 3L))
+  }
+
+  expect_quick_identical(fn, list())
+})
+
+test_that("array() reshape accepts literal dim vectors in the AST", {
+  dim_const <- c(2L, 3L, 4L)
+  fn <- eval(bquote(function(x) {
+    declare(type(x = integer(2L, 3L, 4L)))
+    array(as.double(x), dim = .(dim_const))
+  }))
+
+  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 dim as a literal sequence (2:4)", {
+  fn <- function(x) {
+    declare(type(x = integer(2L, 3L, 4L)))
+    array(as.double(x), dim = 2: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 dim passed as a variable bound to a literal sequence", {
+  fn <- function(x) {
+    declare(type(x = integer(2L, 3L, 4L)))
+    d <- 2:4
+    array(as.double(x), dim = d)
+  }
+
+  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 fails early when recycling would be required", {
+  fn <- function() {
+    array(c(1L, 2L, 3L), dim = c(2L, 2L))
+  }
+
+  # Fortran `reshape()` errors when the source is too short; quickr should stop
+  # before generating uncompilable code.
+  expect_error(quick(fn), "does not support recycling")
+})
+
+test_that("array() fill reshape handles dim expressions that lower to comma-containing Fortran", {
+  fn <- function(y, x) {
+    declare(type(y = double(NA, NA)), type(x = double(nrow(y), ncol(y))))
+
+    # `dim(x)` uses the dims declared above, which dims2f() lowers to
+    # `size(y, 1), size(y, 2)` (commas inside expressions). Codegen must not
+    # split on commas in Fortran output.
+    array(integer(nrow(y) * ncol(y)), dim = dim(x))
+  }
+
+  set.seed(1)
+  y <- matrix(runif(6), 2, 3)
+  x <- y
+  expect_quick_identical(fn, list(y, x))
+})
+
+test_that("array() reshape supports dim = 1 for non-scalar data", {
+  fn <- function(x) {
+    declare(type(x = integer(2L, 3L, 4L)))
+    # Rank-1 length-1 arrays are scalar-like in quickr; index the first element
+    # to compare against base R without relying on `dim` attributes.
+    array(as.double(x), dim = 1L)[1]
+  }
+
+  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() forwards hoist when data needs hoisted temporaries", {
+  fn <- function(x, y) {
+    declare(type(x = double(2L, 2L)), type(y = double(2L, 2L)))
+    array((x + y)[, 1], dim = c(2L, 1L))
+  }
+
+  set.seed(1)
+  x <- matrix(runif(4), 2, 2)
+  y <- matrix(runif(4), 2, 2)
+  expect_quick_identical(fn, list(x, y))
+})
+
+test_that("array() rejects empty dim vectors (dim=c())", {
+  fn <- function(x) {
+    declare(type(x = double(2L, 2L)))
+    array(as.double(x), dim = c())
+  }
+
+  # Base R errors here ("'dims' cannot be of length 0"); quickr should fail
+  # early too, rather than emitting rank-mismatched Fortran.
+  expect_error(quick(fn), "dim")
+})

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

@@ -13,4 +13,3 @@ test_that("r2f() attaches `r` metadata for bind(c) logical symbols", {
   expect_true(inherits(a_var, Variable))
   expect_identical(a_var@r, quote(m))
 })
-