[4/15] Fix result modes in ifelse(), t(), and diag(); enforce ifelse() branch shapes

R/r2f-conditionals.R

@@ -1,13 +1,108 @@
 # r2f-conditionals.R
 # Handlers for vectorized conditionals: ifelse
 
+# --- Local Helpers ---
+
+ifelse_branch_shape_msg <- paste0(
+  "ifelse() `yes` and `no` must be scalars or match the shape of `test`; ",
+  "R-style recycling is not supported"
+)
+
+# Three-valued conformability verdict for one axis of an ifelse() branch
+# against `test`: ok+known (no guard), not-ok+known (compile error), or
+# unknown (runtime guard). NA dims are always unknown: two unknown lengths
+# are not the same quantity.
+ifelse_axis_verdict <- function(test_dim, branch_dim) {
+  if (is_wholenumber(test_dim) && is_wholenumber(branch_dim)) {
+    return(list(
+      ok = identical(as.integer(test_dim), as.integer(branch_dim)),
+      unknown = FALSE
+    ))
+  }
+  if (!is_scalar_na(test_dim) && !is_scalar_na(branch_dim)) {
+    test_norm <- fortranize_expr_symbols(test_dim)
+    branch_norm <- fortranize_expr_symbols(branch_dim)
+    if (identical(test_norm, branch_norm)) {
+      return(list(ok = TRUE, unknown = FALSE))
+    }
+  }
+  list(ok = TRUE, unknown = TRUE)
+}
+
+# Enforce the shape contract for one ifelse() branch: scalars broadcast
+# natively; a non-scalar branch must match `test`'s shape, because
+# merge() requires conformable arguments and a runtime mismatch would
+# read past the shorter branch. Statically unequal dims are a compile
+# error; symbolic dims get a statement-level runtime size guard, emitted into
+# `hoist` -- always a live hoist context, since r2f() substitutes a fresh one
+# before dispatching to any handler.
+check_ifelse_branch_shape <- function(branch, mask, hoist, scope) {
+  if (passes_as_scalar(branch@value)) {
+    return(invisible())
+  }
+  if (branch@value@rank != mask@value@rank) {
+    stop(ifelse_branch_shape_msg, call. = FALSE)
+  }
+  unknown_axes <- integer()
+  for (axis in seq_len(mask@value@rank)) {
+    verdict <- ifelse_axis_verdict(
+      dim_or_one(mask, axis),
+      dim_or_one(branch, axis)
+    )
+    if (!verdict$ok) {
+      stop(ifelse_branch_shape_msg, call. = FALSE)
+    }
+    if (verdict$unknown) {
+      unknown_axes <- c(unknown_axes, axis)
+    }
+  }
+  if (!length(unknown_axes)) {
+    return(invisible())
+  }
+  # size() is an inquiry, so applying it to operand expression text does
+  # not evaluate the operands.
+  condition <- str_flatten(
+    map_chr(
+      unknown_axes,
+      function(axis) glue("size({branch}, {axis}) /= size({mask}, {axis})")
+    ),
+    " .or. "
+  )
+  emit_quickr_error_if(condition, ifelse_branch_shape_msg, hoist, scope)
+  invisible()
+}
+
 # --- Handlers ---
 
