make TableCollection() aligned with Python version

Author: LynxJinyangii

RcppTskit/NEWS.md

@@ -43,7 +43,8 @@ and releases adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html
 - Added `rtsk_node_table_get_row()` and `TableCollection$node_table_get_row()`
   to retrieve node-table rows by 0-based row index.
 - Added `rtsk_table_collection_sort()` and `TableCollection$sort()` to sort
-  table collections with 0-based `edge_start` semantics.
+  table collections with 0-based `edge_start`, `site_start`, and
+  `mutation_start` semantics.
 - Added low-level variant iterators
   (`rtsk_variant_iterator_init()`/`rtsk_variant_iterator_next()`) and a
   user-facing `TreeSequence$variants()` method to iterate over decoded

RcppTskit/R/Class-TableCollection.R

@@ -103,39 +103,28 @@ TableCollection <- R6Class(
       TreeSequence$new(xptr = ts_xptr)
     },
 
-    # TODO: add site_start and mutation_start
     #' @description Sort this table collection in place.
     #' @param edge_start integer scalar edge-table start row index (0-based).
-    #' # TODO: remove this argument since Python code does not have it
-    #' @param no_check_integrity logical; when \code{TRUE}, pass
-    #'   \code{TSK_NO_CHECK_INTEGRITY} to \code{tskit C}.
+    #' @param site_start integer scalar site-table start row index (0-based).
+    #' @param mutation_start integer scalar mutation-table start row index
+    #'   (0-based).
     #' @details See the \code{tskit Python} equivalent at
     #'   \url{https://tskit.dev/tskit/docs/latest/python-api.html#tskit.TableCollection.sort}.
     #' @return No return value; called for side effects.
     #' @examples
     #' ts_file <- system.file("examples/test.trees", package = "RcppTskit")
     #' tc <- tc_load(ts_file)
     #' tc$sort()
