From 40b59b205b2c773d822dc5632f64186e3d29442f Mon Sep 17 00:00:00 2001
From: Hadley Wickham <h.wickham@gmail.com>
Date: Fri, 9 Sep 2022 12:18:40 -0500
Subject: [PATCH] Use vec_is_list() in map_depth()

Fixes #920
---
 NEWS.md       |  3 +++
 R/map-if-at.R | 43 ++++++++++++++++++++++---------------------
 man/map_if.Rd |  9 +++++++++
 3 files changed, 34 insertions(+), 21 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index 754f08da..02d62419 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -51,6 +51,9 @@
 
 ## Features and fixes
 
+* `map_depth()` now uses `vec_is_list()` to determine if there's more depth
+  to recurse into, as opposed to `!is_atomic(.x)` (#920).
+
 * New `list_update()` which is similar to `list_modify()` but doesn't work
   recursively (#822).
 
diff --git a/R/map-if-at.R b/R/map-if-at.R
index 1f4ea758..80b9d3e2 100644
--- a/R/map-if-at.R
+++ b/R/map-if-at.R
@@ -95,6 +95,15 @@ map_at <- function(.x, .at, .f, ...) {
 #'
 #' # Equivalent to:
 #' map(x, map, paste, collapse = "/")
+#'
+#' # When ragged is TRUE, `.f()` will also be passed leaves at depth < `.depth`
+#' x <- list(1, list(1, list(1, list(1, 1))))
+#' str(x)
+#' str(map_depth(x, 4, ~ length(unlist(.x)), .ragged = TRUE))
+#' str(map_depth(x, 3, ~ length(unlist(.x)), .ragged = TRUE))
+#' str(map_depth(x, 2, ~ length(unlist(.x)), .ragged = TRUE))
+#' str(map_depth(x, 1, ~ length(unlist(.x)), .ragged = TRUE))
+#' str(map_depth(x, 0, ~ length(unlist(.x)), .ragged = TRUE))
 #' @export
 map_depth <- function(.x, .depth, .f, ..., .ragged = FALSE) {
   if (!is_integerish(.depth, n = 1, finite = TRUE)) {
@@ -116,28 +125,20 @@ map_depth_rec <- function(.x,
                           .atomic) {
   if (.depth < 0) {
     abort("Invalid depth")
-  }
-
-  if (.atomic) {
-    if (!.ragged) {
-      abort("List not deep enough")
+  } else if (.depth == 0) {
+    .f(.x, ...)
+  } else if (.depth == 1) {
+    map(.x, .f, ...)
+  } else {
+    if (vec_is_list(.x)) {
+      map(.x, map_depth_rec, .depth - 1, .f, ..., .ragged = .ragged)
+    } else {
+      if (.ragged) {
+        map(.x, .f, ...)
+      } else {
+        abort("List not deep enough")
+      }
     }
-    return(map(.x, .f, ...))
   }
 
-  if (.depth == 0) {
-    return(.f(.x, ...))
-  }
-
-  if (.depth == 1) {
-    return(map(.x, .f, ...))
-  }
-
-  # Should this be replaced with a generic way of figuring out atomic
-  # types?
-  .atomic <- is_atomic(.x)
-
-  map(.x, function(x) {
-    map_depth_rec(x, .depth - 1, .f, ..., .ragged = .ragged, .atomic = .atomic)
-  })
 }
diff --git a/man/map_if.Rd b/man/map_if.Rd
index 21f14b23..f34b227d 100644
--- a/man/map_if.Rd
+++ b/man/map_if.Rd
@@ -100,6 +100,15 @@ map_depth(x, 2, paste, collapse = "/")
 
 # Equivalent to:
 map(x, map, paste, collapse = "/")
+
+# When ragged is TRUE, `.f()` will also be passed leaves at depth < `.depth`
+x <- list(1, list(1, list(1, list(1, 1))))
+str(x)
+str(map_depth(x, 4, ~ length(unlist(.x)), .ragged = TRUE))
+str(map_depth(x, 3, ~ length(unlist(.x)), .ragged = TRUE))
+str(map_depth(x, 2, ~ length(unlist(.x)), .ragged = TRUE))
+str(map_depth(x, 1, ~ length(unlist(.x)), .ragged = TRUE))
+str(map_depth(x, 0, ~ length(unlist(.x)), .ragged = TRUE))
 }
 \seealso{
 Other map variants: