facet_grid() evaluates facets before adding margins (#5944)

* evaluate facets before adding margins

* add test

* add news bullet

Author: teunbrand

NEWS.md

@@ -1,5 +1,7 @@
 # ggplot2 (development version)
 
+* Fixed bug in `facet_grid(margins = TRUE)` when using expresssions 
+  (@teunbrand, #1864).
 * `geom_step()` now supports the `orientation` argument (@teunbrand, #5936).
 * `position_dodge()` and `position_jitterdodge()` now have a `reverse` argument 
   (@teunbrand, #3610)

R/facet-grid-.R

@@ -297,12 +297,21 @@ FacetGrid <- ggproto("FacetGrid", Facet,
       return(data)
     }
 
-    # Compute faceting values and add margins
-    margin_vars <- list(intersect(names(rows), names(data)),
-      intersect(names(cols), names(data)))
-    data <- reshape_add_margins(data, margin_vars, params$margins)
-
+    # Compute faceting values
     facet_vals <- eval_facets(c(rows, cols), data, params$.possible_columns)
+    if (nrow(facet_vals) == nrow(data)) {
+      # Margins are computed on evaluated faceting values (#1864).
+      facet_vals <- reshape_add_margins(
+        # We add an index column to track data recycling
+        vec_cbind(facet_vals, .index = seq_len(nrow(facet_vals))),
+        list(intersect(names(rows), names(facet_vals)),
+             intersect(names(cols), names(facet_vals))),
+        params$margins
+      )
+      # Apply recycling on original data to fit margins
+      data <- vec_slice(data, facet_vals$.index)
+      facet_vals$.index <- NULL
+    }
 
     # If any faceting variables are missing, add them in by
     # duplicating the data

tests/testthat/test-facet-map.R

@@ -22,6 +22,10 @@ test_that("margins add extra data", {
   loc <- panel_map_one(facet_grid(a~b, margins = "b"), df)
 
   expect_equal(nrow(loc), nrow(df) * 2)
+
+  # For variables including computation (#1864)
+  loc <- panel_map_one(facet_grid(a ~ I(b + 1), margins = TRUE), df)
+  expect_equal(nrow(loc), nrow(df) * 4)
 })
 
 test_that("grid: missing facet columns are duplicated", {