-    sort = function(edge_start = 0L, no_check_integrity = FALSE) {
+    sort = function(edge_start = 0L, site_start = 0L, mutation_start = 0L) {
       validate_row_index(edge_start, "edge_start")
-      # TODO: remove this argument since Python code does not have it
-      if (
-        !is.logical(no_check_integrity) ||
-          length(no_check_integrity) != 1L ||
-          is.na(no_check_integrity)
-      ) {
-        stop("no_check_integrity must be TRUE/FALSE!")
-      }
-      options <- if (isTRUE(no_check_integrity)) {
-        as.integer(rtsk_const_tsk_no_check_integrity())
-      } else {
-        0L
-      }
+      validate_row_index(site_start, "site_start")
+      validate_row_index(mutation_start, "mutation_start")
       rtsk_table_collection_sort(
         tc = self$xptr,
         edge_start = as.integer(edge_start),
-        # TODO: remove this argument since Python code does not have it
-        options = options
+        site_start = as.integer(site_start),
+        mutation_start = as.integer(mutation_start),
+        options = 0L
       )
     },
 

RcppTskit/R/RcppExports.R

@@ -223,8 +223,8 @@ rtsk_table_collection_drop_index <- function(tc, options = 0L) {
     invisible(.Call(`_RcppTskit_rtsk_table_collection_drop_index`, tc, options))
 }
 
-rtsk_table_collection_sort <- function(tc, edge_start = 0L, options = 0L) {
-    invisible(.Call(`_RcppTskit_rtsk_table_collection_sort`, tc, edge_start, options))
+rtsk_table_collection_sort <- function(tc, edge_start = 0L, site_start = 0L, mutation_start = 0L, options = 0L) {
+    invisible(.Call(`_RcppTskit_rtsk_table_collection_sort`, tc, edge_start, site_start, mutation_start, options))
 }
 
 rtsk_table_collection_summary <- function(tc) {

RcppTskit/inst/include/RcppTskit_public.hpp

@@ -63,7 +63,8 @@ Rcpp::String rtsk_table_collection_get_file_uuid(SEXP tc);
 bool rtsk_table_collection_has_index(SEXP tc, int options = 0);
 void rtsk_table_collection_build_index(SEXP tc, int options = 0);
 void rtsk_table_collection_drop_index(SEXP tc, int options = 0);
-void rtsk_table_collection_sort(SEXP tc, int edge_start = 0, int options = 0);
+void rtsk_table_collection_sort(SEXP tc, int edge_start = 0, int site_start = 0,
+                                int mutation_start = 0, int options = 0);
 Rcpp::List rtsk_table_collection_summary(SEXP tc);
 Rcpp::List rtsk_table_collection_metadata_length(SEXP tc);
 int rtsk_individual_table_add_row(

RcppTskit/man/TableCollection.Rd

@@ -594,14 +594,16 @@ BEGIN_RCPP
 END_RCPP
 }
 // rtsk_table_collection_sort
-void rtsk_table_collection_sort(SEXP tc, int edge_start, int options);
-RcppExport SEXP _RcppTskit_rtsk_table_collection_sort(SEXP tcSEXP, SEXP edge_startSEXP, SEXP optionsSEXP) {
+void rtsk_table_collection_sort(SEXP tc, int edge_start, int site_start, int mutation_start, int options);
+RcppExport SEXP _RcppTskit_rtsk_table_collection_sort(SEXP tcSEXP, SEXP edge_startSEXP, SEXP site_startSEXP, SEXP mutation_startSEXP, SEXP optionsSEXP) {
 BEGIN_RCPP
     Rcpp::RNGScope rcpp_rngScope_gen;
     Rcpp::traits::input_parameter< SEXP >::type tc(tcSEXP);
     Rcpp::traits::input_parameter< int >::type edge_start(edge_startSEXP);
+    Rcpp::traits::input_parameter< int >::type site_start(site_startSEXP);
+    Rcpp::traits::input_parameter< int >::type mutation_start(mutation_startSEXP);
     Rcpp::traits::input_parameter< int >::type options(optionsSEXP);
-    rtsk_table_collection_sort(tc, edge_start, options);
+    rtsk_table_collection_sort(tc, edge_start, site_start, mutation_start, options);
     return R_NilValue;
 END_RCPP
 }
@@ -899,7 +901,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"_RcppTskit_rtsk_table_collection_has_index", (DL_FUNC) &_RcppTskit_rtsk_table_collection_has_index, 2},
     {"_RcppTskit_rtsk_table_collection_build_index", (DL_FUNC) &_RcppTskit_rtsk_table_collection_build_index, 2},
     {"_RcppTskit_rtsk_table_collection_drop_index", (DL_FUNC) &_RcppTskit_rtsk_table_collection_drop_index, 2},
-    {"_RcppTskit_rtsk_table_collection_sort", (DL_FUNC) &_RcppTskit_rtsk_table_collection_sort, 3},
+    {"_RcppTskit_rtsk_table_collection_sort", (DL_FUNC) &_RcppTskit_rtsk_table_collection_sort, 5},
     {"_RcppTskit_rtsk_table_collection_summary", (DL_FUNC) &_RcppTskit_rtsk_table_collection_summary, 1},
     {"_RcppTskit_rtsk_table_collection_metadata_length", (DL_FUNC) &_RcppTskit_rtsk_table_collection_metadata_length, 1},
     {"_RcppTskit_rtsk_individual_table_add_row", (DL_FUNC) &_RcppTskit_rtsk_individual_table_add_row, 5},

RcppTskit/src/RcppExports.cpp

@@ -1365,14 +1365,16 @@ void rtsk_table_collection_drop_index(SEXP tc, int options = 0) {
   // # nocov end
 }
 
-// TODO: model this after
-// https://github.com/tskit-dev/tskit/blob/2f26dc6f0d033cdf7d6fb1adfbab96468dde8831/python/_tskitmodule.c#L4526
 // PUBLIC, wrapper for tsk_table_collection_sort
 // @title Sort a table collection
 // @param tc an external pointer to table collection as a
 //   \code{tsk_table_collection_t} object.
 // @param edge_start integer scalar edge-table start row index (0-based) used in
 //   sorting bookmark.
+// @param site_start integer scalar site-table start row index (0-based) used in
+//   sorting bookmark.
+// @param mutation_start integer scalar mutation-table start row index (0-based)
+//   used in sorting bookmark.
 // @param options passed to \code{tskit C}; this wrapper supports
 //   \code{TSK_NO_CHECK_INTEGRITY}.
 // @details This function calls
@@ -1383,19 +1385,36 @@ void rtsk_table_collection_drop_index(SEXP tc, int options = 0) {
 // tc_xptr <- RcppTskit:::rtsk_table_collection_load(ts_file)
 // RcppTskit:::rtsk_table_collection_sort(tc_xptr)
 // [[Rcpp::export]]
-void rtsk_table_collection_sort(SEXP tc, int edge_start = 0, int options = 0) {
+void rtsk_table_collection_sort(SEXP tc, int edge_start = 0, int site_start = 0,
+                                int mutation_start = 0, int options = 0) {
   if (Rcpp::IntegerVector::is_na(edge_start)) {
     Rcpp::stop(
         "edge_start must not be NA_integer_ in rtsk_table_collection_sort");
   }
   if (edge_start < 0) {
     Rcpp::stop("edge_start must be >= 0 in rtsk_table_collection_sort");
   }
+  if (Rcpp::IntegerVector::is_na(site_start)) {
+    Rcpp::stop(
+        "site_start must not be NA_integer_ in rtsk_table_collection_sort");
+  }
+  if (site_start < 0) {
+    Rcpp::stop("site_start must be >= 0 in rtsk_table_collection_sort");
+  }
+  if (Rcpp::IntegerVector::is_na(mutation_start)) {
+    Rcpp::stop(
+        "mutation_start must not be NA_integer_ in rtsk_table_collection_sort");
+  }
+  if (mutation_start < 0) {
+    Rcpp::stop("mutation_start must be >= 0 in rtsk_table_collection_sort");
+  }
   const tsk_flags_t flags = validate_supported_options(
       options, kTableSortSupportedFlags, "rtsk_table_collection_sort");
   rtsk_table_collection_t tc_xptr(tc);
   tsk_bookmark_t start = {0};
   start.edges = static_cast<tsk_size_t>(edge_start);
+  start.sites = static_cast<tsk_size_t>(site_start);
+  start.mutations = static_cast<tsk_size_t>(mutation_start);
   int ret = tsk_table_collection_sort(tc_xptr, &start, flags);
   if (ret != 0) {
     Rcpp::stop(tsk_strerror(ret));

RcppTskit/src/RcppTskit.cpp

@@ -290,15 +290,28 @@ test_that("table_collection_sort wrapper validates inputs and sorts in place", {
     rtsk_table_collection_sort(tc_xptr, edge_start = -1L),
     regexp = "edge_start must be >= 0 in rtsk_table_collection_sort"
   )
+  expect_error(
+    rtsk_table_collection_sort(tc_xptr, site_start = NA_integer_),
+    regexp = "site_start must not be NA_integer_ in rtsk_table_collection_sort"
+  )
+  expect_error(
+    rtsk_table_collection_sort(tc_xptr, site_start = -1L),
+    regexp = "site_start must be >= 0 in rtsk_table_collection_sort"
+  )
+  expect_error(
+    rtsk_table_collection_sort(tc_xptr, mutation_start = NA_integer_),
+    regexp = "mutation_start must not be NA_integer_ in rtsk_table_collection_sort"
+  )
+  expect_error(
+    rtsk_table_collection_sort(tc_xptr, mutation_start = -1L),
+    regexp = "mutation_start must be >= 0 in rtsk_table_collection_sort"
+  )
   expect_error(
     rtsk_table_collection_sort(tc_xptr, options = bitwShiftL(1L, 4)),
     regexp = "only supports options"
   )
   expect_no_error(rtsk_table_collection_sort(tc_xptr))
-  expect_no_error(rtsk_table_collection_sort(
-    tc_xptr,
-    options = as.integer(rtsk_const_tsk_no_check_integrity())
-  ))
+  expect_no_error(rtsk_table_collection_sort(tc_xptr, 0L, 0L, 0L))
 
   expect_error(
     tc$sort(edge_start = NA_integer_),
@@ -309,11 +322,19 @@ test_that("table_collection_sort wrapper validates inputs and sorts in place", {
     regexp = "edge_start must be a non-NA zero or positive integer scalar!"
   )
   expect_error(
-    tc$sort(no_check_integrity = NA),
-    regexp = "no_check_integrity must be TRUE/FALSE!"
+    tc$sort(site_start = NA_integer_),
+    regexp = "site_start must be a non-NA zero or positive integer scalar!"
+  )
+  expect_error(
+    tc$sort(mutation_start = NA_integer_),
+    regexp = "mutation_start must be a non-NA zero or positive integer scalar!"
   )
   expect_no_error(tc$sort())
-  expect_no_error(tc$sort(no_check_integrity = TRUE))
+  expect_no_error(tc$sort(
+    edge_start = 0L,
+    site_start = 0L,
+    mutation_start = 0L
+  ))
 })
 
 test_that("individual_table_add_row wrapper expands the table collection and handles inputs", {