/
traverse.R
72 lines (60 loc) · 1.97 KB
/
traverse.R
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
#' Apply a function hierarchically to a forest
#'
#' Apply a function hierarchically to a forest in the climbing or descending direction.
#'
#' @param .x A forest
#' @param .f A function, formula, or vector (not necessarily atomic).
#' @param ... Additional arguments passed on to the mapped function.
#' @param .climb Climbing or descending?
#'
#' @return A forest.
#'
#' @export
traverse <- function(.x, .f, ...,
.climb = FALSE) {
.f <- purrr::as_mapper(.f, ...)
.x$nodes <- traverse_impl(.x$nodes, .f,
.climb = .climb)
.x
}
traverse_impl <- function(nodes, .f,
.climb = FALSE) {
node_names <- nodes$.$name
node_parents <- nodes$.$parent
node_data <- drop_node(nodes)
grps <- vec_group_loc(node_parents)
grps <- vec_slice(grps, !vec_detect_missing(grps$key))
grps <- vec_slice(grps,
vec_order(grps$key,
direction = if (.climb) "desc" else "asc"))
rle <- vec_group_rle(vec_slice(node_names, grps$key))
sizes_rle <- field(rle, "length")
inits_rle <- cumsum(sizes_rle) - sizes_rle
loc <- vec_seq_along(sizes_rle)
for (i in loc) {
size_rle <- sizes_rle[[i]]
rle_locs <- seq_len(size_rle)
grp <- vec_slice(grps, inits_rle[[i]] + rle_locs)
grp_parent <- grp$key
grp_children <- grp$loc
parents <- vec_slice(node_data, grp_parent)
parents <- vec_chop(parents)
children <- vec_chop(node_data, grp_children)
new_node_data <- vec_init(list_of(.ptype = node_data), size_rle)
for (j in rle_locs) {
if (.climb) {
new_node_data[[j]] <- .f(children[[j]], parents[[j]])
} else {
new_node_data[[j]] <- .f(parents[[j]], children[[j]])
}
}
new_node_data <- rbind_check(!!!new_node_data)
if (.climb) {
vec_slice(node_data, vec_c(!!!grp_children)) <- new_node_data
} else {
vec_slice(node_data, grp_parent) <- new_node_data
}
}
nodes[-1L] <- node_data
nodes
}