From 4fbc8575d80e54e1973ac58e979c2300aec21394 Mon Sep 17 00:00:00 2001
From: Teun van den Brand <49372158+teunbrand@users.noreply.github.com>
Date: Mon, 23 Sep 2024 14:01:15 +0200
Subject: [PATCH] Mask `stage()` at the expression level (#6110)

* evaluation helper

* use helper

* generalise `substitute_aes()`

* substitute expressions instead of wrangling data masks

* add test

* make a comment for the next person to trip over this

* add news bullet
---
 NEWS.md                              |  2 ++
 R/aes-evaluation.R                   | 36 ++++++++++++++++++++++++++++
 R/aes.R                              |  7 ++++--
 R/geom-.R                            | 14 ++++-------
 R/layer.R                            | 20 +++++-----------
 tests/testthat/test-aes-calculated.R | 25 +++++++++++++++++++
 6 files changed, 79 insertions(+), 25 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index 8d7f8760d0..f369879868 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,5 +1,7 @@
 # ggplot2 (development version)
 
+* Fixed bug where the `ggplot2::`-prefix did not work with `stage()` 
+  (@teunbrand, #6104).
 * New `get_labs()` function for retrieving completed plot labels 
   (@teunbrand, #6008).
 * Built-in `theme_*()` functions now have `ink` and `paper` arguments to control
diff --git a/R/aes-evaluation.R b/R/aes-evaluation.R
index 8e47ebcd1e..30d0b3f501 100644
--- a/R/aes-evaluation.R
+++ b/R/aes-evaluation.R
@@ -358,3 +358,39 @@ make_labels <- function(mapping) {
   }
   Map(default_label, names(mapping), mapping)
 }
+
+eval_aesthetics <- function(aesthetics, data, mask = NULL) {
+
+  env <- child_env(base_env())
+
+  # Here we mask functions, often to replace `stage()` with context appropriate
+  # functions `stage_calculated()`/`stage_scaled()`.
+  if (length(mask) > 0) {
+    aesthetics <- substitute_aes(aesthetics, mask_function, mask = mask)
+  }
+
+  evaled <- lapply(aesthetics, eval_tidy, data = data, env = env)
+  names(evaled) <- names(aesthetics)
+  compact(rename_aes(evaled))
+}
+
+# `mask` is a list of functions where `names(mask)` indicate names of functions
+# that need to be replaced, and `mask[[i]]` is the function to replace it
+# with.
+mask_function <- function(x, mask) {
+  if (!is.call(x)) {
+    return(x)
+  }
+  nms <- names(mask)
+  x[-1] <- lapply(x[-1], mask_function, mask = mask)
+  if (!is_call(x, nms)) {
+    return(x)
+  }
+  for (nm in nms) {
+    if (is_call(x, nm)) {
+      x[[1]] <- mask[[nm]]
+      return(x)
+    }
+  }
+}
+
diff --git a/R/aes.R b/R/aes.R
index d739289b0e..0829f64daf 100644
--- a/R/aes.R
+++ b/R/aes.R
@@ -198,9 +198,12 @@ rename_aes <- function(x) {
   }
   x
 }
-substitute_aes <- function(x) {
+
+# `x` is assumed to be a strict list of quosures;
+# it should have no non-quosure constants in it, even though `aes()` allows it.
+substitute_aes <- function(x, fun = standardise_aes_symbols, ...) {
   x <- lapply(x, function(aesthetic) {
-    as_quosure(standardise_aes_symbols(quo_get_expr(aesthetic)), env = environment(aesthetic))
+    as_quosure(fun(quo_get_expr(aesthetic), ...), env = environment(aesthetic))
   })
   class(x) <- "uneval"
   x
diff --git a/R/geom-.R b/R/geom-.R
index c5a1ab275d..bb409ba0cf 100644
--- a/R/geom-.R
+++ b/R/geom-.R
@@ -157,12 +157,10 @@ 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)
+      modified_aes <- eval_aesthetics(
+        substitute_aes(modifiers), data,
+        mask = list(stage = stage_scaled)
+      )
 
       # Check that all output are valid data
       nondata_modified <- check_nondata_cols(modified_aes)
@@ -177,11 +175,9 @@ Geom <- ggproto("Geom",
         ))
       }
 
