Fix hoisting in unbraced for-loop bodies

Braced loop bodies give each statement its own hoist target, but a single-expression for-loop body inherited the target of the enclosing statement. Body-local setup could therefore be emitted before the loop, where it executed only once and could reference an uninitialized loop variable.

Give both index- and value-iteration bodies a fresh hoist target. Iterable setup remains outside the loop, while guards and temporaries required by the body are emitted inside it.

Author: mns-nordicals

R/r2f-control-flow.R

@@ -77,7 +77,7 @@ r2f_handlers[["while"]] <- function(args, scope, ...) {
 }
 
 # ---- for ----
-r2f_handlers[["for"]] <- function(args, scope, ...) {
+r2f_handlers[["for"]] <- function(args, scope, ..., hoist = NULL) {
   .[var, iterable, body] <- args
   stopifnot(is.symbol(var))
   var <- as.character(var)
@@ -174,7 +174,10 @@ r2f_handlers[["for"]] <- function(args, scope, ...) {
       previous_openmp <- enter_openmp_scope(scope)
       on.exit(exit_openmp_scope(scope, previous_openmp), add = TRUE)
     }
-    body <- r2f(body, scope, ...)
+    # The body is a distinct execution region and needs its own hoist target.
+    # Otherwise a single-expression body reuses the enclosing statement's
+    # target and emits loop-dependent setup before the loop.
+    body <- r2f(body, scope, ..., hoist = NULL)
     check_pending_parallel_consumed(scope)
     loop_stmts <- str_flatten_lines(glue("{var_name} = {element_expr}"), body)
 
@@ -214,12 +217,14 @@ r2f_handlers[["for"]] <- function(args, scope, ...) {
   }
   scope[[var]] <- loop_var
 
-  iterable <- r2f_for_iterable(iterable, scope, ...)
+  iterable <- r2f_for_iterable(iterable, scope, ..., hoist = hoist)
   if (!is.null(parallel)) {
     previous_openmp <- enter_openmp_scope(scope)
     on.exit(exit_openmp_scope(scope, previous_openmp), add = TRUE)
   }
-  body <- r2f(body, scope, ...)
+  # See the value-iteration path above: body-local setup must run inside the
+  # loop even when the R body is not wrapped in braces.
+  body <- r2f(body, scope, ..., hoist = NULL)
   check_pending_parallel_consumed(scope)
 
   directives <- openmp_directives(parallel)

tests/testthat/test-for-iterables.R

@@ -288,3 +288,58 @@ test_that("for() supports rev() for value iteration", {
   expect_quick_identical(rev_values, 1:3, c(1L, 2L, 3L, 4L))
   expect_quick_identical(rev_matrix_values, matrix(1:6, nrow = 2))
 })
+
+test_that("unbraced for() bodies keep hoisted setup inside index loops", {
+  rolling_sum <- function(x, weights) {
+    declare(type(x = double(NA)), type(weights = double(NA)))
+    out <- double(length(x) - length(weights) + 1L)
+    n <- length(weights)
+    # fmt: skip
+    for (i in seq_along(out))
+      out[i] <- sum(x[i:(i + n - 1L)] * weights)
+    out
+  }
+
+  seq_lengths <- function(k) {
+    declare(type(k = integer(1)))
+    out <- 0L
+    # fmt: skip
+    for (i in 1:3)
+      out <- out + length(seq(i, i + 3L, by = k))
+    out
+  }
+
+  expect_quick_identical(
+    rolling_sum,
+    list(as.double(1:8), c(0.25, 0.75))
+  )
+  expect_quick_identical(seq_lengths, 1L, 2L)
+
+  qseq_lengths := quick(seq_lengths)
+  expect_error(
+    qseq_lengths(0L),
+    "invalid '(to - from)/by'",
+    fixed = TRUE
+  )
+  expect_error(
+    qseq_lengths(-1L),
+    "wrong sign in 'by' argument",
+    fixed = TRUE
+  )
+})
+
+test_that("unbraced for() bodies keep hoisted setup inside value loops", {
+  sum_floors <- function(starts) {
+    declare(type(starts = integer(NA)))
+    out <- 0
+    # fmt: skip
+    for (i in starts)
+      out <- out + floor(i + 0.5)
+    out
+  }
+
+  expect_quick_identical(
+    sum_floors,
+    c(1L, 3L, 5L)
+  )
+})