tweaks for recurse

rstudio · Oct 6, 2024 · cd3d5cd · cd3d5cd
1 parent 855593e
commit cd3d5cd
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 34 deletions.
diff --git a/inst/ext/renv.c b/inst/ext/renv.c
@@ -19,6 +19,10 @@ static const int _STRSXP = STRSXP;
 static const int _VECSXP = VECSXP;
 static const int _ENVSXP = ENVSXP;
 
+// Initialized in R_init_renv
+static SEXP s_callbacksym;
+static SEXP s_objectsym;
+
 static SEXP renv_call_expect(SEXP node,
                              SEXP package,
                              SEXP methods)
@@ -283,53 +287,77 @@ static SEXP enumerate(SEXP x,
   return R_NilValue;
 }
 
-static SEXP recurse_impl(SEXP object,
-                         SEXP objectsym,
-                         SEXP callback,
-                         SEXP callbacksym,
-                         SEXP envir)
+static SEXP recurse(SEXP object,
+                    SEXP callback,
+                    SEXP envir)
 {
-  if (object != R_MissingArg)
+  SEXP symbol, expr, frame = R_NilValue;
+  SEXP dots = Rf_findVarInFrame(envir, R_DotsSymbol);
+  if (TYPEOF(callback) == CLOSXP && dots == R_MissingArg)
   {
-    Rf_defineVar(objectsym, object, envir);
-    SEXP call = Rf_lang3(callbacksym, objectsym, R_DotsSymbol);
-    R_forceAndCall(call, 1, envir);
+    symbol = TAG(FORMALS(callback));
+    expr = BODY(callback);
+    frame = PROTECT(R_NewEnv(CLOENV(callback), 0, 29));
   }
 
-  switch (TYPEOF(object))
-  {
-  case VECSXP:
-  case EXPRSXP:
-  {
-    for (R_xlen_t i = 0, n = Rf_xlength(object); i < n; i++)
-      recurse_impl(VECTOR_ELT(object, i), objectsym, callback, callbacksym, envir);
-    break;
-  }
+  const int size = 16384;
+  SEXP queue[size];
+  queue[0] = object;
 
-  case LISTSXP:
-  case LANGSXP:
+  int index = 0;
+  int slot = 1;
+
+  while (index != slot)
   {
-    while (object != R_NilValue)
+    object = queue[index++];
+    index = index % size;
+
+    if (object != R_MissingArg)
     {
-      recurse_impl(CAR(object), objectsym, callback, callbacksym, envir);
-      object = CDR(object);
+      if (frame == R_NilValue)
+      {
+        Rf_defineVar(s_objectsym, object, envir);
+        SEXP call = Rf_lang3(s_callbacksym, s_objectsym, R_DotsSymbol);
+        R_forceAndCall(call, 1, envir);
+      }
+      else
+      {
+        Rf_defineVar(symbol, object, frame);
+        Rf_eval(expr, frame);
+      }
+    }
+
+    switch (TYPEOF(object))
+    {
+    case VECSXP:
+    case EXPRSXP:
+    {
+      for (R_xlen_t i = 0, n = Rf_xlength(object); i < n; i++)
+      {
+        queue[slot++] = VECTOR_ELT(object, i);
+        slot = slot % size;
+      }
+      break;
+    }
+
+    case LISTSXP:
+    case LANGSXP:
+    {
+      while (object != R_NilValue)
+      {
+        queue[slot++] = CAR(object);
+        slot = slot % size;
+        object = CDR(object);
+      }
+      break;
+    }
     }
-    break;
-  }
   }
 
+  UNPROTECT(frame != R_NilValue ? 1 : 0);
   return R_NilValue;
 }
 
-static SEXP recurse(SEXP object,
-                    SEXP callback,
-                    SEXP envir)
-{
-  SEXP callbacksym = Rf_install("callback");
-  SEXP objectsym = Rf_install("object");
-  return recurse_impl(object, objectsym, callback, callbacksym, envir);
-}
-
 // Init ----
 
 static const R_CallMethodDef callEntries[] = {
@@ -342,6 +370,9 @@ static const R_CallMethodDef callEntries[] = {
 
 void R_init_renv(DllInfo* dllInfo)
 {
+  s_callbacksym = Rf_install("callback");
+  s_objectsym = Rf_install("object");
+
   R_registerRoutines(dllInfo, NULL, callEntries, NULL, NULL);
   R_useDynamicSymbols(dllInfo, FALSE);
 }
diff --git a/tests/testthat/test-recurse.R b/tests/testthat/test-recurse.R
@@ -21,10 +21,25 @@ test_that("recurse() can handle lists", {
 
   expect_equal(items, list(1, 2, 3, 4))
 
+  items <- list()
+  recurse(data, function(el, ignored) {
+    if (is.numeric(el))
+      items[[length(items) + 1L]] <<- el
+  }, ignored = 42)
+
+  expect_equal(items, list(1, 2, 3, 4))
+
 })
 
 test_that("recurse() can handle dots", {
 
+  counter <- 0L
+  recurse(list(1, list(2, list(3, list(4, list(5))))), function(node) {
+    if (is.list(node))
+      counter <<- counter + 1L
+  })
+  expect_equal(counter, 5L)
+
   counter <- 0L
   recurse(list(1, list(2, list(3, list(4, list(5))))), function(node, extra) {
     expect_equal(extra, 42)