Simplify legend layout logic (#5648)

* extract title logic in `legend_add_title()`

* `guide_custom()` uses title logic

* simplify legend layout

* borrow legend's setup_elements method

* make `add_title` a ggproto method

* fix messy merge conflict

* document `add_title` method

* `unique()` --> `unique0()`

Author: teunbrand

DESCRIPTION

@@ -79,7 +79,7 @@ Config/testthat/edition: 3
 Encoding: UTF-8
 LazyData: true
 Roxygen: list(markdown = TRUE)
-RoxygenNote: 7.2.3
+RoxygenNote: 7.3.0
 Collate: 
     'ggproto.R'
     'ggplot-global.R'

NAMESPACE

@@ -718,7 +718,10 @@ import(scales)
 import(vctrs)
 importFrom(glue,glue)
 importFrom(glue,glue_collapse)
+importFrom(grid,arrow)
+importFrom(grid,unit)
 importFrom(lifecycle,deprecated)
+importFrom(scales,alpha)
 importFrom(stats,setNames)
 importFrom(tibble,tibble)
 importFrom(utils,.DollarNames)

R/guide-.R

@@ -154,6 +154,9 @@ new_guide <- function(..., available_aes = "any", super) {
 #'   methods, the measurements from `measure_grobs()` and layout from
 #'   `arrange_layout()` to finalise the guide.
 #'
+#' - `add_title` Adds the title to a gtable, taking into account the size
+#'   of the title as well as the gtable size.
+#'
 #' @rdname ggplot2-ggproto
 #' @format NULL
 #' @usage NULL
@@ -416,6 +419,69 @@ Guide <- ggproto(
 
   draw_early_exit = function(self, params, elements) {
     zeroGrob()
+  },
+
+  add_title = function(gtable, title, position, just) {
+    if (is.zero(title)) {
+      return(gtable)
+    }
+
+    title_width_cm  <- width_cm(title)
+    title_height_cm <- height_cm(title)
+
+    # Add extra row/col for title
+    gtable <- switch(
+      position,
+      top    = gtable_add_rows(gtable, unit(title_height_cm, "cm"), pos =  0),
+      right  = gtable_add_cols(gtable, unit(title_width_cm,  "cm"), pos = -1),
+      bottom = gtable_add_rows(gtable, unit(title_height_cm, "cm"), pos = -1),
+      left   = gtable_add_cols(gtable, unit(title_width_cm,  "cm"), pos =  0)
+    )
+
+    # Add title
+    args <- switch(
+      position,
+      top    = list(t =  1, l =  1, r = -1, b =  1),
+      right  = list(t =  1, l = -1, r = -1, b = -1),
+      bottom = list(t = -1, l =  1, r = -1, b = -1),
+      left   = list(t =  1, l =  1, r =  1, b = -1),
+    )
+    gtable <- inject(gtable_add_grob(
+      x = gtable, grobs = title, !!!args, z = -Inf, name = "title", clip = "off"
+    ))
+
+    if (position %in% c("top", "bottom")) {
+
+      if (any(unitType(gtable$widths) == "null")) {
+        # Don't need to add extra title size for stretchy legends
+        return(gtable)
+      }
+      table_width <- sum(width_cm(gtable$widths))
+      extra_width <- max(0, title_width_cm - table_width)
+      if (extra_width == 0) {
+        return(gtable)
+      }
+      extra_width <- unit((c(1, -1) * just$hjust + c(0, 1)) * extra_width, "cm")
+      gtable <- gtable_add_cols(gtable, extra_width[1], pos =  0)
+      gtable <- gtable_add_cols(gtable, extra_width[2], pos = -1)
+
+    } else {
+
+      if (any(unitType(gtable$heights) == "null")) {
+        # Don't need to add extra title size for stretchy legends
+        return(gtable)
+      }
+      table_height <- sum(height_cm(gtable$heights))
+      extra_height <- max(0, title_height_cm - table_height)
+      if (extra_height == 0) {
+        return(gtable)
+      }
+      extra_height <- unit((c(-1, 1) * just$vjust + c(1, 0)) * extra_height, "cm")
+      gtable <- gtable_add_rows(gtable, extra_height[1], pos =  0)
+      gtable <- gtable_add_rows(gtable, extra_height[2], pos = -1)
+    }
+
+    gtable
   }
 )
 

R/guide-custom.R

@@ -90,89 +90,37 @@ GuideCustom <- ggproto(
     params
   },
 
-  setup_elements = function(params, elements, theme) {
-    theme <- add_theme(theme, params$theme)
-    title_position <- theme$legend.title.position %||% switch(
-      params$direction, vertical = "top", horizontal = "left"
-    )
-    title_position <- arg_match0(
-      title_position, .trbl, arg_nm = "legend.title.position"
-    )
-    theme$legend.title.position <- title_position
-    theme$legend.key.spacing <- theme$legend.key.spacing %||% unit(5.5, "pt")
-    gap <- calc_element("legend.key.spacing", theme)
-
-    margin <- calc_element("text", theme)$margin
-    title <- theme(text = element_text(
-      hjust = 0, vjust = 0.5,
-      margin = position_margin(title_position, margin, gap)
-    ))
-    elements$title <- calc_element("legend.title", add_theme(theme, title))
-    Guide$setup_elements(params, elements, theme)
-  },
-
   draw = function(self, theme, position = NULL, direction = NULL,
                   params = self$params) {
 
-    if (is.zero(params$grob)) {
-      return(zeroGrob())
+    # Render title
+    params <- replace_null(params, position = position, direction = direction)
+    elems <- GuideLegend$setup_elements(params, self$elements, theme)
+    if (!is.waive(params$title) && !is.null(params$title)) {
+      title <- self$build_title(params$title, elems, params)
+    } else {
+      title <- zeroGrob()
     }
 
-    # Render title
-    params$direction <- params$direction %||% direction
-    elems <- self$setup_elements(params, self$elements, theme)
-    elems <- self$override_elements(params, elems, theme)
+    title_position <- elems$title_position
 
     # Start with putting the main grob in a gtable
     width  <- convertWidth(params$width, "cm", valueOnly = TRUE)
     height <- convertHeight(params$height, "cm", valueOnly = TRUE)
     gt <- gtable(widths = unit(width, "cm"), heights = unit(height, "cm"))
     gt <- gtable_add_grob(gt, params$grob, t = 1, l = 1, clip = "off")
 
-    # Render title
-    if (!is.waive(params$title) && !is.null(params$title)) {
-      title <- self$build_title(params$title, elems, params)
-    } else {
-      title <- zeroGrob()
-    }
 
-    # Add title
-    if (!is.zero(title)) {
-      common_args <- list(name = "title", clip = "off", grobs = title)
-      if (elems$title_position == "top") {
-        gt <- gtable_add_rows(gt, unit(height_cm(title), "cm"), pos = 0)
-        gt <- inject(gtable_add_grob(gt, t = 1, l = 1, !!!common_args))
-      } else if (elems$title_position == "bottom") {
-        gt <- gtable_add_rows(gt, unit(height_cm(title), "cm"), pos = -1)
-        gt <- inject(gtable_add_grob(gt, t = -1, l = 1, !!!common_args))
-      } else if (elems$title_position == "left") {
-        gt <- gtable_add_cols(gt, unit(width_cm(title), "cm"), pos = 0)
-        gt <- inject(gtable_add_grob(gt, t = 1, l = 1, !!!common_args))
-      } else if (elems$title_position == "right") {
-        gt <- gtable_add_cols(gt, unit(width_cm(title), "cm"), pos = -1)
-        gt <- inject(gtable_add_grob(gt, t = 1, l = -1, !!!common_args))
-      }
-
-      # Add extra space for large titles
-      extra_width  <- max(0, width_cm(title) - width)
-      extra_height <- max(0, height_cm(title) - height)
-      just <- with(elems$title, rotate_just(angle, hjust, vjust))
-      hjust <- just$hjust
-      vjust <- just$vjust
-      if (elems$title_position %in% c("top", "bottom")) {
-        gt <- gtable_add_cols(gt, unit(extra_width * hjust, "cm"), pos = 0)
-        gt <- gtable_add_cols(gt, unit(extra_width * (1 - hjust), "cm"), pos = -1)
-      } else {
-        gt <- gtable_add_rows(gt, unit(extra_height * (1 - vjust), "cm"), pos = 0)
-        gt <- gtable_add_rows(gt, unit(extra_height * vjust, "cm"), pos = -1)
-      }
-    }
+    gt <- self$add_title(
+      gt, title, title_position,
+      with(elems$title, rotate_just(angle, hjust, vjust))
+    )
 
     # Add padding and background
     gt <- gtable_add_padding(gt, elems$margin)
-    background <- element_grob(elems$background)
+
     gt <- gtable_add_grob(
-      gt, background,
+      gt, elems$background,
       t = 1, l = 1, r = -1, b = -1,
       z = -Inf, clip = "off"
     )

R/guide-legend.R

@@ -471,43 +471,11 @@ GuideLegend <- ggproto(
     )
     heights <- head(vec_interleave(!!!heights), -1)
 
-    has_title <- !is.zero(grobs$title)
-
-    if (has_title) {
-      # Measure title
-      title_width  <- width_cm(grobs$title)
-      title_height <- height_cm(grobs$title)
-
-      # Titles are assumed to have sufficient size when keys are null units
-      extra_width <-
-        if (isTRUE(elements$stretch_x)) 0 else max(0, title_width - sum(widths))
-      extra_height <-
-        if (isTRUE(elements$stretch_y)) 0 else max(0, title_height - sum(heights))
-
-      just <- with(elements$title, rotate_just(angle, hjust, vjust))
-
-      # Combine title with rest of the sizes based on its position
-      widths <- switch(
-        elements$title_position,
-        "left"  = c(title_width, widths),
-        "right" = c(widths, title_width),
-        c(extra_width * just$hjust, widths, extra_width * (1 - just$hjust))
-      )
-      heights <- switch(
-        elements$title_position,
-        "top"    = c(title_height, heights),
-        "bottom" = c(heights, title_height),
-        c(extra_height * (1 - just$vjust), heights, extra_height * just$vjust)
-      )
-    }
-
     list(
       widths  = widths,
       heights = heights,
       padding = elements$padding,
-      has_title = has_title,
-      label_position = elements$text_position,
-      title_position = elements$title_position
+      label_position = elements$text_position
     )
   },
 
@@ -521,97 +489,40 @@ GuideLegend <- ggproto(
       row <- ceiling(break_seq / dim[2L])
       col <- (break_seq - 1L) %% dim[2L] + 1L
     } else {
-      df <- mat_2_df(arrayInd(break_seq, dim), c("R", "C"))
-      row <- df$R
-      col <- df$C
+      row <- (break_seq - 1L) %% dim[1L] + 1L
+      col <- ceiling(break_seq / dim[1L])
     }
-    # Make spacing for padding / gaps. For example: because first gtable cell
-    # will be padding, first item will be at [2, 2] position. Then the
-    # second item-row will be [4, 2] because [3, 2] will be a gap cell.
-    key_row <- label_row <- row * 2
-    key_col <- label_col <- col * 2
+    # Account for spacing in between keys
+    key_row <- row * 2 - 1
+    key_col <- col * 2 - 1
 
     # Make gaps for key-label spacing depending on label position
-    switch(
-      sizes$label_position,
-      "top" = {
-        key_row   <- key_row + row
-        label_row <- key_row - 1
-      },
-      "bottom" = {
-        key_row   <- key_row + row - 1
-        label_row <- key_row + 1
-      },
-      "left" = {
-        key_col   <- key_col + col
-        label_col <- key_col - 1
-      },
-      "right" = {
-        key_col   <- key_col + col - 1
-        label_col <- key_col + 1
-      }
+    position <- sizes$label_position
+    key_row <- key_row + switch(position, top  = row, bottom = row - 1, 0)
+    lab_row <- key_row + switch(position, top  = -1,  bottom = 1,       0)
+    key_col <- key_col + switch(position, left = col, right  = col - 1, 0)
+    lab_col <- key_col + switch(position, left = -1,  right  = 1,       0)
+
+    data_frame0(
+      key_row = key_row, key_col = key_col,
+      label_row = lab_row, label_col = lab_col
     )
-
-    # Offset layout based on title position
-    if (sizes$has_title) {
-      position <- sizes$title_position
-      if (position != "right") {
-        key_col   <- key_col   + 1
-        label_col <- label_col + 1
-      }
-      if (position != "bottom") {
-        key_row   <- key_row   + 1
-        label_row <- label_row + 1
-      }
-      nrow <- length(sizes$heights)
-      ncol <- length(sizes$widths)
-      title_row <- switch(position, top  = 1, bottom = nrow, seq_len(nrow)) + 1
-      title_col <- switch(position, left = 1, right  = ncol, seq_len(ncol)) + 1
-    } else {
-      title_row <- NA
-      title_col <- NA
-    }
-
-    df <- cbind(df, key_row, key_col, label_row, label_col)
-
-    list(layout = df, title_row = title_row, title_col = title_col)
   },
 
-  assemble_drawing = function(grobs, layout, sizes, params, elements) {
-    widths <- unit(c(sizes$padding[4], sizes$widths, sizes$padding[2]), "cm")
+  assemble_drawing = function(self, grobs, layout, sizes, params, elements) {
+
+    widths <- unit(sizes$widths, "cm")
     if (isTRUE(elements$stretch_x)) {
-      widths[unique(layout$layout$key_col)] <- elements$key_width
+      widths[unique0(layout$key_col)] <- elements$key_width
     }
 
-    heights <- unit(c(sizes$padding[1], sizes$heights, sizes$padding[3]), "cm")
+    heights <- unit(sizes$heights, "cm")
     if (isTRUE(elements$stretch_y)) {
-      heights[unique(layout$layout$key_row)] <- elements$key_height
+      heights[unique0(layout$key_row)] <- elements$key_height
     }
 
     gt <- gtable(widths = widths, heights = heights)
 
-    # Add background
-    if (!is.zero(elements$background)) {
-      gt <- gtable_add_grob(
-        gt, elements$background,
-        name = "background", clip = "off",
-        t = 1, r = -1, b = -1, l =1
-      )
-    }
-
-    # Add title
-    if (!is.zero(grobs$title)) {
-      gt <- gtable_add_grob(
-        gt, grobs$title,
-        name = "title", clip = "off",
-        t = min(layout$title_row), r = max(layout$title_col),
-        b = max(layout$title_row), l = min(layout$title_col)
-      )
-    }
-
-    # Extract appropriate part of layout
-    layout   <- layout$layout
-
     # Add keys
     if (!is.zero(grobs$decor)) {
       n_key_layers <- params$n_key_layers %||% 1L
@@ -640,6 +551,21 @@ GuideLegend <- ggproto(
       )
     }
 
+    gt <- self$add_title(
+      gt, grobs$title, elements$title_position,
+      with(elements$title, rotate_just(angle, hjust, vjust))
+    )
+
+    gt <- gtable_add_padding(gt, unit(elements$padding, "cm"))
+
+    # Add background
+    if (!is.zero(elements$background)) {
+      gt <- gtable_add_grob(
+        gt, elements$background,
+        name = "background", clip = "off",
+        t = 1, r = -1, b = -1, l =1, z = -Inf
+      )
+    }
     gt
   }
 )