ag_break() can be caught by the wrong ag_for().

[t-kalinowski]

break and next signaling needs to be refactored.

ag_while() and ag_for() should still setup a condition handler to catch control-flow conditions, even if they dispatched to the primitive. Probably just a simple tryCatch() or withCallingHandlers() should work.

This is because in a non-eager graphing context, a break emitted from an ag_if() branch will not be caught by correct loop, if the loop was dispatched to the primitive and is not being traced onto the graph.

The correct thing to do here is throw an error, telling users that only traced loops can capture traced loop-control-flow-statements, and telling them to either untrace the break or to trace the for.

test_that("`break` is caught by the correct `for`", {

  # This test currently fails
  # ag_for() and ag_while() need to setup condition handlers even
  # if they dispatch to the primative.
  skip_if_no_tensorflow()

  fn <- function(x) {
    i <- x[1]
    r <- 10L
    for(i in x) {
      for(r in 10:13)
        if(i >= 1)
          break
    }
    list(i, r)
  }
  ag_fn <- autograph(fn)
  tf_ag_fn <- tf_function(ag_fn)

  fn(1:3)
  ag_fn(1:3)
  ag_fn(as_tensor(1:3))
  tf_ag_fn(as_tensor(1:3))

  expect_error(ag_fn(as_tensor(1:3)))
  expect_error(tf_ag_fn(as_tensor(1:3)))
})