Avoid clone in vec_cast() with shaped objects when nothing changes (#2109)

* Only clone with shaped casting if dims/dim names change

So the self cast remains free

* Add golden tests

* NEWS bullet

Author: DavisVaughan

NEWS.md

@@ -1,5 +1,7 @@
 # vctrs (development version)
 
+* `vec_cast()` with arrays no longer clones when no casting is required (#2006).
+
 * `vec_rank()` now throws an improved error on non-vector types, like `NULL` (#1967).
 
 * `vec_ptype_common()` has gained a `.finalise` argument that defaults to `TRUE`. Setting this to `FALSE` lets you opt out of prototype finalisation, which allows `vec_ptype_common()` to act like `vec_ptype()` and `vec_ptype2()`, which don't finalise. This can be useful in some advanced common type determination cases (#2100).

src/cast.c

@@ -20,17 +20,18 @@ r_obj* ffi_cast(r_obj* x,
 r_obj* vec_cast_opts(const struct cast_opts* opts) {
   r_obj* x = opts->x;
   r_obj* to = opts->to;
-  struct vctrs_arg* x_arg = opts->p_x_arg;
-  struct vctrs_arg* to_arg = opts->p_to_arg;
+  struct vctrs_arg* p_x_arg = opts->p_x_arg;
+  struct vctrs_arg* p_to_arg = opts->p_to_arg;
+  struct r_lazy call = opts->call;
 
   if (x == r_null) {
     // Allow both `vec_cast(NULL, <vector>)` and `vec_cast(NULL, NULL)`
-    obj_check_vector(to, VCTRS_ALLOW_NULL_yes, to_arg, opts->call);
+    obj_check_vector(to, VCTRS_ALLOW_NULL_yes, p_to_arg, call);
     return x;
   }
   if (to == r_null) {
     // Allow `vec_cast(<vector>, NULL)`
-    obj_check_vector(x, VCTRS_ALLOW_NULL_no, x_arg, opts->call);
+    obj_check_vector(x, VCTRS_ALLOW_NULL_no, p_x_arg, call);
     return x;
   }
 
@@ -42,10 +43,10 @@ r_obj* vec_cast_opts(const struct cast_opts* opts) {
   }
 
   if (x_type == VCTRS_TYPE_scalar) {
-    stop_scalar_type(x, x_arg, opts->call);
+    stop_scalar_type(x, p_x_arg, call);
   }
   if (to_type == VCTRS_TYPE_scalar) {
-    stop_scalar_type(to, to_arg, opts->call);
+    stop_scalar_type(to, p_to_arg, call);
   }
 
   r_obj* out = r_null;
@@ -64,7 +65,26 @@ r_obj* vec_cast_opts(const struct cast_opts* opts) {
   }
 
   if (has_dim(x) || has_dim(to)) {
-    out = vec_shape_broadcast(out, opts);
+    r_obj* x_dim = r_dim(x);
+    r_obj* x_dim_names = r_dim_names(x);
+
+    // Ensure `out` has the shape of `x`.
+    // Native casting doesn't propagate shape.
+    if (
+      !equal_object(r_dim(out), x_dim) ||
+      !equal_object(r_dim_names(out), x_dim_names)
+    ) {
+      out = KEEP(r_clone_referenced(out));
+      r_attrib_poke_dim(out, x_dim);
+      r_attrib_poke_dim_names(out, x_dim_names);
+      FREE(1);
+    }
+    KEEP(out);
+
+    // Broadcast `out` to the shape of `to`
+    out = vec_shape_broadcast(out, to, p_x_arg, p_to_arg, call);
+
+    FREE(1);
   }
 
   FREE(1);

src/shape.c

@@ -186,30 +186,31 @@ int dim2(
 
 // -----------------------------------------------------------------------------
 
-r_obj* vec_shape_broadcast(r_obj* out, const struct cast_opts* p_opts) {
-  r_obj* r_x_arg = KEEP(vctrs_arg(p_opts->p_x_arg));
-  r_obj* r_to_arg = KEEP(vctrs_arg(p_opts->p_to_arg));
-  r_obj* call = KEEP(r_lazy_eval(p_opts->call));
-
-  out = KEEP(r_clone_referenced(out));
-
-  r_attrib_poke_dim(out, r_dim(p_opts->x));
-  r_attrib_poke_dim_names(out, r_dim_names(p_opts->x));
+r_obj* vec_shape_broadcast(
+  r_obj* x,
+  r_obj* to,
+  struct vctrs_arg* p_x_arg,
+  struct vctrs_arg* p_to_arg,
+  struct r_lazy call
+) {
+  r_obj* ffi_x_arg = KEEP(vctrs_arg(p_x_arg));
+  r_obj* ffi_to_arg = KEEP(vctrs_arg(p_to_arg));
+  r_obj* ffi_call = KEEP(r_lazy_eval(call));
 
-  out = vctrs_eval_mask5(
+  r_obj* out = vctrs_eval_mask5(
     r_sym("shape_broadcast"),
     r_syms.x,
-    out,
+    x,
     r_sym("to"),
-    p_opts->to,
+    to,
     syms.x_arg,
-    r_x_arg,
+    ffi_x_arg,
     syms.to_arg,
-    r_to_arg,
+    ffi_to_arg,
     r_syms.call,
-    call
+    ffi_call
   );
 
-  FREE(4);
+  FREE(3);
   return out;
 }

src/shape.h

@@ -2,7 +2,6 @@
 #define VCTRS_SHAPE_H
 
 #include "vctrs-core.h"
-#include "cast.h"
 
 // Computes the common shape of `x` and `y` and attaches it as the
 // dimensions of `ptype`. If `x` and `y` are both atomic with `NULL` dimensions,
@@ -15,7 +14,12 @@ r_obj* vec_shaped_ptype(
   struct vctrs_arg* p_y_arg
 );
 
-r_obj* vec_shape_broadcast(r_obj* out, const struct cast_opts* p_opts);
-
+r_obj* vec_shape_broadcast(
+  r_obj* x,
+  r_obj* to,
+  struct vctrs_arg* p_x_arg,
+  struct vctrs_arg* p_to_arg,
+  struct r_lazy call
+);
 
 #endif

tests/testthat/_snaps/cast.md

@@ -91,3 +91,17 @@
       Error:
       ! Can't convert `TRUE` <logical> to <vctrs_unspecified>.
 
+# casting performs expected allocations
+
+    Code
+      x <- matrix(rep(1L, 100), ncol = 2)
+      with_memory_prof(vec_cast(x, x))
+    Output
+      [1] 0B
+    Code
+      x <- matrix(rep(1L, 100), ncol = 2)
+      y <- matrix(rep(1, 100), ncol = 2)
+      with_memory_prof(vec_cast(x, y))
+    Output
+      [1] 848B
+

tests/testthat/test-cast.R

@@ -357,3 +357,18 @@ test_that("can cast to unspecified `NA` with `vec_cast()` and `vec_cast_common()
   # so this technically works (but again, it is an edge case)
   expect_identical(vec_cast_common(TRUE, .to = unspecified(1)), list(TRUE))
 })
+
+# Golden tests -------------------------------------------------------
+
+test_that("casting performs expected allocations", {
+  expect_snapshot({
+    # No allocations when shape doesn't change (#2006)
+    x <- matrix(rep(1L, 1e2), ncol = 2)
+    with_memory_prof(vec_cast(x, x))
+
+    # One allocation when type changes
+    x <- matrix(rep(1L, 1e2), ncol = 2)
+    y <- matrix(rep(1, 1e2), ncol = 2)
+    with_memory_prof(vec_cast(x, y))
+  })
+})