Add ZM set/drop trans

NAMESPACE

@@ -311,6 +311,8 @@ export(wk_crs_inherit)
 export(wk_crs_output)
 export(wk_debug)
 export(wk_debug_filter)
+export(wk_drop_m)
+export(wk_drop_z)
 export(wk_flatten)
 export(wk_flatten_filter)
 export(wk_format)
@@ -331,8 +333,11 @@ export(wk_problems)
 export(wk_problems_handler)
 export(wk_restore)
 export(wk_set_crs)
+export(wk_set_m)
+export(wk_set_z)
 export(wk_trans_affine)
 export(wk_trans_inverse)
+export(wk_trans_set)
 export(wk_transform)
 export(wk_transform_filter)
 export(wk_translate)

R/set.R

@@ -0,0 +1,60 @@
+
+#' Set coordinate values
+#'
+#' @inheritParams wk_handle
+#' @param z,m A vector of Z or M values applied feature-wise and recycled
+#'    along `handleable`. Use `NA` to keep the existing value of a given
+#'    feature.
+#' @param value An [xy()], [xyz()], [xym()], or [xyzm()] of coordinates
+#'   used to replace values in the input. Use `NA` to keep the existing
+#'   value.
+#' @param use_z,use_m Used to declare the output type. Use `TRUE` to
+#'   ensure the output has that dimension, `FALSE` to ensure it does not,
+#'   and `NA` to leave the dimension unchanged.
+#' @param
+#'
+#' @export
+#'
+#' @examples
+#' wk_set_z(wkt("POINT (0 1)"), 2)
+#' wk_set_m(wkt("POINT (0 1)"), 2)
+#' wk_drop_z(wkt("POINT ZM (0 1 2 3)"))
+#' wk_drop_m(wkt("POINT ZM (0 1 2 3)"))
+#'
+wk_set_z <- function(handleable, z, ...) {
+  wk_set_base(handleable, wk_trans_set(xyz(NA, NA, z), use_z = TRUE), ...)
+}
+
+#' @rdname wk_set_z
+#' @export
+wk_set_m <- function(handleable, m, ...) {
+  wk_set_base(handleable, wk_trans_set(xym(NA, NA, m), use_m = TRUE), ...)
+}
+
+#' @rdname wk_set_z
+#' @export
+wk_drop_z <- function(handleable, ...) {
+  wk_set_base(handleable, wk_trans_set(xy(NA, NA), use_z = FALSE), ...)
+}
+
+#' @rdname wk_set_z
+#' @export
+wk_drop_m <- function(handleable, ...) {
+  wk_set_base(handleable, wk_trans_set(xy(NA, NA), use_m = FALSE), ...)
+}
+
+#' @rdname wk_set_z
+#' @export
+wk_trans_set <- function(value, use_z = NA, use_m = NA) {
+  value <- as_xy(value)
+  value <- as_xy(value, dims = c("x", "y", "z", "m"))
+  new_wk_trans(
+    .Call(wk_c_trans_set_new, value, as.logical(use_z)[1], as.logical(use_m)[1]),
+    "wk_trans_set"
+  )
+}
+
+wk_set_base <- function(handleable, trans, ...) {
+  result <- wk_handle(handleable, wk_transform_filter(wk_writer(handleable), trans), ...)
+  wk_set_crs(wk_restore(handleable, result), wk_crs(handleable))
+}

inst/include/wk-v1-impl.c

@@ -151,7 +151,8 @@ wk_trans_t* wk_trans_create() {
     Rf_error("Failed to alloc wk_trans_t*"); // # nocov
   }
 
-  trans->flags_out = WK_FLAG_DIMS_UNKNOWN;
+  trans->use_z = NA_INTEGER;
+  trans->use_m = NA_INTEGER;
 
   trans->xyzm_out_min[0] = R_NegInf;
   trans->xyzm_out_min[1] = R_NegInf;

inst/include/wk-v1.h

