Preserve input mode in t() and diag()

R's t() and diag() preserve their input's type, but the handlers
unconditionally cast to double: t(m) and diag(m) on an integer matrix
returned doubles, and the constructor forms diag(x) / diag(x, nrow,
ncol) lost typeof(x) too. Only the identity forms (diag(n),
diag(nrow = n)) are double in R, which is what quickr already emits
for them.

Drop the maybe_cast_double() calls and carry the input mode through
the result Variable, the hoisted temporaries, and the zero fill in
diag_matrix(). can_use_output() gains a `mode` argument (default
"double"; all other callers unchanged) so an in-place destination is
only used when its declared mode matches, and infer_dest_diag()
reports the input's mode instead of hard-coding double -- a
double-inferred dest would have mislabeled an integer result's
declaration.

The transpose casts in unwrap_transpose_arg() stay, now with a
comment saying why: that path only feeds matrix products, which
always return double in R.

Author: mns-nordicals

R/r2f-matrix-blas.R

@@ -261,12 +261,13 @@ can_use_output <- function(
   input_names = character(),
   expected_dims = NULL,
   context,
-  allow_alias = character()
+  allow_alias = character(),
+  mode = "double"
 ) {
   if (is.null(dest)) {
     return(FALSE)
   }
-  if (!identical(dest@mode, "double")) {
+  if (!identical(dest@mode, mode)) {
     return(FALSE)
   }
   assert_dest_dims_compatible(dest, expected_dims, context)
@@ -1190,7 +1191,8 @@ 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)
@@ -1204,14 +1206,15 @@ diag_extract <- function(x, scope, hoist, dest = NULL, context = "diag") {
       dest,
       input_names = x_name,
       expected_dims = list(diag_len),
-      context = context
+      context = context,
+      mode = x@value@mode
     )
   ) {
     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))
     out_name <- out_var@name
   }
 
@@ -1241,9 +1244,20 @@ 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
+  zero <- switch(
+    mode,
+    double = "0.0_c_double",
+    integer = "0_c_int",
+    logical = ".false.",
+    complex = "(0.0_c_double, 0.0_c_double)",
+    stop(context, " does not support mode ", mode, call. = FALSE)
+  )
+
   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 +1270,19 @@ diag_matrix <- function(
       dest,
       input_names = x_name,
       expected_dims = list(nrow, ncol),
-      context = context
+      context = context,
+      mode = mode
     )
   ) {
     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))
     out_name <- out_var@name
   }
 
-  hoist$emit(glue("{out_name} = 0.0_c_double"))
+  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,14 @@ 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(x@mode, list(diag_len)))
   }
 
   # Case: x is a scalar literal (identity matrix of that size)
@@ -330,7 +335,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 +349,12 @@ infer_dest_diag <- function(args, scope) {
       if (is.null(ncol)) {
         ncol <- nrow
       }
-      return(Variable("double", list(nrow, ncol)))
+      return(Variable(x@mode, list(nrow, ncol)))
     }
     # 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(x@mode, list(len, len)))
     }
   }
 

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)

tests/testthat/_snaps/matrix.md

@@ -231,3 +231,79 @@
         return out;
       }
 
