From b9d289f0d915589ef834839af6e41722377d99f2 Mon Sep 17 00:00:00 2001
From: Teun van den Brand <tahvdbrand@gmail.com>
Date: Mon, 16 Sep 2024 15:51:46 +0200
Subject: [PATCH] simplify stage masking

---
 R/aes-evaluation.R | 24 +++++++++++++++++-------
 R/geom-.R          |  9 +++------
 R/layer.R          | 13 +++++--------
 3 files changed, 25 insertions(+), 21 deletions(-)

diff --git a/R/aes-evaluation.R b/R/aes-evaluation.R
index 5554b54772..a791f40635 100644
--- a/R/aes-evaluation.R
+++ b/R/aes-evaluation.R
@@ -206,16 +206,26 @@ from_theme <- function(x) {
   x
 }
 
+get_stage <- local({
+  stage <- "start"
+  function() {stage}
+})
+
+set_stage <- function(new_stage) {
+  old_stage <- environment(get_stage)$stage
+  environment(get_stage)$stage <- new_stage
+  invisible(old_stage)
+}
+
 #' @rdname aes_eval
 #' @export
 stage <- function(start = NULL, after_stat = NULL, after_scale = NULL) {
-  start
-}
-stage_calculated <- function(start = NULL, after_stat = NULL, after_scale = NULL) {
-  after_stat
-}
-stage_scaled <- function(start = NULL, after_stat = NULL, after_scale = NULL) {
-  after_scale
+  switch(
+    get_stage(),
+    after_stat = after_stat,
+    after_scale = after_scale,
+    start
+  )
 }
 
 # Regex to determine if an identifier refers to a calculated aesthetic
diff --git a/R/geom-.R b/R/geom-.R
index 5b6a2af09d..ca9ecd7408 100644
--- a/R/geom-.R
+++ b/R/geom-.R
@@ -157,12 +157,9 @@ Geom <- ggproto("Geom",
     # This order means that they will have access to all default aesthetics
     if (length(modifiers) != 0) {
       # Set up evaluation environment
-      env <- child_env(baseenv(), after_scale = after_scale)
-      # Mask stage with stage_scaled so it returns the correct expression
-      stage_mask <- child_env(emptyenv(), stage = stage_scaled)
-      mask <- new_data_mask(as_environment(data, stage_mask), stage_mask)
-      mask$.data <- as_data_pronoun(mask)
-      modified_aes <- lapply(substitute_aes(modifiers),  eval_tidy, mask, env)
+      set_stage("after_scale")
+      env <- child_env(baseenv())
+      modified_aes <- lapply(substitute_aes(modifiers),  eval_tidy, data, env)
 
       # Check that all output are valid data
       nondata_modified <- check_nondata_cols(modified_aes)
diff --git a/R/layer.R b/R/layer.R
index 8acb438c9e..3dc3202b00 100644
--- a/R/layer.R
+++ b/R/layer.R
@@ -58,8 +58,8 @@
 #'   `NA`, the default, includes if any aesthetics are mapped.
 #'   `FALSE` never includes, and `TRUE` always includes.
 #'   It can also be a named logical vector to finely select the aesthetics to
-#'   display. To include legend keys for all levels, even 
-#'   when no data exists, use `TRUE`.  If `NA`, all levels are shown in legend, 
+#'   display. To include legend keys for all levels, even
+#'   when no data exists, use `TRUE`.  If `NA`, all levels are shown in legend,
 #'   but unobserved levels are omitted.
 #' @param inherit.aes If `FALSE`, overrides the default aesthetics,
 #'   rather than combining with them. This is most useful for helper functions
@@ -303,6 +303,7 @@ Layer <- ggproto("Layer", NULL,
     }
 
     # Evaluate aesthetics
+    set_stage("start")
     env <- child_env(baseenv(), stage = stage)
     evaled <- lapply(aesthetics, eval_tidy, data = data, env = env)
     evaled <- compact(evaled)
@@ -386,13 +387,9 @@ Layer <- ggproto("Layer", NULL,
     data_orig <- plot$scales$backtransform_df(data)
 
     # Add map stat output to aesthetics
-    env <- child_env(baseenv(), stat = stat, after_stat = after_stat)
-    stage_mask <- child_env(emptyenv(), stage = stage_calculated)
-    mask <- new_data_mask(as_environment(data_orig, stage_mask), stage_mask)
-    mask$.data <- as_data_pronoun(mask)
-
+    set_stage("after_stat")
     new <- substitute_aes(new)
-    stat_data <- lapply(new, eval_tidy, mask, env)
+    stat_data <- lapply(new, eval_tidy, data_orig, env)
 
     # Check that all columns in aesthetic stats are valid data
     nondata_stat_cols <- check_nondata_cols(stat_data)