@@ -83,7 +83,8 @@ typedef struct {
 
 typedef struct {
     void* trans_data;
-    uint32_t flags_out;
+    int use_z;
+    int use_m;
     double xyzm_out_min[4];
     double xyzm_out_max[4];
     void (*vector_start)(void* trans_data);

src/cpp11.cpp

@@ -52,6 +52,7 @@ extern SEXP wk_c_sfc_writer_new();
 extern SEXP wk_c_trans_affine_as_matrix(SEXP);
 extern SEXP wk_c_trans_affine_new(SEXP);
 extern SEXP wk_c_trans_filter_new(SEXP, SEXP);
+extern SEXP wk_c_trans_set_new(SEXP, SEXP, SEXP);
 extern SEXP wk_c_vector_meta_handler_new();
 extern SEXP wk_c_vertex_filter_new(SEXP, SEXP);
 extern SEXP wk_c_wkb_writer_new(SEXP, SEXP);
@@ -82,6 +83,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"wk_c_trans_affine_as_matrix",  (DL_FUNC) &wk_c_trans_affine_as_matrix,  1},
     {"wk_c_trans_affine_new",        (DL_FUNC) &wk_c_trans_affine_new,        1},
     {"wk_c_trans_filter_new",        (DL_FUNC) &wk_c_trans_filter_new,        2},
+    {"wk_c_trans_set_new",           (DL_FUNC) &wk_c_trans_set_new,           3},
     {"wk_c_vector_meta_handler_new", (DL_FUNC) &wk_c_vector_meta_handler_new, 0},
     {"wk_c_vertex_filter_new",       (DL_FUNC) &wk_c_vertex_filter_new,       2},
     {"wk_c_wkb_writer_new",          (DL_FUNC) &wk_c_wkb_writer_new,          2},

src/handle-wkt.cpp

@@ -705,18 +705,20 @@ cpp11::sexp wk_cpp_handle_wkt(cpp11::strings wkt, cpp11::sexp xptr, bool reveal_
   WKHandlerXPtr cppHandler(xptr);
   WKTStreamingHandler streamer(cppHandler);
 
-  cppHandler.vector_start(&globalMeta);
+  int result = cppHandler.vector_start(&globalMeta);
 
-  for (R_xlen_t i = 0; i < n_features; i++) {
-    if (((i + 1) % 1000) == 0) cpp11::check_user_interrupt();
+  if (result != WK_ABORT) {
+    for (R_xlen_t i = 0; i < n_features; i++) {
+      if (((i + 1) % 1000) == 0) cpp11::check_user_interrupt();
 
-    try {
-      if (streamer.readFeature(&globalMeta, wkt[i], i) == WK_ABORT) {
-        break;
-      }
-    } catch (WKParseException& e) {
-      if (cppHandler.error(e.what()) == WK_ABORT) {
-        break;
+      try {
+        if (streamer.readFeature(&globalMeta, wkt[i], i) == WK_ABORT) {
+          break;
+        }
+      } catch (WKParseException& e) {
+        if (cppHandler.error(e.what()) == WK_ABORT) {
+          break;
+        }
       }
     }
   }

src/trans-affine.c

@@ -31,6 +31,7 @@ SEXP wk_c_trans_affine_new(SEXP trans_matrix) {
     double* trans_matrix_ptr = REAL(trans_matrix);
     double* t = (double*) malloc(6 * sizeof(double));
     if (t == NULL) {
+        free(trans); // # nocov
         Rf_error("Failed to alloc double[6]"); // # nocov
     }
     t[0] = trans_matrix_ptr[0];

src/trans-set.c

@@ -0,0 +1,67 @@
+
+#define R_NO_REMAP
+#include <R.h>
+#include <Rinternals.h>
+#include "wk-v1.h"
+#include <stdlib.h>
+#include <memory.h>
+
+typedef struct {
+    double* xyzm[4];
+    R_xlen_t n;
+} wk_trans_set_t;
+
+int wk_trans_set_trans(R_xlen_t feature_id, double* xyzm_in, double* xyzm_out, void* trans_data) {
+    wk_trans_set_t* data = (wk_trans_set_t*) trans_data;
+    R_xlen_t set_id = feature_id % data->n;
+    double set;
+    for (int i = 0; i < 4; i++) {
+        set = data->xyzm[i][set_id];
+        if (ISNA(set)) {
+            xyzm_out[i] = xyzm_in[i];
+        } else {
+            xyzm_out[i] = set;
+        }
+    }
+    return WK_CONTINUE;
+}
+
+void wk_trans_set_finalize(void* trans_data) {
+    free(trans_data);
+}
+
+SEXP wk_c_trans_set_new(SEXP xy, SEXP use_z, SEXP use_m) {
+    if (Rf_xlength(xy) != 4 || TYPEOF(xy) != VECSXP) {
+        Rf_error("`xy` must be an xyzm() object"); // # nocov
+    }
+
+    // prepare data for C struct / validate args
+    int use_z_int = LOGICAL(use_z)[0];
+    int use_m_int = LOGICAL(use_m)[0];
+    R_xlen_t n = Rf_xlength(VECTOR_ELT(xy, 0));
+    double* xyzm[4];
+    for (int i = 0; i < 4; i++) {
+        xyzm[i] = REAL(VECTOR_ELT(xy, i));
+    }
+
+    // create the wk_trans object
+    wk_trans_t* trans = wk_trans_create();
+    trans->trans = &wk_trans_set_trans;
+    trans->finalizer = &wk_trans_set_finalize;
+
+    wk_trans_set_t* data = (wk_trans_set_t*) malloc(sizeof(wk_trans_set_t));
+    if (data == NULL) {
+        free(trans); // # nocov
+        Rf_error("Failed to alloc wk_trans_set_t"); // # nocov
+    }
+
+    trans->use_z = use_z_int;
+    trans->use_m = use_m_int;
+    memcpy(data->xyzm, xyzm, 4 * sizeof(void*));
+    data->n = n;
+
+    trans->trans_data = data;
+
+    // keep the xy as a tag because we need the pointers to stay valid
+    return wk_trans_create_xptr(trans, xy, R_NilValue);
+}

src/transform.c

@@ -26,11 +26,30 @@ int wk_trans_filter_vector_start(const wk_vector_meta_t* meta, void* handler_dat
   trans_filter_t* trans_filter = (trans_filter_t*) handler_data;
 
   memcpy(&(trans_filter->vector_meta), meta, sizeof(wk_vector_meta_t));
+
+  // bounds are no longer valid
   trans_filter->vector_meta.flags &= ~WK_FLAG_HAS_BOUNDS;
-  if (!(trans_filter->trans->flags_out & WK_FLAG_DIMS_UNKNOWN)) {
+
+  // set the output dimensions NA_INTEGER means "leave alone"
+  int dims_maybe_unknown = 0;
+  if (trans_filter->trans->use_z == 1) {
+    trans_filter->vector_meta.flags |= WK_FLAG_HAS_Z;
+  } else if (trans_filter->trans->use_z == 0) {
     trans_filter->vector_meta.flags &= ~WK_FLAG_HAS_Z;
+  } else {
+    dims_maybe_unknown = 1;
+  }
+
+  if (trans_filter->trans->use_m == 1) {
+    trans_filter->vector_meta.flags |= WK_FLAG_HAS_M;
+  } else if (trans_filter->trans->use_m == 0) {
     trans_filter->vector_meta.flags &= ~WK_FLAG_HAS_M;
-    trans_filter->vector_meta.flags |= trans_filter->trans->flags_out;
+  } else {
+    dims_maybe_unknown = 1;
+  }
+
+  if (!dims_maybe_unknown) {
+    trans_filter->vector_meta.flags &= ~WK_FLAG_DIMS_UNKNOWN;
   }
 
   trans_filter->feature_id = -1;
@@ -71,11 +90,18 @@ int wk_trans_filter_geometry_start(const wk_meta_t* meta, uint32_t part_id, void
   wk_meta_t* new_meta = trans_filter->meta + trans_filter->recursive_level;
   memcpy(new_meta, meta, sizeof(wk_meta_t));
   new_meta->flags &= ~WK_FLAG_HAS_BOUNDS;
-  if (!(trans_filter->trans->flags_out & WK_FLAG_DIMS_UNKNOWN)) {
+  if (trans_filter->trans->use_z == 1) {
+    new_meta->flags |= WK_FLAG_HAS_Z;
+  } else if (trans_filter->trans->use_z == 0) {
     new_meta->flags &= ~WK_FLAG_HAS_Z;
+  }
+
+  if (trans_filter->trans->use_m == 1) {
+    new_meta->flags |= WK_FLAG_HAS_M;
+  } else if (trans_filter->trans->use_m == 0) {
     new_meta->flags &= ~WK_FLAG_HAS_M;
-    new_meta->flags |= trans_filter->trans->flags_out;
   }
+
   return trans_filter->next->geometry_start(new_meta, part_id, trans_filter->next->handler_data);
 }
 
@@ -109,13 +135,13 @@ int wk_trans_filter_coord(const wk_meta_t* meta, const double* coord, uint32_t c
     trans_filter->xyzm_in[3] = coord[3];
   } else if (meta->flags & WK_FLAG_HAS_Z) {
     trans_filter->xyzm_in[2] = coord[2];
-    trans_filter->xyzm_in[3] = NA_REAL;
+    trans_filter->xyzm_in[3] = R_NaN;
   } else if (new_meta->flags & WK_FLAG_HAS_M) {
-    trans_filter->xyzm_in[2] = NA_REAL;
+    trans_filter->xyzm_in[2] = R_NaN;
     trans_filter->xyzm_in[3] = coord[2];
   } else {
-    trans_filter->xyzm_in[2] = NA_REAL;
-    trans_filter->xyzm_in[3] = NA_REAL;
+    trans_filter->xyzm_in[2] = R_NaN;
+    trans_filter->xyzm_in[3] = R_NaN;
   }
 
   int result = trans_filter->trans->trans(

tests/testthat/test-set.R

@@ -0,0 +1,47 @@
+
+test_that("wk_set_(z|m)() works", {
+  expect_identical(wk_set_z(wkt("POINT (0 1)"), 2), wkt("POINT Z (0 1 2)"))
+  expect_identical(wk_set_m(wkt("POINT (0 1)"), 2), wkt("POINT M (0 1 2)"))
+  expect_identical(wk_set_z(wkt("POINT M (0 1 3)"), 2), wkt("POINT ZM (0 1 2 3)"))
+  expect_identical(wk_set_m(wkt("POINT Z (0 1 3)"), 2), wkt("POINT ZM (0 1 3 2)"))
+  expect_identical(wk_set_z(wkt("POINT ZM (0 1 2 3)"), 7), wkt("POINT ZM (0 1 7 3)"))
+  expect_identical(wk_set_m(wkt("POINT ZM (0 1 2 3)"), 7), wkt("POINT ZM (0 1 2 7)"))
+})
+
+test_that("wk_drop_(z|m) works", {
+  expect_identical(wk_drop_z(wkt("POINT ZM (0 1 2 3)")), wkt("POINT M (0 1 3)"))
+  expect_identical(wk_drop_m(wkt("POINT ZM (0 1 2 3)")), wkt("POINT Z (0 1 2)"))
+})
+
+test_that("wk_trans_set() is vectorized", {
+  expect_identical(
+    wk_handle(
+      rep(wkt("POINT Z (0 0 0)"), 4),
+      wk_transform_filter(wkt_writer(), wk_trans_set(xyz(NA, NA, c(1, 2)), use_z = TRUE))
+    ),
+    rep(wkt(c("POINT Z (0 0 1)", "POINT Z (0 0 2)")), 2)
+  )
+})
+
+test_that("wk_trans_set() can set ZM values at the same time", {
+  expect_identical(
+    wk_handle(
+      wkt("POINT (0 0)"),
+      wk_transform_filter(
+        wkt_writer(),
+        wk_trans_set(xyzm(NA, NA, 1, 2), use_z = TRUE, use_m = TRUE)
+      )
+    ),
+    wkt("POINT ZM (0 0 1 2)")
+  )
+})
+
+test_that("wk_trans_set() can set XY values", {
+  expect_identical(
+    wk_handle(
+      wkt("POINT Z (0 0 0)"),
+      wk_transform_filter(wkt_writer(), wk_trans_set(xy(1, 2)))
+    ),
+    wkt("POINT Z (1 2 0)")
+  )
+})