draft purrr_continue()

tidyverse · Sep 5, 2023 · b60133f · b60133f
1 parent ac4f5a9
commit b60133f
Show file tree

Hide file tree

Showing 6 changed files with 70 additions and 19 deletions.
diff --git a/R/map.R b/R/map.R
@@ -169,13 +169,46 @@ map_ <- function(.type,
   names <- vec_names(.x)
 
   .f <- as_mapper(.f, ...)
-
   i <- 0L
+  print(i)
+
+  the$last_map_index <- NULL
+  the$last_map_results <- NULL
+  the$last_map <- list(
+    env = current_env(),
+    call = expr(call_with_cleanup(map_impl, environment(), .type, .progress, n, names, i, the))
+  )
+
   with_indexed_errors(
     i = i,
     names = names,
     error_call = .purrr_error_call,
-    call_with_cleanup(map_impl, environment(), .type, .progress, n, names, i)
+    call_with_cleanup(map_impl, environment(), .type, .progress, n, names, i, the)
+  )
+}
+
+purrr_continue <- function(.f = NULL) {
+  # add 0 to force a copy
+  i <- the$last_map_index + 0L
+  last_map_index <- the$last_map_index
+  last_map_results <- the$last_map_results
+
+  env2 <- the$last_map$env
+  env2$.f <- .f %||% the$last_map$.f
+  env2$i <- i
+
+  new_map_results <- with_indexed_errors(
+    i = i,
+    names = env2$names,
+    error_call = env2$.purrr_error_call,
+    rlang::eval_bare(the$last_map$call, env = env2)
+  )
+
+  idx <- seq2(1, last_map_index)
+  vctrs::vec_assign(
+    new_map_results,
+    idx,
+    vctrs::vec_slice(last_map_results, idx)
   )
 }
 

diff --git a/R/map2.R b/R/map2.R
@@ -79,7 +79,7 @@ map2_ <- function(.type,
     i = i,
     names = names,
     error_call = .purrr_error_call,
-    call_with_cleanup(map2_impl, environment(), .type, .progress, n, names, i)
+    call_with_cleanup(map2_impl, environment(), .type, .progress, n, names, i, the)
   )
 }
 

diff --git a/R/package-purrr.R b/R/package-purrr.R
@@ -7,3 +7,6 @@
 "_PACKAGE"
 
 the <- new_environment()
+the$last_map_results <- NULL
+the$last_map_index <- NULL
+the$last_map <- NULL
diff --git a/R/pmap.R b/R/pmap.R
@@ -128,7 +128,7 @@ pmap_ <- function(.type,
     i = i,
     names = names,
     error_call = .purrr_error_call,
-    call_with_cleanup(pmap_impl, environment(), .type, .progress, n, names, i, call_names, call_n)
+    call_with_cleanup(pmap_impl, environment(), .type, .progress, n, names, i, call_names, call_n, the)
   )
 }
 

diff --git a/src/init.c b/src/init.c
@@ -13,9 +13,9 @@
 extern SEXP coerce_impl(SEXP, SEXP);
 extern SEXP pluck_impl(SEXP, SEXP, SEXP, SEXP);
 extern SEXP flatten_impl(SEXP);
-extern SEXP map_impl(SEXP, SEXP, SEXP, SEXP, SEXP, SEXP);
-extern SEXP map2_impl(SEXP, SEXP, SEXP, SEXP, SEXP, SEXP);
-extern SEXP pmap_impl(SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP);
+extern SEXP map_impl(SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP);
+extern SEXP map2_impl(SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP);
+extern SEXP pmap_impl(SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP);
 extern SEXP transpose_impl(SEXP, SEXP);
 extern SEXP vflatten_impl(SEXP, SEXP);
 
