Support rectangular solve and qr.solve compilation

R/r2f-matrix-blas.R

@@ -649,8 +649,8 @@ lapack_solve <- function(
   assert_rank2_matrix(A, paste0(context, " expects a matrix for `a`"))
 
   a_dims <- matrix_dims(A)
-  assert_square_matrix(a_dims$rows, a_dims$cols, context)
-  n <- a_dims$rows
+  m <- a_dims$rows
+  n <- a_dims$cols
 
   b_rank <- B@value@rank
   assert_rhs_rank(
@@ -667,15 +667,15 @@ lapack_solve <- function(
   if (b_rank == 1L) {
     b_len <- dim_or_one(B, 1L)
     assert_conformable_dims(
-      n,
+      m,
       b_len,
       context = context,
       err_msg = paste0("non-conformable arguments in ", context)
     )
   } else {
     b_rows <- dim_or_one(B, 1L)
     assert_conformable_dims(
-      n,
+      m,
       b_rows,
       context = context,
       err_msg = paste0("non-conformable arguments in ", context)
@@ -685,15 +685,89 @@ lapack_solve <- function(
   A_name <- ensure_blas_operand_name(A, hoist)
   B_input_name <- ensure_blas_operand_name(B, hoist)
 
-  A_work <- hoist$declare_tmp(mode = "double", dims = list(n, n))
+  nrhs <- if (b_rank == 1L) 1L else dim_or_one(B, 2L)
+
+  square <- check_conformable(m, n)
+  if (square$ok && !square$unknown) {
+    A_work <- hoist$declare_tmp(mode = "double", dims = list(m, m))
+    hoist$emit(glue("{A_work@name} = {A_name}"))
+
+    expected_dims <- if (b_rank == 1L) list(n) else list(n, nrhs)
+    writes_to_dest <- FALSE
+    if (
+      can_use_output(
+        dest,
+        input_names = c(A_name, B_input_name),
+        expected_dims = expected_dims,
+        context = context,
+        allow_alias = B_input_name
+      )
+    ) {
+      out_var <- dest
+      out_name <- dest@name
+      writes_to_dest <- TRUE
+    } else {
+      out_var <- hoist$declare_tmp(mode = "double", dims = expected_dims)
+      out_name <- out_var@name
+    }
+    hoist$emit(glue("{out_name} = {B_input_name}"))
+
+    ipiv <- hoist$declare_tmp(mode = "integer", dims = list(m))
+    info <- hoist$declare_tmp(mode = "integer", dims = NULL)
+
+    hoist$emit(glue(
+      "call dgesv({blas_int(m)}, {blas_int(nrhs)}, {A_work@name}, {blas_int(m)}, {ipiv@name}, {out_name}, {blas_int(m)}, {info@name})"
+    ))
+
+    out <- Fortran(out_name, out_var)
+    if (writes_to_dest) {
+      attr(out, "writes_to_dest") <- TRUE
+    }
+    return(out)
+  }
+
+  A_work <- hoist$declare_tmp(mode = "double", dims = list(m, n))
   hoist$emit(glue("{A_work@name} = {A_name}"))
 
+  max_mn <- call("max", m, n)
+
+  B_work <- hoist$declare_tmp(mode = "double", dims = list(max_mn, nrhs))
+  m_f <- dims2f(list(m), scope)
+  if (!nzchar(m_f)) {
+    m_f <- "1"
+  }
+  n_f <- dims2f(list(n), scope)
+  if (!nzchar(n_f)) {
+    n_f <- "1"
+  }
+  nrhs_f <- dims2f(list(nrhs), scope)
+  if (!nzchar(nrhs_f)) {
+    nrhs_f <- "1"
+  }
+  hoist$emit(glue("{B_work@name} = 0.0_c_double"))
+  if (b_rank == 1L) {
+    hoist$emit(glue("{B_work@name}(1:{m_f}, 1) = {B_input_name}"))
+  } else {
+    hoist$emit(glue("{B_work@name}(1:{m_f}, 1:{nrhs_f}) = {B_input_name}"))
+  }
+
+  info <- hoist$declare_tmp(mode = "integer", dims = NULL)
+
+  mn <- call("min", m, n)
+  lwork <- call("max", 1L, call("+", mn, call("max", mn, nrhs)))
+  work <- hoist$declare_tmp(mode = "double", dims = list(lwork))
+
+  hoist$emit(glue(
+    "call dgels('N', {blas_int(m)}, {blas_int(n)}, {blas_int(nrhs)}, {A_work@name}, {blas_int(m)}, {B_work@name}, {blas_int(max_mn)}, {work@name}, {blas_int(lwork)}, {info@name})"
+  ))
+
+  expected_dims <- if (b_rank == 1L) list(n) else list(n, nrhs)
   writes_to_dest <- FALSE
   if (
     can_use_output(
       dest,
       input_names = c(A_name, B_input_name),
-      expected_dims = B@value@dims,
+      expected_dims = expected_dims,
       context = context,
       allow_alias = B_input_name
     )
@@ -702,18 +776,19 @@ lapack_solve <- function(
     out_name <- dest@name
     writes_to_dest <- TRUE
   } else {
-    out_var <- hoist$declare_tmp(mode = "double", dims = B@value@dims)
+    out_var <- hoist$declare_tmp(mode = "double", dims = expected_dims)
     out_name <- out_var@name
   }
-  hoist$emit(glue("{out_name} = {B_input_name}"))
 
-  ipiv <- hoist$declare_tmp(mode = "integer", dims = list(n))
-  info <- hoist$declare_tmp(mode = "integer", dims = NULL)
-  nrhs <- if (b_rank == 1L) 1L else dim_or_one(B, 2L)
-
-  hoist$emit(glue(
-    "call dgesv({blas_int(n)}, {blas_int(nrhs)}, {A_work@name}, {blas_int(n)}, {ipiv@name}, {out_name}, {blas_int(n)}, {info@name})"
-  ))
+  if (b_rank == 1L) {
+    if (passes_as_scalar(out_var)) {
+      hoist$emit(glue("{out_name} = {B_work@name}(1, 1)"))
+    } else {
+      hoist$emit(glue("{out_name} = {B_work@name}(1:{n_f}, 1)"))
+    }
+  } else {
+    hoist$emit(glue("{out_name} = {B_work@name}(1:{n_f}, 1:{nrhs_f})"))
+  }
 
   out <- Fortran(out_name, out_var)
   if (writes_to_dest) {

R/r2f-matrix-infer.R

@@ -180,21 +180,22 @@ infer_dest_solve <- function(args, scope) {
     return(NULL)
   }
   a_dims <- matrix_dims_var(A)
-  n <- a_dims$rows
+  n_rows <- a_dims$rows
+  n_cols <- a_dims$cols
 
   b_arg <- args$b %||% if (length(args) >= 2L) args[[2L]] else NULL
   if (is.null(b_arg)) {
-    return(Variable("double", list(n, n)))
+    return(Variable("double", list(n_rows, n_cols)))
   }
   B <- infer_symbol_var(b_arg, scope)
   if (is.null(B)) {
     return(NULL)
   }
   if (B@rank == 1L) {
-    return(Variable("double", list(n)))
+    return(Variable("double", list(n_cols)))
   }
   if (B@rank == 2L) {
-    return(Variable("double", list(n, B@dims[[2L]])))
+    return(Variable("double", list(n_cols, B@dims[[2L]])))
   }
   NULL
 }

R/r2f-matrix.R

@@ -571,6 +571,37 @@ register_r2f_handler(
   dest_infer = infer_dest_solve
 )
 
+register_r2f_handler(
+  "qr.solve",
+  function(args, scope, ..., hoist = NULL, dest = NULL) {
+    a_arg <- args$a %||% args[[1L]]
+    if (is.null(a_arg) || is_missing(a_arg)) {
+      stop("qr.solve() expects `a`", call. = FALSE)
+    }
+    if (!is.null(args$tol) && !is_missing(args$tol)) {
+      stop("qr.solve() does not support tol yet", call. = FALSE)
+    }
+
+    b_arg <- args$b %||% if (length(args) >= 2L) args[[2L]] else NULL
+    if (is.null(b_arg) || is_missing(b_arg)) {
+      stop("qr.solve() expects `b`", call. = FALSE)
+    }
+
+    A <- r2f(a_arg, scope, ..., hoist = hoist)
+    B <- r2f(b_arg, scope, ..., hoist = hoist)
+    lapack_solve(
+      A = A,
+      B = B,
+      scope = scope,
+      hoist = hoist,
+      dest = dest,
+      context = "qr.solve"
+    )
+  },
+  dest_supported = TRUE,
+  dest_infer = infer_dest_solve
+)
+
 register_r2f_handler(
   "chol",
   function(args, scope, ..., hoist = NULL, dest = NULL) {

R/r2f.R

@@ -752,10 +752,12 @@ r2f_handlers[["["]] <- function(
   # during symbol lowering as `(x/=0)`. When indexing, we must subscript the
   # underlying storage first, then convert the indexed value/section to logical.
   if (var@value@mode == "logical" && logical_as_int(var@value)) {
-    designator <- glue("{var@value@name}({str_flatten_commas(idxs)})")
+    base_name <- var@value@name %||% stop("missing array name for subscripting")
+    designator <- glue("{base_name}({str_flatten_commas(idxs)})")
     Fortran(glue("({designator} /= 0)"), outval)
   } else {
-    Fortran(glue("{var}({str_flatten_commas(idxs)})"), outval)
+    base_name <- var@value@name %||% stop("missing array name for subscripting")
+    Fortran(glue("{base_name}({str_flatten_commas(idxs)})"), outval)
   }
 }
 

R/subroutine.R

@@ -3,6 +3,7 @@ new_fortran_subroutine <- function(
   closure,
   parent = environment(closure)
 ) {
+  check_fortran_subroutine_name_valid(name)
   check_all_var_names_valid(closure)
 
   # translate body, and populate scope with variables
@@ -102,7 +103,7 @@ new_fortran_subroutine <- function(
   }
   subroutine <- glue(
     "
-subroutine {name}({str_flatten_commas(fsub_arg_names)}) bind(c)
+subroutine {name}({str_flatten_commas(fsub_arg_names)}) {fortran_bind_clause(name)}
   use iso_c_binding, only: {str_flatten_commas(used_iso_bindings)}
   implicit none
 
@@ -125,6 +126,43 @@ end subroutine
   )
 }
 
+check_fortran_subroutine_name_valid <- function(name) {
+  if (!is_string(name) || !nzchar(name)) {
+    stop("`name` must be a non-empty string.", call. = FALSE)
+  }
+  # Must be valid in:
+  # - Fortran: used as the procedure name.
+  # - C: used as the binding label and referenced from the C bridge.
+  if (!grepl("^[A-Za-z][A-Za-z0-9_]*$", name)) {
+    suggestion <- gsub("[^A-Za-z0-9_]", "_", name)
+    if (!grepl("^[A-Za-z]", suggestion)) {
+      suggestion <- paste0("quick_", suggestion)
+    }
+    stop(
+      "Invalid `quick()` name: '",
+      name,
+      "'. The name must match the pattern '^[A-Za-z][A-Za-z0-9_]*$' ",
+      "(letters, digits, underscore; starting with a letter). ",
+      "Suggested name: '",
+      suggestion,
+      "'.",
+      call. = FALSE
+    )
+  }
+  invisible(TRUE)
+}
+
+fortran_bind_clause <- function(name) {
+  # Fortran is case-insensitive, but C symbol names are case-sensitive.
+  # When `bind(c)` is used without an explicit `name=`, some toolchains
+  # canonicalize the external symbol name (e.g., lowercase), which can cause
+  # the C bridge to look up the wrong symbol for mixed-case procedure names.
+  if (identical(name, tolower(name))) {
+    return("bind(c)")
+  }
+  glue('bind(c, name = "{name}")')
+}
+
 insert_fortran_line_continuations <- function(
   code,
   preserve_attributes = TRUE