-      names(modified_aes) <- names(rename_aes(modifiers))
-
       modified_aes <- cleanup_mismatched_data(modified_aes, nrow(data), "after_scale")
 
-      modified_aes <- data_frame0(!!!compact(modified_aes))
+      modified_aes <- data_frame0(!!!modified_aes)
 
       data <- data_frame0(!!!defaults(modified_aes, data))
     }
diff --git a/R/layer.R b/R/layer.R
index 3fd89cf3f7..b10c230e1d 100644
--- a/R/layer.R
+++ b/R/layer.R
@@ -307,10 +307,7 @@ Layer <- ggproto("Layer", NULL,
     }
 
     # Evaluate aesthetics
-    env <- child_env(baseenv(), stage = stage)
-    evaled <- lapply(aesthetics, eval_tidy, data = data, env = env)
-    evaled <- compact(evaled)
-
+    evaled <- eval_aesthetics(aesthetics, data)
     plot$scales$add_defaults(evaled, plot$plot_env)
 
     # Check for discouraged usage in mapping
@@ -390,14 +387,10 @@ 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)
-
-    new <- substitute_aes(new)
-    stat_data <- lapply(new, eval_tidy, mask, env)
-
+    stat_data <- eval_aesthetics(
+      substitute_aes(new), data_orig,
+      mask = list(stage = stage_calculated)
+    )
     # Check that all columns in aesthetic stats are valid data
     nondata_stat_cols <- check_nondata_cols(stat_data)
     if (length(nondata_stat_cols) > 0) {
@@ -411,8 +404,7 @@ Layer <- ggproto("Layer", NULL,
       ))
     }
 
-    names(stat_data) <- names(new)
-    stat_data <- data_frame0(!!!compact(stat_data))
+    stat_data <- data_frame0(!!!stat_data)
 
     # Add any new scales, if needed
     plot$scales$add_defaults(stat_data, plot$plot_env)
diff --git a/tests/testthat/test-aes-calculated.R b/tests/testthat/test-aes-calculated.R
index b453af02f5..3ac8e06dbe 100644
--- a/tests/testthat/test-aes-calculated.R
+++ b/tests/testthat/test-aes-calculated.R
@@ -99,3 +99,28 @@ test_that("A deprecated warning is issued when stat(var) or ..var.. is used", {
   p2 <- ggplot(NULL, aes(..bar..))
   expect_snapshot_warning(b2 <- ggplot_build(p2))
 })
+
+test_that("functions can be masked", {
+
+  foo <- function(x) x + 10
+  bar <- function(x) x * 10
+
+  data <- data.frame(val = 10)
+  mapping <- aes(x = val, y = foo(20))
+
+  evaled  <- eval_aesthetics(mapping, data = data, mask = list())
+  expect_equal(evaled, list(x = 10, y = 30))
+
+  evaled <- eval_aesthetics(mapping, data = data, mask = list(foo = bar))
+  expect_equal(evaled, list(x = 10, y = 200))
+
+  # Test namespace-prefixed evaluation (#6104)
+  mapping <- aes(x = val, y = ggplot2::stage(10, 20, 30))
+  evaled <- eval_aesthetics(mapping, data = data, mask = list())
+  expect_equal(evaled, list(x = 10, y = 10))
+  evaled <- eval_aesthetics(mapping, data = data, mask = list(stage = stage_calculated))
+  expect_equal(evaled, list(x = 10, y = 20))
+  evaled <- eval_aesthetics(mapping, data = data, mask = list(stage = stage_scaled))
+  expect_equal(evaled, list(x = 10, y = 30))
+
+})