From 15ffbd50d60e2e1fd91f0ab891384f44ae8a17cb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kirill=20M=C3=BCller?= <krlmlr+r@mailbox.org>
Date: Sun, 21 Aug 2022 07:52:54 +0200
Subject: [PATCH] pre_visit_one() and post_visit_one()

---
 R/initialize.R                   |  2 +-
 R/nested-to-tree.R               |  6 ++---
 R/relevel.R                      |  6 ++---
 R/transform-block.R              |  2 +-
 R/visit.R                        | 39 ++++++++++++++++++++++++++++++++
 vignettes/customizing_styler.Rmd |  2 +-
 6 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/R/initialize.R b/R/initialize.R
index 281f22d2e..88eba5b9e 100644
--- a/R/initialize.R
+++ b/R/initialize.R
@@ -10,7 +10,7 @@
 #'   {
 #'     string_to_format <- "call( 3)"
 #'     pd <- styler:::compute_parse_data_nested(string_to_format)
-#'     styler:::pre_visit(pd, c(default_style_guide_attributes))
+#'     styler:::pre_visit_one(pd, default_style_guide_attributes)
 #'   }
 #' )
 #' @export
diff --git a/R/nested-to-tree.R b/R/nested-to-tree.R
index d43f48fb9..b6b3c30d7 100644
--- a/R/nested-to-tree.R
+++ b/R/nested-to-tree.R
@@ -8,7 +8,7 @@
 #' @keywords internal
 create_tree <- function(text, structure_only = FALSE) {
   compute_parse_data_nested(text, transformers = NULL) %>%
-    pre_visit(c(default_style_guide_attributes)) %>%
+    pre_visit_one(default_style_guide_attributes) %>%
     create_tree_from_pd_with_default_style_attributes(structure_only)
 }
 
@@ -36,8 +36,8 @@ create_tree_from_pd_with_default_style_attributes <- function(pd,
 #'     {
 #'       code <- "a <- function(x) { if(x > 1) { 1+1 } else {x} }"
 #'       nested_pd <- styler:::compute_parse_data_nested(code)
-#'       initialized <- styler:::pre_visit(
-#'         nested_pd, c(default_style_guide_attributes)
+#'       initialized <- styler:::pre_visit_one(
+#'         nested_pd, default_style_guide_attributes
 #'       )
 #'       styler:::create_node_from_nested_root(initialized,
 #'         structure_only = FALSE
diff --git a/R/relevel.R b/R/relevel.R
index 92ed14965..c8984a5d1 100644
--- a/R/relevel.R
+++ b/R/relevel.R
@@ -12,8 +12,7 @@
 #' @param pd_nested A nested parse table to partially flatten.
 #' @keywords internal
 flatten_operators <- function(pd_nested) {
-  pd_nested %>%
-    post_visit(c(flatten_operators_one))
+  post_visit_one(pd_nested, flatten_operators_one)
 }
 
 #' Flatten one level of nesting with its child
@@ -144,8 +143,7 @@ wrap_expr_in_expr <- function(pd) {
 #' @keywords internal
 relocate_eq_assign <- function(pd) {
   if (parser_version_get() < 2) {
-    pd %>%
-      post_visit(c(relocate_eq_assign_nest))
+    post_visit_one(pd, relocate_eq_assign_nest)
   } else {
     pd
   }
diff --git a/R/transform-block.R b/R/transform-block.R
index a188c4a4e..6c9c06e73 100644
--- a/R/transform-block.R
+++ b/R/transform-block.R
@@ -28,7 +28,7 @@ parse_transform_serialize_r_block <- function(pd_nested,
                                               base_indention) {
   if (!all(pd_nested$is_cached, na.rm = TRUE) || !cache_is_activated()) {
     transformed_pd <- apply_transformers(pd_nested, transformers)
-    flattened_pd <- post_visit(transformed_pd, list(extract_terminals)) %>%
+    flattened_pd <- post_visit_one(transformed_pd, extract_terminals) %>%
       enrich_terminals(transformers$use_raw_indention) %>%
       apply_ref_indention() %>%
       set_regex_indention(
diff --git a/R/visit.R b/R/visit.R
index 9206af2f1..a05ff54f8 100644
--- a/R/visit.R
+++ b/R/visit.R
@@ -36,6 +36,25 @@ pre_visit <- function(pd_nested, funs) {
   pd_nested
 }
 
+#' @rdname visit
+#' @keywords internal
+pre_visit_one <- function(pd_nested, fun) {
+  if (is.null(pd_nested)) {
+    return()
+  }
+  pd_nested <- fun(pd_nested)
+
+  children <- pd_nested$child
+  for (i in seq_along(children)) {
+    child <- children[[i]]
+    if (!is.null(child)) {
+      children[[i]] <- pre_visit_one(child, fun)
+    }
+  }
+  pd_nested$child <- children
+  pd_nested
+}
+
 #' @rdname visit
 #' @keywords internal
 post_visit <- function(pd_nested, funs) {
@@ -58,6 +77,26 @@ post_visit <- function(pd_nested, funs) {
   visit_one(pd_nested, funs)
 }
 
+#' @rdname visit
+#' @keywords internal
+post_visit_one <- function(pd_nested, fun) {
+  if (is.null(pd_nested)) {
+    return()
+  }
+  force(fun)
+
+  children <- pd_nested$child
+  for (i in seq_along(children)) {
+    child <- children[[i]]
+    if (!is.null(child)) {
+      children[[i]] <- post_visit_one(child, fun)
+    }
+  }
+  pd_nested$child <- children
+
+  fun(pd_nested)
+}
+
 #' Transform a flat parse table with a list of transformers
 #'
 #' Uses [Reduce()] to apply each function of `funs` sequentially to
diff --git a/vignettes/customizing_styler.Rmd b/vignettes/customizing_styler.Rmd
index 7ec808c8b..ccf5d519c 100644
--- a/vignettes/customizing_styler.Rmd
+++ b/vignettes/customizing_styler.Rmd
@@ -40,7 +40,7 @@ As the name says, this function removes spaces after the opening parenthesis. Bu
 ```{r}
 string_to_format <- "call( 3)"
 pd <- styler:::compute_parse_data_nested(string_to_format) %>%
-  styler:::pre_visit(c(default_style_guide_attributes))
+  styler:::pre_visit_one(default_style_guide_attributes)
 pd$child[[1]] %>%
   select(token, terminal, text, newlines, spaces)
 ```