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

DESCRIPTION

@@ -27,7 +27,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).
+
 * Async functions and generators can now be R6 methods (#63).
 
 * A memory leak that crept back in (see #36) was fixed with a more robust approach.

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 <- environment()$generator_env
@@ -289,7 +291,12 @@ generator0 <- function(fn, type = "generator") {
         out <- evalq(envir = user_env, {
           base::evalq(envir = rlang::wref_key(!!weak_env), {
             defer(if (exited) cleanup())
+            # `setup()` teardowns fire at every step end (suspend/return/error).
+            # Registered as a separate, later `defer()` so they run before
+            # `cleanup()` (LIFO) yet can't prevent `cleanup()` if one errors.
+            defer(run_step_teardowns())
             env_poke_exits(user_env, exits)
+            run_setups()
             !!state_machine
           })
         })
@@ -347,6 +354,9 @@ new_generator_env <- function(parent, info) {
   env$exits <- NULL
   env$exited <- TRUE
   env$.last_value <- NULL
+  env$setups <- list()
+  env$setup_ids <- integer()
+  env$step_teardowns <- list()
 
   with(env, {
     user <- function(expr) {
@@ -360,6 +370,58 @@ new_generator_env <- function(parent, info) {
       exited <<- FALSE
       exits <<- env_poke_exits(user_env, NULL)
     }
+
+    run_one_setup <- function(expr) {
+      e <- new.env(parent = user_env)
+      harvester <- new_function(
+        pairlist2(),
+        block(
+          expr,
+          quote(`.__coro_setup_exits__` <- sys.on.exit()),
+          quote(on.exit()),
+          quote(`.__coro_setup_exits__`)
+        ),
+        env = e
+      )
+      td <- list(env = e, exits = harvester())
+      step_teardowns <<- c(step_teardowns, list(td))
+    }
+
+    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 (td in rev(step_teardowns)) {
+        if (is_null(td$exits)) {
+          next
+        }
+        exprs <- if (is_call(td$exits, "{")) as.list(td$exits)[-1] else list(td$exits)
+        for (ex in exprs) {
+          tryCatch(
+            eval_bare(ex, td$env),
+            error = function(cnd) if (is_null(err)) err <<- cnd
+          )
+        }
+      }
+      step_teardowns <<- list()
+      if (!is_null(err)) {
+        cnd_signal(err)
+      }
+    }
   })
 
   if (!is_null(info$async_ops)) {
@@ -501,5 +563,6 @@ utils::globalVariables(c(
   "user",
   "exits",
   "suspend",
-  "generator_env"
+  "generator_env",
+  "do_setup"
 ))

R/parser.R

@@ -154,6 +154,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 +271,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 +286,7 @@ expr_type_impl <- function(expr) {
     `{` = ,
     `yield` = ,
     `await` = ,
+    `setup` = ,
     `return` = ,
     `if` = ,
     `repeat` = ,
@@ -466,6 +476,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 +1055,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,

R/setup.R

@@ -0,0 +1,57 @@
+#' Set up per-step state in a coroutine
+#'
+#' @description
+#' `setup()` registers an expression that runs at the start of **every** step of
+#' a [generator()] or [async()] function: the initial entry and every resumption
+#' after a [yield()] or [await()]. Any `withr::defer()`, `withr::local_*()`, or
+#' [on.exit()] registered while `expr` runs is torn down at the **end of that
+#' step** (the next `yield`/`await`, a `return`, normal completion, an error, or
+#' early close).
+#'
+#' This gives setup/teardown parity for each step of a coroutine, unlike a
+#' top-level `on.exit()` which only fires once when the whole function exits.
+#'
+#' @details
+#' - `expr` runs in an isolated environment: it can read the function's
+#'   arguments, locals, and lexical scope, and mutate external state
+#'   (`the$x <- 1`, `<<-`), but **plain assignments stay local to `expr`** and are
+#'   not visible to the function body.
+#' - When execution reaches a `setup()` call it is registered and runs for the
+#'   current step; re-encountering the *same* `setup()` call (e.g. on a later loop
+#'   iteration) is a no-op. A `setup()` inside an `if`/loop registers only if and
+#'   when reached.
+#' - Multiple `setup()` calls stack and run in registration order each step; their
+#'   teardowns fire in reverse order.
+#' - `setup()` cannot contain `yield()`/`await()` and its result cannot be
+#'   assigned.
+#'
+#' Like [yield()] and [await()], `setup()` is a syntactic construct recognised by
+#' the coroutine compiler. Calling it directly (outside a coroutine body) is an
+#' error.
+#'
+#' @param expr An expression to run at the start of each step.
+#'
+#' @seealso [generator()], [async()], [yield()], [await()].
+#' @examples
+#' the <- new.env()
+#' the$x <- 0
+#'
+#' gen <- generator(function() {
+#'   setup({
+#'     old_x <- the$x
+#'     the$x <- 1
+#'     withr::defer(the$x <- old_x)
+#'   })
+#'   yield(the$x)   # 1 while the step runs
+#'   yield(the$x)   # 1 again: setup re-ran for this step
+#' })
+#'
+#' g <- gen()
+#' g()        # 1
+#' the$x      # 0 — restored at the end of the step
+#' g()        # 1
+#' the$x      # 0
+#' @export
+setup <- function(expr) {
+  abort("`setup()` can't be called directly or within function arguments.")
+}

_pkgdown.yml

@@ -18,6 +18,7 @@ reference:
   - generator
   - gen
   - yield
+  - setup
   - loop
   - collect
 

tests/testthat/_snaps/setup.md

@@ -0,0 +1,41 @@
+# setup() compiles to a do_setup() state
+
+    Code
+      generator_body(function() {
+        setup({
+          old <- the$x
+          withr::defer(the$x <- old)
+        })
+        yield(1)
+      })
+    Output
+      {
+          if (exhausted) {
+              return(invisible(exhausted()))
+          }
+          repeat {
+              switch(state[[1L]], `1` = {
+                  do_setup(1L, quote({
+                      old <- the$x
+                      withr::defer(the$x <- old)
+                  }))
+                  state[[1L]] <- 2L
+              }, `2` = {
+                  user({
+                      1
+                  })
+                  state[[1L]] <- 3L
+                  suspend()
+                  return(last_value())
+              }, `3` = {
+                  .last_value <- if (missing(arg)) exhausted() else arg
+                  state[[1L]] <- 4L
+              }, `4` = {
+                  exhausted <- TRUE
+                  return(exhausted())
+              })
+          }
+          exhausted <- TRUE
+          invisible(exhausted())
+      }
+