+# t() and diag() preserve integer mode
+
+    Code
+      fn
+    Output
+      function(m) {
+          declare(type(m = integer(3, 3)))
+          out <- diag(m)
+          out
+        }
+      <environment: 0x0>
+    Code
+      cat(fsub)
+    Output
+      subroutine fn(m, out) bind(c)
+        use iso_c_binding, only: c_int
+        implicit none
+      
+        ! manifest start
+        ! args
+        integer(c_int), intent(in) :: m(3, 3)
+        integer(c_int), intent(out) :: out(3)
+        ! manifest end
+      
+      
+        block
+          integer(c_int) :: btmp1_
+      
+          do btmp1_ = 1_c_int, int(3, kind=c_int)
+            out(btmp1_) = m(btmp1_, btmp1_)
+          end do
+        end block
+      end subroutine
+    Code
+      cat(cwrapper)
+    Output
+      #define R_NO_REMAP
+      #include <R.h>
+      #include <Rinternals.h>
+      
+      
+      extern void fn(const int* const m__, int* const out__);
+      
+      SEXP fn_(SEXP _args) {
+        // m
+        _args = CDR(_args);
+        SEXP m = CAR(_args);
+        if (TYPEOF(m) != INTSXP) {
+          Rf_error("typeof(m) must be 'integer', not '%s'", Rf_type2char(TYPEOF(m)));
+        }
+        const int* const m__ = INTEGER(m);
+        const int* const m__dim_ = ({
+        SEXP dim_ = Rf_getAttrib(m, R_DimSymbol);
+        if (Rf_length(dim_) != 2) Rf_error(
+          "m must be a 2D-array, but length(dim(m)) is %i",
+          (int) Rf_length(dim_));
+        INTEGER(dim_);});
+        const int m__dim_1_ = m__dim_[0];
+        const int m__dim_2_ = m__dim_[1];
+        
+        if (m__dim_1_ != 3)
+          Rf_error("dim(m)[1] must be 3, not %0.f",
+                    (double)m__dim_1_);
+        if (m__dim_2_ != 3)
+          Rf_error("dim(m)[2] must be 3, not %0.f",
+                    (double)m__dim_2_);
+        const R_xlen_t out__len_ = 3;
+        SEXP out = PROTECT(Rf_allocVector(INTSXP, out__len_));
+        int* out__ = INTEGER(out);
+        
+        fn(m__, out__);
+        
+        UNPROTECT(1);
+        return out;
+      }
+

tests/testthat/test-matrix.R

@@ -298,3 +298,68 @@ test_that("indexing function like transposed expressions hoists temporaries that
   x <- matrix(runif(25), 5, 5)
   expect_quick_identical(fn, list(x = x))
 })
+
+test_that("t() and diag() preserve integer mode", {
+  tfn <- function(m) {
+    declare(type(m = integer(2, 3)))
+    t(m)
+  }
+  expect_quick_equal(tfn, list(matrix(1:6, 2, 3))) # typeof integer
+
+  tvec <- function(x) {
+    declare(type(x = integer(3)))
+    t(x)
+  }
+  expect_quick_equal(tvec, list(1:3)) # R: 1 x 3 integer matrix
+
+  dfn <- function(m) {
+    declare(type(m = integer(3, 3)))
+    diag(m)
+  }
+  expect_quick_equal(dfn, list(matrix(1:9, 3, 3))) # R: integer vector
+
+  # same, through the inferred-destination path (out <- diag(m))
+  dfn2 <- function(m) {
+    declare(type(m = integer(3, 3)))
+    out <- diag(m)
+    out
+  }
+  expect_translation_snapshots(dfn2)
+  expect_quick_equal(dfn2, list(matrix(1:9, 3, 3)))
+
+  dvec <- function(x) {
+    declare(type(x = integer(3)))
+    diag(x)
+  }
+  expect_quick_equal(dvec, list(1:3)) # R: integer matrix
+
+  # x recycled along the diagonal of a non-square result
+  drect <- function(x) {
+    declare(type(x = integer(2)))
+    diag(x, 3L, 4L)
+  }
+  expect_quick_equal(drect, list(1:2))
+
+  # identity forms stay double, as in R
+  dident <- function() {
+    out <- diag(3L)
+    out
+  }
+  expect_quick_equal(dident, list())
+})
+
+test_that("t() and diag() preserve logical mode", {
+  m <- matrix(c(TRUE, FALSE, TRUE, TRUE), 2, 2)
+
+  tfn <- function(m) {
+    declare(type(m = logical(2, 2)))
+    t(m)
+  }
+  expect_quick_equal(tfn, list(m))
+
+  dfn <- function(m) {
+    declare(type(m = logical(2, 2)))
+    diag(m)
+  }
+  expect_quick_equal(dfn, list(m))
+})