[CRIU] [PATCH 1/3] pstree: add for_each_pssubtree_item helper to walk only root's subtree

[Pavel Tikhomirov]

Also add skip_descendants flag to skip diveing in children, to optimize
walk if we do not have to visit some sub-trees. Will use it all in next
patch.

Signed-off-by: Pavel Tikhomirov <ptikhomirov at virtuozzo.com>
---
 criu/include/pstree.h |  3 +++
 criu/pstree.c         | 14 ++++++++++----
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/criu/include/pstree.h b/criu/include/pstree.h
index f441d53..21b6a99 100644
--- a/criu/include/pstree.h
+++ b/criu/include/pstree.h
@@ -93,6 +93,9 @@ extern struct pstree_item *pstree_item_next(struct pstree_item *item);
 #define for_each_pstree_item(pi) \
 	for (pi = root_item; pi != NULL; pi = pstree_item_next(pi))
 
+#define for_each_pssubtree_item(pi, root) \
+	for (pi = root; pi != NULL; pi = pssubtree_item_next(pi, root, false))
+
 extern bool restore_before_setsid(struct pstree_item *child);
 extern int prepare_pstree(void);
 extern int prepare_dummy_pstree(void);
diff --git a/criu/pstree.c b/criu/pstree.c
index 833b3d0..7c74607 100644
--- a/criu/pstree.c
+++ b/criu/pstree.c
@@ -232,13 +232,13 @@ void init_pstree_helper(struct pstree_item *ret)
 	task_entries->nr_helpers++;
 }
 
-/* Deep first search on children */
-struct pstree_item *pstree_item_next(struct pstree_item *item)
+/* Search only root's subtree, possibly skipping descendants */
+struct pstree_item *pssubtree_item_next(struct pstree_item *item, struct pstree_item *root, bool skip_descendants)
 {
-	if (!list_empty(&item->children))
+	if (!skip_descendants && !list_empty(&item->children))
 		return list_first_entry(&item->children, struct pstree_item, sibling);
 
-	while (item->parent) {
+	while (item->parent && item != root) {
 		if (item->sibling.next != &item->parent->children)
 			return list_entry(item->sibling.next, struct pstree_item, sibling);
 		item = item->parent;
@@ -247,6 +247,12 @@ struct pstree_item *pstree_item_next(struct pstree_item *item)
 	return NULL;
 }
 
+/* Deep first search on children */
+struct pstree_item *pstree_item_next(struct pstree_item *item)
+{
+	return pssubtree_item_next(item, NULL, false);
+}
+
 /* Preorder traversal of pstree item */
 int preorder_pstree_traversal(struct pstree_item *item, int (*f)(struct pstree_item *))
 {
-- 
2.9.3