Allow for early termination of the visitor

R/visit.R

@@ -1,29 +1,40 @@
 #' Visits the set of nodes and all of their children, invoking the callback for each visited node
 #'
 #' @param nodes The list or array of nodes to visit
-#' @param callback The callback function to invoke for each node
+#' @param callback The callback function to invoke for each node. The callback should return FALSE to stop visiting the children of the node, or anything else to continue.
+#'
+#' @return FALSE if the visitor was stopped by the callback.
 #'
 #' @export
 visit_nodes <- function(nodes, callback) {
   if (!is.null(nodes)) {
     for (node in nodes) {
-      visit_node(node, callback)
+      res <- visit_node(node, callback)
+      if (isFALSE(res)) {
+        return(FALSE)
+      }
     }
   }
 }
 
 #' Visits the given node and all of their children, invoking the callback for each visited node
 #'
 #' @param node The node to visit
-#' @param callback The callback function to invoke for each node
+#' @param callback The callback function to invoke for each node. The callback should return FALSE to stop visiting the children of the node, or anything else to continue.
+#'
+#' @return FALSE if the visitor was stopped by the callback.
 #'
 #' @export
 visit_node <- function(node, callback) {
   if (is.null(node)) {
     return()
   }
 
-  callback(node)
+  res <- callback(node)
+  # Exit early if the callback returns FALSE
+  if (isFALSE(res)) {
+    return(FALSE)
+  }
 
   # same logic as the builtin visitor (while explicitly specifying if an entry is a single node or a list)
   # https://github.com/Code-Inspect/flowr/blob/main/src/r-bridge/lang-4.x/ast/model/processing/visitor.ts#L22