From 98e35c02c9b4abcf28efaca671b7fec7ad357bf4 Mon Sep 17 00:00:00 2001
From: Sven Verdoolaege <skimo@kotnet.org>
Date: Tue, 3 Dec 2013 17:06:03 +0100
Subject: [PATCH] add isl_schedule_tree_band_set_partial_schedule

This will be useful for the isl_schedule_node_group function
that will be introduced in a subsequent commit.

Signed-off-by: Sven Verdoolaege <skimo@kotnet.org>
---
 isl_schedule_band.c | 20 ++++++++++++++++++++
 isl_schedule_band.h |  3 +++
 isl_schedule_tree.c | 23 +++++++++++++++++++++++
 isl_schedule_tree.h |  3 +++
 4 files changed, 49 insertions(+)

diff --git a/isl_schedule_band.c b/isl_schedule_band.c
index 647f02fb..cd36c524 100644
--- a/isl_schedule_band.c
+++ b/isl_schedule_band.c
@@ -338,6 +338,26 @@ __isl_give isl_multi_union_pw_aff *isl_schedule_band_get_partial_schedule(
 	return band ? isl_multi_union_pw_aff_copy(band->mupa) : NULL;
 }
 
+/* Replace the schedule of "band" by "schedule".
+ */
+__isl_give isl_schedule_band *isl_schedule_band_set_partial_schedule(
+	__isl_take isl_schedule_band *band,
+	__isl_take isl_multi_union_pw_aff *schedule)
+{
+	band = isl_schedule_band_cow(band);
+	if (!band || !schedule)
+		goto error;
+
+	isl_multi_union_pw_aff_free(band->mupa);
+	band->mupa = schedule;
+
+	return band;
+error:
+	isl_schedule_band_free(band);
+	isl_multi_union_pw_aff_free(schedule);
+	return NULL;
+}
+
 /* Return the loop AST generation type for the band member of "band"
  * at position "pos".
  */
diff --git a/isl_schedule_band.h b/isl_schedule_band.h
index a2975a9d..63a624dd 100644
--- a/isl_schedule_band.h
+++ b/isl_schedule_band.h
@@ -62,6 +62,9 @@ __isl_give isl_schedule_band *isl_schedule_band_intersect_domain(
 	__isl_take isl_schedule_band *band, __isl_take isl_union_set *domain);
 __isl_give isl_multi_union_pw_aff *isl_schedule_band_get_partial_schedule(
 	__isl_keep isl_schedule_band *band);
+__isl_give isl_schedule_band *isl_schedule_band_set_partial_schedule(
+	__isl_take isl_schedule_band *band,
+	__isl_take isl_multi_union_pw_aff *schedule);
 enum isl_ast_loop_type isl_schedule_band_member_get_ast_loop_type(
 	__isl_keep isl_schedule_band *band, int pos);
 __isl_give isl_schedule_band *isl_schedule_band_member_set_ast_loop_type(
diff --git a/isl_schedule_tree.c b/isl_schedule_tree.c
index 0b6f030c..9849b61f 100644
--- a/isl_schedule_tree.c
+++ b/isl_schedule_tree.c
@@ -968,6 +968,29 @@ __isl_give isl_multi_union_pw_aff *isl_schedule_tree_band_get_partial_schedule(
 	return isl_schedule_band_get_partial_schedule(tree->band);
 }
 
+/* Replace the schedule of the band tree root by "schedule".
+ */
+__isl_give isl_schedule_tree *isl_schedule_tree_band_set_partial_schedule(
+	__isl_take isl_schedule_tree *tree,
+	__isl_take isl_multi_union_pw_aff *schedule)
+{
+	tree = isl_schedule_tree_cow(tree);
+	if (!tree || !schedule)
+		goto error;
+
+	if (tree->type != isl_schedule_node_band)
+		isl_die(isl_schedule_tree_get_ctx(tree), isl_error_invalid,
+			"not a band node", return NULL);
+	tree->band = isl_schedule_band_set_partial_schedule(tree->band,
+								schedule);
+
+	return tree;
+error:
+	isl_schedule_tree_free(tree);
+	isl_multi_union_pw_aff_free(schedule);
+	return NULL;
+}
+
 /* Return the loop AST generation type for the band member
  * of the band tree root at position "pos".
  */
diff --git a/isl_schedule_tree.h b/isl_schedule_tree.h
index 7b40be09..ac1254f4 100644
--- a/isl_schedule_tree.h
+++ b/isl_schedule_tree.h
@@ -103,6 +103,9 @@ __isl_give isl_schedule_tree *isl_schedule_tree_band_intersect_domain(
 	__isl_take isl_schedule_tree *tree, __isl_take isl_union_set *domain);
 __isl_give isl_multi_union_pw_aff *isl_schedule_tree_band_get_partial_schedule(
 	__isl_keep isl_schedule_tree *tree);
+__isl_give isl_schedule_tree *isl_schedule_tree_band_set_partial_schedule(
+	__isl_take isl_schedule_tree *tree,
+	__isl_take isl_multi_union_pw_aff *schedule);
 enum isl_ast_loop_type isl_schedule_tree_band_member_get_ast_loop_type(
 	__isl_keep isl_schedule_tree *tree, int pos);
 __isl_give isl_schedule_tree *isl_schedule_tree_band_member_set_ast_loop_type(
-- 
2.11.4.GIT