feat: per-step setup/teardown with setup()

DESCRIPTION

@@ -28,7 +28,8 @@ Suggests:
     R6,
     reticulate,
     rmarkdown,
-    testthat (>= 3.0.0)
+    testthat (>= 3.0.0),
+    withr
 VignetteBuilder:
     knitr
 Config/Needs/website: tidyverse/tidytemplate

NAMESPACE

@@ -21,5 +21,6 @@ export(gen)
 export(generator)
 export(is_exhausted)
 export(loop)
+export(setup)
 export(yield)
 import(rlang)

NEWS.md

@@ -1,5 +1,10 @@
 # coro (development version)
 
+* New `setup()` registers an expression that runs at the start of every step of
+  a generator or async function, with any `withr::defer()`/`on.exit()` it
+  registers torn down at the end of that step. This gives per-step setup/teardown
+  parity, including around `await()`/`yield()` (#68).
+
 * Fixed performance issue related to JIT compilation of generator/async instances (#71). The JIT compiler is now able to properly cache the instances, making repeated instantiations fast, and preventing the JIT from kicking in on every call.
 
 * Async functions and generators can now be R6 methods (#63).

R/generator.R

@@ -152,6 +152,8 @@ generator0 <- function(fn, type = "generator") {
         exited <- NULL
         cleanup <- NULL
         close_active_iterators <- NULL
+        run_setups <- NULL
+        run_step_teardowns <- NULL
         `.__generator_env__.` <- NULL
       }
 
@@ -274,8 +276,12 @@ generator0 <- function(fn, type = "generator") {
 
         out <- evalq(envir = user_env, {
           base::evalq(envir = .__generator_env__., {
+            # LIFO, independent execution
             defer(if (exited) cleanup())
+            # `setup()` teardowns fire at every step end (suspend/return/error).
+            defer(run_step_teardowns())
             env_poke_exits(user_env, exits)
+            run_setups()
             !!state_machine
           })
         })
@@ -353,6 +359,8 @@ new_generator_env <- function(parent, info) {
     }
   })
 
+  init_setup_runtime(env)
+
   if (!is_null(info$async_ops)) {
     env$then <- info$async_ops$then
     env$as_promise <- info$async_ops$as_promise
@@ -361,6 +369,80 @@ new_generator_env <- function(parent, info) {
   env
 }
 
+# Installs the `setup()` lifecycle into the generator's runtime `env`:
+# registration + per-call-site deduping (`do_setup()`), the capture-and-clear
+# harvest of each setup body's teardowns (`run_one_setup()`), re-running the
+# registered setups at the start of each step (`run_setups()`), and replaying
+# the harvested teardowns at step end (`run_step_teardowns()`). Kept separate so
+# `new_generator_env()` stays focused on env construction.
+init_setup_runtime <- function(env) {
+  env$setups <- list()
+  env$setup_ids <- integer()
+  env$step_teardowns <- list()
+
+  with(env, {
+    run_one_setup <- function(expr) {
+      # Run the setup body in a throwaway closure whose frame is an isolated
+      # child of `user_env`. Capture its `on.exit()`/`withr::defer()`
+      # registrations (without firing them) plus the frame itself, so the
+      # teardowns can be replayed at step end in the same environment where the
+      # body's locals live (a bare `on.exit(x <- old)` references those locals).
+      harvester <- new_function(
+        pairlist2(),
+        block(
+          expr,
+          quote(`.__coro_setup_td__` <- list(
+            env = environment(),
+            exits = sys.on.exit()
+          )),
+          quote(on.exit()),
+          quote(`.__coro_setup_td__`)
+        ),
+        env = new.env(parent = user_env)
+      )
+      step_teardowns <<- c(step_teardowns, list(harvester()))
+    }
+
+    do_setup <- function(id, expr) {
+      if (id %in% setup_ids) {
+        return(invisible(NULL))
+      }
+      setup_ids <<- c(setup_ids, id)
+      setups <<- c(setups, list(expr))
+      run_one_setup(expr)
+    }
+
+    run_setups <- function() {
+      step_teardowns <<- list()
+      for (expr in setups) {
+        run_one_setup(expr)
+      }
+    }
+
+    run_step_teardowns <- function() {
+      err <- NULL
+      for (teardown in rev(step_teardowns)) {
+        if (is_null(teardown$exits)) {
+          next
+        }
+        exprs <- if (is_call(teardown$exits, "{")) as.list(teardown$exits)[-1] else list(teardown$exits)
+        for (expr in exprs) {
+          tryCatch(
+            eval_bare(expr, teardown$env),
+            error = function(cnd) if (is_null(err)) err <<- cnd
+          )
+        }
+      }
+      step_teardowns <<- list()
+      if (!is_null(err)) {
+        cnd_signal(err)
+      }
+    }
+  })
+
+  invisible(env)
+}
+
 env_bind_arg <- function(env, arg, frame = caller_env()) {
   if (identical(arg, "...")) {
     env[["..."]] <- env_get(
@@ -492,5 +574,6 @@ utils::globalVariables(c(
   "user",
   "exits",
   "suspend",
-  "generator_env"
+  "generator_env",
+  "do_setup"
 ))

R/parser.R

@@ -1,4 +1,6 @@
 walk_states <- function(expr, info) {
+  check_setup_calls(expr)
+
   continue <- function(counter, last) {
     # Break if last
     if (last) 0L else counter() + 1L
@@ -24,6 +26,33 @@ walk_states <- function(expr, info) {
     invisible(exhausted())
   })
 }
+
+# `setup()` may not appear inside a loop. Mixing per-step registration with
+# iteration (and the branching/yields a loop body may contain) has opaque
+# semantics. For per-iteration setup and teardown, factor the loop body into
+# its own generator and delegate to it, e.g. `for (x in step(i)) yield(x)`.
+# Nested `function` definitions are their own coroutines and are not descended
+# into. Setup bodies are user code, not a nested `setup()`, so we stop there too.
+check_setup_calls <- function(expr, in_loop = FALSE) {
+  if (!is_call(expr) || is_call(expr, "function")) {
+    return(invisible())
+  }
+  if (is_setup_call(expr)) {
+    if (in_loop) {
+      abort(c(
+        "Can't use `setup()` within a loop.",
+        i = "For per-iteration setup and teardown, move the loop body into its own generator and delegate to it with `for (x in step(i)) yield(x)`."
+      ))
+    }
+    return(invisible())
+  }
+  in_loop <- in_loop || is_call(expr, c("for", "while", "repeat"))
+  for (arg in as.list(expr)[-1]) {
+    check_setup_calls(arg, in_loop)
+  }
+  invisible()
+}
+
 walk_loop_states <- function(body, states, counter, info) {
   continue <- function(counter, last) {
     # Go back to state 1 of loop body if last
@@ -154,6 +183,15 @@ expr_states <- function(expr, counter, continue, last, return, info) {
       info = info,
       assign_var = yield_lhs(expr, assign)
     ),
+    `setup` = setup_state(
+      preamble = NULL,
+      expr = node_cadr(expr),
+      counter = counter,
+      continue = continue,
+      last = last,
+      return = return,
+      info = info
+    ),
     `return` = return_state(
       expr = user_call(strip_explicit_return(expr)),
       counter = counter,
@@ -262,7 +300,7 @@ expr_type_impl <- function(expr) {
   head <- node_car(expr)
   if (is_call(head, "::") && identical(head[[2]], quote(coro))) {
     head <- head[[3]]
-    if (!as_string(head) %in% c("yield", "await", "await_each")) {
+    if (!as_string(head) %in% c("yield", "await", "await_each", "setup")) {
       return(default)
     }
   }
@@ -277,6 +315,7 @@ expr_type_impl <- function(expr) {
     `{` = ,
     `yield` = ,
     `await` = ,
+    `setup` = ,
     `return` = ,
     `if` = ,
     `repeat` = ,
@@ -466,6 +505,19 @@ block_states <- function(block, counter, continue, last, return, info) {
         ))
         next
       },
+      `setup` = {
+        skip()
+        push_states(setup_state(
+          preamble = collect(),
+          expr = node_cadr(expr),
+          counter = counter,
+          continue = continue,
+          last = last,
+          return = return,
+          info = info
+        ))
+        next
+      },
       `repeat` = {
         node_poke_car(node, "repeat")
         push_states(loop_states(
@@ -1032,6 +1084,42 @@ next_state <- function(preamble, counter, info) {
   state
 }
 
+setup_state <- function(preamble, expr, counter, continue, last, return, info) {
+  check_setup_body(expr)
+
+  i <- counter()
+  next_i <- continue(counter, last)
+  depth <- machine_depth(counter)
+
+  block <- expr({
+    !!!preamble %&&% list(user_call(preamble))
+    do_setup(!!i, !!call2("quote", expr))
+    !!continue_call(next_i, depth)
+  })
+  state <- new_state(block, NULL, tag = i)
+  counter(inc = 1L)
+
+  state
+}
+
+# `setup()` bodies run as plain code outside the state machine, so they can't
+# suspend. Reject `yield()`/`await()`/`await_each()` at compile time. Don't
+# descend into nested `function` definitions (a nested generator/async is fine).
+check_setup_body <- function(expr) {
+  if (!is_call(expr) || is_call(expr, "function")) {
+    return(invisible())
+  }
+  for (nm in c("yield", "await", "await_each")) {
+    if (is_call(expr, nm, ns = c("", "coro"))) {
+      abort(sprintf("Can't use `%s()` within `setup()`.", nm))
+    }
+  }
+  for (arg in as.list(expr)[-1]) {
+    check_setup_body(arg)
+  }
+  invisible()
+}
+
 try_catch_states <- function(
   preamble,
   expr,