From f813582196caeca02b56cdd3fe6e62db4f300429 Mon Sep 17 00:00:00 2001
From: Sven Verdoolaege <skimo@kotnet.org>
Date: Thu, 1 Aug 2013 10:47:03 +0200
Subject: [PATCH] add isl_schedule_map_schedule_node

Signed-off-by: Sven Verdoolaege <skimo@kotnet.org>
---
 doc/user.pod           | 12 ++++++++++--
 include/isl/schedule.h |  4 ++++
 isl_schedule.c         | 24 ++++++++++++++++++++++++
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/doc/user.pod b/doc/user.pod
index 6630a339..4bf502a4 100644
--- a/doc/user.pod
+++ b/doc/user.pod
@@ -7420,11 +7420,19 @@ 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
+The following functions allows for a depth-first post-order
+traversal of the nodes in a schedule tree or
+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.h>
+	__isl_give isl_schedule *isl_schedule_map_schedule_node(
+		__isl_take isl_schedule *schedule,
+		__isl_give isl_schedule_node *(*fn)(
+			__isl_take isl_schedule_node *node,
+			void *user), void *user);
+
 	#include <isl/schedule_node.h>
 	__isl_give isl_schedule_node *
 	isl_schedule_node_map_descendant(
diff --git a/include/isl/schedule.h b/include/isl/schedule.h
index 3f07c220..f29cfb49 100644
--- a/include/isl/schedule.h
+++ b/include/isl/schedule.h
@@ -90,6 +90,10 @@ __isl_give isl_union_set *isl_schedule_get_domain(
 
 int isl_schedule_foreach_schedule_node(__isl_keep isl_schedule *sched,
 	int (*fn)(__isl_keep isl_schedule_node *node, void *user), void *user);
+__isl_give isl_schedule *isl_schedule_map_schedule_node(
+	__isl_take isl_schedule *schedule,
+	__isl_give isl_schedule_node *(*fn)(
+		__isl_take isl_schedule_node *node, void *user), void *user);
 
 __isl_give isl_band_list *isl_schedule_get_band_forest(
 	__isl_keep isl_schedule *schedule);
diff --git a/isl_schedule.c b/isl_schedule.c
index 2abf101d..85a8e7c4 100644
--- a/isl_schedule.c
+++ b/isl_schedule.c
@@ -337,6 +337,30 @@ int isl_schedule_foreach_schedule_node(__isl_keep isl_schedule *sched,
 	return r;
 }
 
+/* Traverse the node of "sched" 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 *isl_schedule_map_schedule_node(
+	__isl_take isl_schedule *schedule,
+	__isl_give isl_schedule_node *(*fn)(
+		__isl_take isl_schedule_node *node, void *user), void *user)
+{
+	isl_schedule_node *node;
+
+	node = isl_schedule_get_root(schedule);
+	isl_schedule_free(schedule);
+
+	node = isl_schedule_node_map_descendant(node, fn, user);
+	schedule = isl_schedule_node_get_schedule(node);
+	isl_schedule_node_free(node);
+
+	return schedule;
+}
+
 /* Return an isl_union_map representation of the schedule.
  * If we still have access to the schedule tree, then we return
  * an isl_union_map corresponding to the subtree schedule of the child
-- 
2.11.4.GIT