From 549fae9738494a7a3536998dd499ed6d54017c95 Mon Sep 17 00:00:00 2001
From: Sergio Oller <sergioller@gmail.com>
Date: Sun, 6 Nov 2022 09:03:42 +0100
Subject: [PATCH 1/3] Add helper to get attributes from ggproto method

---
 R/ggproto.r | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/R/ggproto.r b/R/ggproto.r
index d6b0533f01..f93a6bc6f5 100644
--- a/R/ggproto.r
+++ b/R/ggproto.r
@@ -351,3 +351,7 @@ format.ggproto_method <- function(x, ...) {
 
 # proto2 TODO: better way of getting formals for self$draw
 ggproto_formals <- function(x) formals(environment(x)$f)
+
+ggproto_attr <- function(x, which, default = NULL) {
+  attr(environment(x)$f, which = which, exact = TRUE) %||% default
+}

From 99b7c60d753fd928b39b58fc788bbd9d7c15e4d5 Mon Sep 17 00:00:00 2001
From: Sergio Oller <sergioller@gmail.com>
Date: Sun, 6 Nov 2022 09:29:06 +0100
Subject: [PATCH 2/3] FEAT/PERF: ScaleContinuous accepts palette capability
 `may_return_na`

We can extend palettes with attributes to improve the mapping efficiency.

With this commit a palette may define an attribute `may_return_na` to `FALSE`.
If it does, `ScaleContinuous` will assume the palette may not return missing
values, and it will skip checking for those and replacing them.
---
 R/scale-.r | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/R/scale-.r b/R/scale-.r
index b0065ffd52..b4d69aa5ed 100644
--- a/R/scale-.r
+++ b/R/scale-.r
@@ -609,7 +609,14 @@ ScaleContinuous <- ggproto("ScaleContinuous", Scale,
     pal <- self$palette(uniq)
     scaled <- pal[match(x, uniq)]
 
-    ifelse(!is.na(scaled), scaled, self$na.value)
+    # A specific palette can have as attribute "may_return_NA = FALSE"
+    # If it has such attribute, we will skip the ifelse(!is.na(scaled), ...)
+    pal_may_return_na <- ggproto_attr(self$palette, "may_return_NA", default = TRUE)
+    if (pal_may_return_na) {
+      scaled <- ifelse(!is.na(scaled), scaled, self$na.value)
+    }
+
+    scaled
   },
 
   rescale = function(self, x, limits = self$get_limits(), range = limits) {

From d9e577dcb6386625ccf2cf81dcce16d4fa1d1709 Mon Sep 17 00:00:00 2001
From: Sergio Oller <sergioller@gmail.com>
Date: Sun, 6 Nov 2022 12:00:38 +0100
Subject: [PATCH 3/3] TEST: Palette capability "may_return_na"

---
 tests/testthat/test-scale-colour-continuous.R | 29 +++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/tests/testthat/test-scale-colour-continuous.R b/tests/testthat/test-scale-colour-continuous.R
index 10e3ae4dd5..444b27f9db 100644
--- a/tests/testthat/test-scale-colour-continuous.R
+++ b/tests/testthat/test-scale-colour-continuous.R
@@ -18,3 +18,32 @@ test_that("type argument is checked for proper input", {
     scale_colour_continuous(type = "abc")
   )
 })
+
+test_that("palette with may_return_NA=FALSE works as expected", {
+  sc <- scale_fill_continuous()
+  # A palette that may return NAs, will have NAs replaced by the scale's na.value
+  # by the scale:
+  sc$palette <- structure(
+    function(x) {
+      rep(NA_character_, length(x))
+    },
+    may_return_NA = TRUE
+  )
+  sc$na.value <- "red"
+  nat <- sc$map(0.5, limits = c(0, 1))
+  expect_equal(nat, "red")
+
+  # This palette is lying, because it returns NA even though it says it can't.
+  # The scale will not replace the NA values, leading to further errors.
+  # You should not do this in production, but it helps to test:
+  sc <- scale_fill_continuous()
+  sc$palette <- structure(
+    function(x) {
+      rep(NA_character_, length(x))
+    },
+    may_return_NA = FALSE
+  )
+  sc$na.value <- "red"
+  nat <- sc$map(0.5, limits = c(0, 1))
+  expect_equal(nat, NA_character_)
+})