@@ -24,9 +24,9 @@ static const R_CallMethodDef CallEntries[] = {
   {"coerce_impl",           (DL_FUNC) &coerce_impl,    2},
   {"pluck_impl",            (DL_FUNC) &pluck_impl,     4},
   {"flatten_impl",          (DL_FUNC) &flatten_impl,   1},
-  {"map_impl",              (DL_FUNC) &map_impl,       6},
-  {"map2_impl",             (DL_FUNC) &map2_impl,      6},
-  {"pmap_impl",             (DL_FUNC) &pmap_impl,      8},
+  {"map_impl",              (DL_FUNC) &map_impl,       7},
+  {"map2_impl",             (DL_FUNC) &map2_impl,      7},
+  {"pmap_impl",             (DL_FUNC) &pmap_impl,      9},
   {"transpose_impl",        (DL_FUNC) &transpose_impl, 2},
   {"vflatten_impl",         (DL_FUNC) &vflatten_impl,  2},
   {"purrr_eval",            (DL_FUNC) &Rf_eval,        2},

diff --git a/src/map.c b/src/map.c
@@ -25,15 +25,24 @@ SEXP call_loop(SEXP env,
                int n,
                SEXP names,
                int* p_i,
-               int force) {
+               int force,
+               SEXP the_env) {
   SEXP bar = cli_progress_bar(n, progress);
   R_PreserveObject(bar);
   r_call_on_exit((void (*)(void*)) cb_progress_done, (void*) bar);
 
+  SEXP results_symbol = Rf_install("last_map_results");
+  SEXP index_symbol = Rf_install("last_map_index");
+  int start = *p_i;
+  SEXP index = PROTECT(Rf_ScalarInteger(start));
+  Rf_defineVar(index_symbol, index, the_env);
+
   SEXP out = PROTECT(Rf_allocVector(type, n));
+  Rf_defineVar(results_symbol, out, the_env);
   Rf_setAttrib(out, R_NamesSymbol, names);
 
-  for (int i = 0; i < n; ++i) {
+  for (int i = start; i < n; ++i) {
+    SET_INTEGER_ELT(index, 0, i);
     *p_i = i + 1;
 
     if (CLI_SHOULD_TICK) {
@@ -55,7 +64,7 @@ SEXP call_loop(SEXP env,
 
   *p_i = 0;
 
-  UNPROTECT(1);
+  UNPROTECT(2);
   return out;
 }
 
@@ -64,7 +73,8 @@ SEXP map_impl(SEXP env,
               SEXP progress,
               SEXP ffi_n,
               SEXP names,
-              SEXP i) {
+              SEXP i,
+              SEXP the_env) {
   static SEXP call = NULL;
   if (call == NULL) {
     SEXP x_sym = Rf_install(".x");
@@ -95,7 +105,8 @@ SEXP map_impl(SEXP env,
     n,
     names,
     p_i,
-    force
+    force,
+    the_env
   );
 }
 
@@ -104,7 +115,8 @@ SEXP map2_impl(SEXP env,
                SEXP progress,
                SEXP ffi_n,
                SEXP names,
-               SEXP i) {
+               SEXP i,
+               SEXP the_env) {
   static SEXP call = NULL;
   if (call == NULL) {
     SEXP x_sym = Rf_install(".x");
@@ -135,7 +147,8 @@ SEXP map2_impl(SEXP env,
     n,
     names,
     p_i,
-    force
+    force,
+    the_env
   );
 }
 
@@ -146,7 +159,8 @@ SEXP pmap_impl(SEXP env,
                SEXP names,
                SEXP i,
                SEXP call_names,
-               SEXP ffi_call_n) {
+               SEXP ffi_call_n,
+               SEXP the_env) {
   // Construct call like f(.l[[1]][[i]], .l[[2]][[i]], ...)
   //
   // Currently accessing S3 vectors in a list like .l[[c(1, i)]] will not
@@ -203,7 +217,8 @@ SEXP pmap_impl(SEXP env,
     n,
     names,
     p_i,
-    force
+    force,
+    the_env
   );
 
   UNPROTECT(1);