From ec8bbbc9c1db349e84aa0d38593aed85dbd6e552 Mon Sep 17 00:00:00 2001
From: Sven Verdoolaege <skimo@kotnet.org>
Date: Sun, 20 Oct 2013 15:02:59 +0200
Subject: [PATCH] add isl_schedule_node_map_descendant

Signed-off-by: Sven Verdoolaege <skimo@kotnet.org>
---
 doc/user.pod                | 18 ++++++++++++++
 include/isl/schedule_node.h |  4 ++++
 isl_schedule_node.c         | 58 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 80 insertions(+)

diff --git a/doc/user.pod b/doc/user.pod
index b0ec4264..6630a339 100644
--- a/doc/user.pod
+++ b/doc/user.pod
@@ -7420,6 +7420,24 @@ of the given node should be visited.  In particular, if the callback
 returns a positive value, then the children are visited, but if
 the callback returns zero, then the children are not visited.
 
+The following function allows for a depth-first post-order
+traversal of the descendants of a specific node (including the node
+itself), where the user callback is allowed to modify the
+visited node.
+
+	#include <isl/schedule_node.h>
+	__isl_give isl_schedule_node *
+	isl_schedule_node_map_descendant(
+		__isl_take isl_schedule_node *node,
+		__isl_give isl_schedule_node *(*fn)(
+			__isl_take isl_schedule_node *node,
+			void *user), void *user);
+
+The traversal continues from the node returned by the callback function.
+It is the responsibility of the user to ensure that this does not
+lead to an infinite loop.  It is safest to always return a pointer
+to the same position (same ancestors and child positions) as the input node.
+
 Several node types have their own functions for querying
 (and in some cases setting) some node type specific properties.
 
diff --git a/include/isl/schedule_node.h b/include/isl/schedule_node.h
index 7d50bdc6..531314f0 100644
--- a/include/isl/schedule_node.h
+++ b/include/isl/schedule_node.h
@@ -28,6 +28,10 @@ __isl_give isl_schedule *isl_schedule_node_get_schedule(
 
 int isl_schedule_node_foreach_descendant(__isl_keep isl_schedule_node *node,
 	int (*fn)(__isl_keep isl_schedule_node *node, void *user), void *user);
+__isl_give isl_schedule_node *isl_schedule_node_map_descendant(
+	__isl_take isl_schedule_node *node,
+	__isl_give isl_schedule_node *(*fn)(__isl_take isl_schedule_node *node,
+		void *user), void *user);
 
 int isl_schedule_node_get_tree_depth(__isl_keep isl_schedule_node *node);
 int isl_schedule_node_has_parent(__isl_keep isl_schedule_node *node);
diff --git a/isl_schedule_node.c b/isl_schedule_node.c
index ef0162cc..253221b2 100644
--- a/isl_schedule_node.c
+++ b/isl_schedule_node.c
@@ -915,6 +915,64 @@ int isl_schedule_node_foreach_descendant(__isl_keep isl_schedule_node *node,
 	return node ? 0 : -1;
 }
 
+/* Internal data structure for isl_schedule_node_map_descendant.
+ *
+ * "fn" is the user-specified callback function.
+ * "user" is the user-specified argument for the callback.
+ */
+struct isl_schedule_node_postorder_data {
+	__isl_give isl_schedule_node *(*fn)(__isl_take isl_schedule_node *node,
+		void *user);
+	void *user;
+};
+
+/* Callback for "traverse" to enter a node and to move
+ * to the deepest initial subtree that should be traversed
+ * for use in a postorder visit.
+ *
+ * Since we are performing a postorder visit, we only need
+ * to move to the deepest initial leaf here.
+ */
+static __isl_give isl_schedule_node *postorder_enter(
+	__isl_take isl_schedule_node *node, void *user)
+{
+	while (node && isl_schedule_node_has_children(node))
+		node = isl_schedule_node_first_child(node);
+
+	return node;
+}
+
+/* Callback for "traverse" to leave a node
+ * for use in a postorder visit.
+ *
+ * Since we are performing a postorder visit, we need
+ * to call the user callback here.
+ */
+static __isl_give isl_schedule_node *postorder_leave(
+	__isl_take isl_schedule_node *node, void *user)
+{
+	struct isl_schedule_node_postorder_data *data = user;
+
+	return data->fn(node, data->user);
+}
+
+/* Traverse the descendants of "node" (including the node itself)
+ * in depth first postorder, allowing the user to modify the visited node.
+ * The traversal continues from the node returned by the callback function.
+ * It is the responsibility of the user to ensure that this does not
+ * lead to an infinite loop.  It is safest to always return a pointer
+ * to the same position (same ancestors and child positions) as the input node.
+ */
+__isl_give isl_schedule_node *isl_schedule_node_map_descendant(
+	__isl_take isl_schedule_node *node,
+	__isl_give isl_schedule_node *(*fn)(__isl_take isl_schedule_node *node,
+		void *user), void *user)
+{
+	struct isl_schedule_node_postorder_data data = { fn, user };
+
+	return traverse(node, &postorder_enter, &postorder_leave, &data);
+}
+
 /* Return the number of members in the given band node.
  */
 unsigned isl_schedule_node_band_n_member(__isl_keep isl_schedule_node *node)
-- 
2.11.4.GIT