Preserve logical storage through diag()

Author: t-kalinowski

R/r2f-matrix-blas.R

@@ -262,14 +262,22 @@ can_use_output <- function(
   expected_dims = NULL,
   context,
   allow_alias = character(),
-  mode = "double"
+  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, mode)) {
     return(FALSE)
   }
+  if (!identical(logical_as_int(dest), logical_is_c_int)) {
+    return(FALSE)
+  }
   assert_dest_dims_compatible(dest, expected_dims, context)
   output_name <- dest@name
   if (is.null(output_name) || !nzchar(output_name)) {
@@ -293,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
@@ -1199,6 +1208,7 @@ diag_extract <- function(x, scope, hoist, dest = NULL, context = "diag") {
   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 (
@@ -1207,14 +1217,19 @@ diag_extract <- function(x, scope, hoist, dest = NULL, context = "diag") {
       input_names = x_name,
       expected_dims = list(diag_len),
       context = context,
-      mode = x@value@mode
+      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 = x@value@mode, 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
   }
 
@@ -1249,14 +1264,7 @@ diag_matrix <- function(
   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)
-  )
+  logical_is_c_int <- logical_as_int(x@value)
 
   diag_len <- diag_length_expr(nrow, ncol, context)
   x_scalar <- passes_as_scalar(x@value)
@@ -1271,17 +1279,30 @@ diag_matrix <- function(
       input_names = x_name,
       expected_dims = list(nrow, ncol),
       context = context,
-      mode = mode
+      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 = mode, 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
   }
 
+  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)

R/r2f-matrix-infer.R

@@ -318,7 +318,11 @@ infer_dest_diag <- function(args, scope) {
     diag_len <- diag_length_expr(x_dims$rows, x_dims$cols, "diag")
     # diag_extract() preserves x's mode; a double-inferred dest would
     # mislabel an integer diagonal.
-    return(Variable(x@mode, list(diag_len)))
+    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)
@@ -349,12 +353,20 @@ infer_dest_diag <- function(args, scope) {
       if (is.null(ncol)) {
         ncol <- nrow
       }
-      return(Variable(x@mode, 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(x@mode, list(len, len)))
+      return(Variable(
+        mode = x@mode,
+        dims = list(len, len),
+        logical_as_int = logical_as_int(x)
+      ))
     }
   }
 

tests/testthat/_snaps/matrix.md

@@ -307,3 +307,156 @@
         return out;
       }
 
+# diag() preserves integer-backed logical storage in an intermediate
+
+    Code
+      fn
+    Output
+      function(m) {
+          declare(type(m = logical(2, 2)))
+          d <- diag(m)
+          as.integer(d)
+        }
+      <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(2, 2) ! logical
+        integer(c_int), intent(out) :: out_(2)
+      
+        ! locals
+        integer(c_int) :: d(2) ! logical
+        ! manifest end
+      
+      
+        block
+          integer(c_int) :: btmp1_
+      
+          do btmp1_ = 1_c_int, int(2, kind=c_int)
+            d(btmp1_) = m(btmp1_, btmp1_)
+          end do
+        end block
+        out_ = d
+      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) != LGLSXP) {
+          Rf_error("typeof(m) must be 'logical', not '%s'", Rf_type2char(TYPEOF(m)));
+        }
+        const int* const m__ = LOGICAL(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_ != 2)
+          Rf_error("dim(m)[1] must be 2, not %0.f",
+                    (double)m__dim_1_);
+        if (m__dim_2_ != 2)
+          Rf_error("dim(m)[2] must be 2, not %0.f",
+                    (double)m__dim_2_);
+        const R_xlen_t out___len_ = 2;
+        SEXP out_ = PROTECT(Rf_allocVector(INTSXP, out___len_));
+        int* out___ = INTEGER(out_);
+        
+        fn(m__, out___);
+        
+        UNPROTECT(1);
+        return out_;
+      }
+
+# diag() initializes integer-backed logical outputs as integers
+
+    Code
+      fn
+    Output
+      function(x) {
+          declare(type(x = logical(2)))
+          diag(x, 3L, 4L)
+        }
+      <environment: 0x0>
+    Code
+      cat(fsub)
+    Output
+      subroutine fn(x, out_) bind(c)
+        use iso_c_binding, only: c_int
+        implicit none
+      
+        ! manifest start
+        ! args
+        integer(c_int), intent(in) :: x(2) ! logical
+        integer(c_int), intent(out) :: out_(3, 4) ! logical
+        ! manifest end
+      
+      
+        block
+          integer(c_int) :: btmp1_
+      
+          out_ = 0_c_int
+          do btmp1_ = 1_c_int, int(3, kind=c_int)
+            out_(btmp1_, btmp1_) = x(1_c_int + mod(btmp1_ - 1_c_int, int(2, kind=c_int)))
+          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 x__, int* const out___);
+      
+      SEXP fn_(SEXP _args) {
+        // x
+        _args = CDR(_args);
+        SEXP x = CAR(_args);
+        if (TYPEOF(x) != LGLSXP) {
+          Rf_error("typeof(x) must be 'logical', not '%s'", Rf_type2char(TYPEOF(x)));
+        }
+        const int* const x__ = LOGICAL(x);
+        const R_xlen_t x__len_ = Rf_xlength(x);
+        
+        if (x__len_ != 2)
+          Rf_error("length(x) must be 2, not %0.f",
+                    (double)x__len_);
+        const R_xlen_t out___len_ = (3) * (4);
+        SEXP out_ = PROTECT(Rf_allocVector(LGLSXP, out___len_));
+        int* out___ = LOGICAL(out_);
+        {
+          const SEXP _dim_sexp = PROTECT(Rf_allocVector(INTSXP, 2));
+          int* const _dim = INTEGER(_dim_sexp);
+          _dim[0] = 3;
+          _dim[1] = 4;
+          Rf_dimgets(out_, _dim_sexp);
+        }
+        
+        fn(x__, out___);
+        
+        UNPROTECT(2);
+        return out_;
+      }
+

tests/testthat/test-matrix.R

@@ -363,3 +363,31 @@ test_that("t() and diag() preserve logical mode", {
   }
   expect_quick_equal(dfn, list(m))
 })
+
+test_that("diag() preserves integer-backed logical storage in an intermediate", {
+  fn <- function(m) {
+    declare(type(m = logical(2, 2)))
+    d <- diag(m)
+    as.integer(d)
+  }
+
+  m <- matrix(c(TRUE, FALSE, TRUE, FALSE), 2, 2)
+  expect_match(
+    as.character(r2f(fn)),
+    "integer(c_int) :: d(2)",
+    fixed = TRUE
+  )
+  expect_translation_snapshots(fn)
+  expect_quick_identical(fn, list(m))
+})
+
+test_that("diag() initializes integer-backed logical outputs as integers", {
+  fn <- function(x) {
+    declare(type(x = logical(2)))
+    diag(x, 3L, 4L)
+  }
+
+  expect_match(as.character(r2f(fn)), "out_ = 0_c_int", fixed = TRUE)
+  expect_translation_snapshots(fn)
+  expect_quick_identical(fn, list(c(TRUE, FALSE)))
+})