-r2f_handlers[["ifelse"]] <- function(args, scope, ...) {
-  .[mask, tsource, fsource] <- lapply(args, r2f, scope, ...)
+r2f_handlers[["ifelse"]] <- function(args, scope, ..., hoist = NULL) {
+  .[mask, tsource, fsource] <- lapply(args, r2f, scope, ..., hoist = hoist)
+
+  # R: the result is shaped like `test` (branches only contribute values).
+  # A scalar test with array branches is not representable with merge().
+  if (
+    passes_as_scalar(mask@value) &&
+      !(passes_as_scalar(tsource@value) && passes_as_scalar(fsource@value))
+  ) {
+    stop(
+      "ifelse() result takes the shape of `test`; ",
+      "array-valued yes/no with scalar test is not supported",
+      call. = FALSE
+    )
+  }
+
+  # Checked before casts so guards splice the bare operand text.
+  check_ifelse_branch_shape(tsource, mask, hoist, scope)
+  check_ifelse_branch_shape(fsource, mask, hoist, scope)
+
   mask <- booleanize_logical_as_int(mask)
-  # (tsource, fsource, mask)
-  mode <- tsource@value@mode
-  dims <- conform(mask@value, tsource@value, fsource@value)@dims
-  Fortran(glue("merge({tsource}, {fsource}, {mask})"), Variable(mode, dims))
+
+  # merge() requires same-typed branches; promote both to their common mode.
+  promoted <- promote_operands(list(tsource, fsource), context = "ifelse()")
+  .[tsource, fsource] <- promoted$args
+  mode <- promoted$mode
+
+  Fortran(
+    glue("merge({tsource}, {fsource}, {mask})"),
+    Variable(mode, mask@value@dims)
+  )
 }

R/r2f-matrix-blas.R

@@ -261,12 +261,21 @@ can_use_output <- function(
   input_names = character(),
   expected_dims = NULL,
   context,
-  allow_alias = character()
+  allow_alias = character(),
+  mode = "double",
+  logical_is_c_int = FALSE
 ) {
+  stopifnot(
+    is_bool(logical_is_c_int),
+    !logical_is_c_int || identical(mode, "logical")
+  )
   if (is.null(dest)) {
     return(FALSE)
   }
-  if (!identical(dest@mode, "double")) {
+  if (!identical(dest@mode, mode)) {
+    return(FALSE)
+  }
+  if (!identical(logical_as_int(dest), logical_is_c_int)) {
     return(FALSE)
   }
   assert_dest_dims_compatible(dest, expected_dims, context)
@@ -292,7 +301,8 @@ ensure_blas_operand_name <- function(x, hoist) {
   }
   tmp <- hoist$declare_tmp(
     mode = x@value@mode %||% "double",
-    dims = x@value@dims
+    dims = x@value@dims,
+    logical_as_int = logical_as_int(x@value)
   )
   hoist$emit(glue("{tmp@name} = {x}"))
   tmp@name
@@ -1190,28 +1200,36 @@ lapack_chol2inv <- function(
 diag_extract <- function(x, scope, hoist, dest = NULL, context = "diag") {
   assert_hoist_env(hoist)
 
-  x <- maybe_cast_double(x)
+  # R's diag(<matrix>) preserves the input mode; the copy loop is
+  # mode-agnostic.
   assert_rank2_matrix(x, paste0(context, " expects a matrix input"))
 
   x_dims <- matrix_dims(x)
   diag_len <- diag_length_expr(x_dims$rows, x_dims$cols, context)
 
   x_name <- ensure_blas_operand_name(x, hoist)
+  logical_is_c_int <- logical_as_int(x@value)
 
   writes_to_dest <- FALSE
   if (
     can_use_output(
       dest,
       input_names = x_name,
       expected_dims = list(diag_len),
-      context = context
+      context = context,
+      mode = x@value@mode,
+      logical_is_c_int = logical_is_c_int
     )
   ) {
     out_var <- dest
     out_name <- dest@name
     writes_to_dest <- TRUE
   } else {
-    out_var <- hoist$declare_tmp(mode = "double", dims = list(diag_len))
+    out_var <- hoist$declare_tmp(
+      mode = x@value@mode,
+      dims = list(diag_len),
+      logical_as_int = logical_is_c_int
+    )
     out_name <- out_var@name
   }
 
@@ -1241,9 +1259,13 @@ diag_matrix <- function(
 ) {
   assert_hoist_env(hoist)
 
-  x <- maybe_cast_double(x)
+  # R's diag(x, ...) preserves typeof(x). The identity-matrix callers pass
+  # a synthesized 1.0_c_double, which keeps diag(n) double, as in R.
   assert_rank_leq1(x, paste0(context, " expects a vector or scalar input"))
 
+  mode <- x@value@mode
+  logical_is_c_int <- logical_as_int(x@value)
+
   diag_len <- diag_length_expr(nrow, ncol, context)
   x_scalar <- passes_as_scalar(x@value)
   x_len <- if (x_scalar) 1L else dim_or_one(x, 1L)
@@ -1256,18 +1278,32 @@ diag_matrix <- function(
       dest,
       input_names = x_name,
       expected_dims = list(nrow, ncol),
-      context = context
+      context = context,
+      mode = mode,
+      logical_is_c_int = logical_is_c_int
     )
   ) {
     out_var <- dest
     out_name <- dest@name
     writes_to_dest <- TRUE
   } else {
-    out_var <- hoist$declare_tmp(mode = "double", dims = list(nrow, ncol))
+    out_var <- hoist$declare_tmp(
+      mode = mode,
+      dims = list(nrow, ncol),
+      logical_as_int = logical_is_c_int
+    )
     out_name <- out_var@name
   }
 
-  hoist$emit(glue("{out_name} = 0.0_c_double"))
+  zero <- switch(
+    mode,
+    double = "0.0_c_double",
+    integer = "0_c_int",
+    logical = if (logical_as_int(out_var)) "0_c_int" else ".false.",
+    complex = "(0.0_c_double, 0.0_c_double)",
+    stop(context, " does not support mode ", mode, call. = FALSE)
+  )
+  hoist$emit(glue("{out_name} = {zero}"))
 
   idx_i <- hoist$declare_tmp(mode = "integer", dims = NULL)
   value_expr <- if (x_scalar) {

R/r2f-matrix-infer.R

@@ -311,9 +311,18 @@ infer_dest_diag <- function(args, scope) {
 
   # Case: x is a matrix -> extract diagonal (returns vector)
   if (!is.null(x) && x@rank == 2L) {
+    if (is.null(x@mode)) {
+      return(NULL)
+    }
     x_dims <- matrix_dims_var(x)
     diag_len <- diag_length_expr(x_dims$rows, x_dims$cols, "diag")
-    return(Variable("double", list(diag_len)))
+    # diag_extract() preserves x's mode; a double-inferred dest would
+    # mislabel an integer diagonal.
+    return(Variable(
+      mode = x@mode,
+      dims = list(diag_len),
+      logical_as_int = logical_as_int(x)
+    ))
   }
 
   # Case: x is a scalar literal (identity matrix of that size)
@@ -330,7 +339,8 @@ infer_dest_diag <- function(args, scope) {
   }
 
   # Case: x is a vector or scalar, construct diagonal matrix
-  if (!is.null(x) && x@rank <= 1L) {
+  # (diag_matrix() preserves x's mode, matching R)
+  if (!is.null(x) && x@rank <= 1L && !is.null(x@mode)) {
     if (has_nrow || has_ncol) {
       nrow <- if (has_nrow) infer_size(nrow_arg, scope) else NULL
       ncol <- if (has_ncol) infer_size(ncol_arg, scope) else NULL
@@ -343,12 +353,20 @@ infer_dest_diag <- function(args, scope) {
       if (is.null(ncol)) {
         ncol <- nrow
       }
-      return(Variable("double", list(nrow, ncol)))
+      return(Variable(
+        mode = x@mode,
+        dims = list(nrow, ncol),
+        logical_as_int = logical_as_int(x)
+      ))
     }
     # No nrow/ncol: square matrix from vector length
     if (x@rank == 1L) {
       len <- var_dim_or_one(x, 1L)
-      return(Variable("double", list(len, len)))
+      return(Variable(
+        mode = x@mode,
+        dims = list(len, len),
+        logical_as_int = logical_as_int(x)
+      ))
     }
   }
 

R/r2f-matrix-parse.R

@@ -1,6 +1,9 @@
 # Matrix parsing helpers
 
 # Unwrap t() calls to infer transpose flags and normalize scalars/vectors.
+# The double casts here are correct and intentional: this path only feeds
+# matrix-multiplication handlers (%*%, crossprod, ...), and R's matrix
+# products always return double. The standalone t() handler preserves mode.
 unwrap_transpose_arg <- function(arg, scope, ..., hoist) {
   arg_unwrapped <- unwrap_parens(arg)
   if (is_call(arg_unwrapped, quote(t)) && length(arg_unwrapped) == 2L) {

R/r2f-matrix.R

@@ -137,13 +137,14 @@ register_r2f_handler(
 r2f_handlers[["t"]] <- function(args, scope, ..., hoist = NULL) {
   stopifnot(length(args) == 1L)
   x <- r2f(args[[1L]], scope, ..., hoist = hoist)
-  x <- maybe_cast_double(x)
+  # R's t() preserves the input mode. (The transposes feeding matrix
+  # multiplication go through unwrap_transpose_arg(), not this handler.)
   if (x@value@rank == 2) {
-    val <- Variable("double", list(x@value@dims[[2]], x@value@dims[[1]]))
+    val <- Variable(x@value@mode, list(x@value@dims[[2]], x@value@dims[[1]]))
     return(Fortran(glue("transpose({x})"), val))
   } else if (x@value@rank == 1) {
     len <- x@value@dims[[1]]
-    val <- Variable("double", list(1L, len))
+    val <- Variable(x@value@mode, list(1L, len))
     return(Fortran(glue("reshape({x}, [1, int({len})])"), val))
   } else if (x@value@rank == 0) {
     return(x)