isl_stream_read_map: handle "implies" token

Requested-by: Vladimir Klebanov <klebanov@kit.edu>
Signed-off-by: Sven Verdoolaege <skimo@kotnet.org>

4 files changed, 59 insertions, 12 deletions

diff --git a/include/isl/stream.h b/include/isl/stream.h
index d3f65e43..9ad2374b 100644
--- a/include/isl/stream.h
+++ b/include/isl/stream.h
@@ -36,6 +36,7 @@ enum isl_token_type { ISL_TOKEN_ERROR = -1,
 			ISL_TOKEN_STRING,
 			ISL_TOKEN_MAP, ISL_TOKEN_AFF,
 			ISL_TOKEN_CEIL, ISL_TOKEN_FLOOR,
+			ISL_TOKEN_IMPLIES,
 			ISL_TOKEN_LAST };
 
 struct isl_token {
diff --git a/isl_input.c b/isl_input.c
index 480d60d7..8aed9760 100644
--- a/isl_input.c
+++ b/isl_input.c
@@ -660,7 +660,7 @@ static int is_comparator(struct isl_token *tok)
 	}
 }
 
-static struct isl_map *read_disjuncts(struct isl_stream *s,
+static __isl_give isl_map *read_formula(struct isl_stream *s,
 	struct vars *v, __isl_take isl_map *map, int rational);
 static __isl_give isl_pw_aff *accept_extended_affine(struct isl_stream *s,
 	__isl_take isl_space *dim, struct vars *v, int rational);
@@ -750,7 +750,7 @@ static __isl_give isl_pw_aff *accept_extended_affine(struct isl_stream *s,
 
 	isl_stream_push_token(s, tok);
 
-	cond = read_disjuncts(s, v, cond, rational);
+	cond = read_formula(s, v, cond, rational);
 
 	return accept_ternary(s, cond, v, rational);
 }
@@ -1317,7 +1317,7 @@ static __isl_give isl_map *read_exists(struct isl_stream *s,
 	if (isl_stream_eat(s, ':'))
 		goto error;
 
-	map = read_disjuncts(s, v, map, rational);
+	map = read_formula(s, v, map, rational);
 	map = isl_set_unwrap(isl_map_domain(map));
 
 	vars_drop(v, v->n - n);
@@ -1365,7 +1365,7 @@ static int resolve_paren_expr(struct isl_stream *s,
 	    isl_stream_next_token_is(s, ISL_TOKEN_TRUE) ||
 	    isl_stream_next_token_is(s, ISL_TOKEN_FALSE) ||
 	    isl_stream_next_token_is(s, ISL_TOKEN_MAP)) {
-		map = read_disjuncts(s, v, map, rational);
+		map = read_formula(s, v, map, rational);
 		if (isl_stream_eat(s, ')'))
 			goto error;
 		tok->type = ISL_TOKEN_MAP;
@@ -1400,7 +1400,7 @@ static int resolve_paren_expr(struct isl_stream *s,
 
 	isl_stream_push_token(s, tok2);
 
-	map = read_disjuncts(s, v, map, rational);
+	map = read_formula(s, v, map, rational);
 	if (isl_stream_eat(s, ')'))
 		goto error;
 
@@ -1502,6 +1502,46 @@ static struct isl_map *read_disjuncts(struct isl_stream *s,
 	return res;
 }
 
+/* Read a first order formula from "s", add the corresponding
+ * constraints to "map" and return the result.
+ *
+ * In particular, read a formula of the form
+ *
+ *	a
+ *
+ * or
+ *
+ *	a implies b
+ *
+ * where a and b are disjunctions.
+ *
+ * In the first case, map is replaced by
+ *
+ *	map \cap { [..] : a }
+ *
+ * In the second case, it is replaced by
+ *
+ *	(map \setminus { [..] : a}) \cup (map \cap { [..] : b })
+ */
+static __isl_give isl_map *read_formula(struct isl_stream *s,
+	struct vars *v, __isl_take isl_map *map, int rational)
+{
+	isl_map *res;
+
+	res = read_disjuncts(s, v, isl_map_copy(map), rational);
+
+	if (isl_stream_eat_if_available(s, ISL_TOKEN_IMPLIES)) {
+		isl_map *res2;
+
+		res = isl_map_subtract(isl_map_copy(map), res);
+		res2 = read_disjuncts(s, v, map, rational);
+		res = isl_map_union(res, res2);
+	} else
+		isl_map_free(map);
+
+	return res;
+}
+
 static int polylib_pos_to_isl_pos(__isl_keep isl_basic_map *bmap, int pos)
 {
 	if (pos < isl_basic_map_dim(bmap, isl_dim_out))
@@ -1926,7 +1966,7 @@ static __isl_give isl_pw_qpolynomial *read_term(struct isl_stream *s,
 	return pwqp;
 }
 
-static __isl_give isl_map *read_optional_disjuncts(struct isl_stream *s,
+static __isl_give isl_map *read_optional_formula(struct isl_stream *s,
 	__isl_take isl_map *map, struct vars *v, int rational)
 {
 	struct isl_token *tok;
@@ -1939,7 +1979,7 @@ static __isl_give isl_map *read_optional_disjuncts(struct isl_stream *s,
 	if (tok->type == ':' ||
 	    (tok->type == ISL_TOKEN_OR && !strcmp(tok->u.s, "|"))) {
 		isl_token_free(tok);
-		map = read_disjuncts(s, v, map, rational);
+		map = read_formula(s, v, map, rational);
 	} else
 		isl_stream_push_token(s, tok);
 
@@ -1957,7 +1997,7 @@ static struct isl_obj obj_read_poly(struct isl_stream *s,
 	struct isl_set *set;
 
 	pwqp = read_term(s, map, v);
-	map = read_optional_disjuncts(s, map, v, 0);
+	map = read_optional_formula(s, map, v, 0);
 	set = isl_map_range(map);
 
 	pwqp = isl_pw_qpolynomial_intersect_domain(pwqp, set);
@@ -1995,7 +2035,7 @@ static struct isl_obj obj_read_poly_or_fold(struct isl_stream *s,
 	if (isl_stream_eat(s, ')'))
 		goto error;
 
-	set = read_optional_disjuncts(s, set, v, 0);
+	set = read_optional_formula(s, set, v, 0);
 	pwf = isl_pw_qpolynomial_fold_intersect_domain(pwf, set);
 
 	vars_drop(v, v->n - n);
@@ -2041,7 +2081,7 @@ static struct isl_obj obj_read_body(struct isl_stream *s,
 
 	if (isl_stream_next_token_is(s, ':')) {
 		obj.type = isl_obj_set;
-		obj.v = read_optional_disjuncts(s, map, v, rational);
+		obj.v = read_optional_formula(s, map, v, rational);
 		return obj;
 	}
 
@@ -2069,7 +2109,7 @@ static struct isl_obj obj_read_body(struct isl_stream *s,
 		isl_stream_push_token(s, tok);
 	}
 
-	map = read_optional_disjuncts(s, map, v, rational);
+	map = read_optional_formula(s, map, v, rational);
 
 	vars_drop(v, v->n - n);
 
@@ -2746,7 +2786,7 @@ static __isl_give isl_pw_aff *read_pw_aff_with_dom(struct isl_stream *s,
 	if (isl_stream_eat(s, ']'))
 		goto error;
 
-	dom = read_optional_disjuncts(s, dom, v, 0);
+	dom = read_optional_formula(s, dom, v, 0);
 	pwaff = isl_pw_aff_intersect_domain(pwaff, dom);
 
 	return pwaff;
diff --git a/isl_stream.c b/isl_stream.c
index de69f450..c5b15115 100644
--- a/isl_stream.c
+++ b/isl_stream.c
@@ -299,6 +299,8 @@ static enum isl_token_type check_keywords(struct isl_stream *s)
 		return ISL_TOKEN_AND;
 	if (!strcasecmp(s->buffer, "or"))
 		return ISL_TOKEN_OR;
+	if (!strcasecmp(s->buffer, "implies"))
+		return ISL_TOKEN_IMPLIES;
 	if (!strcasecmp(s->buffer, "not"))
 		return ISL_TOKEN_NOT;
 	if (!strcasecmp(s->buffer, "infty"))
diff --git a/isl_test.c b/isl_test.c
index 5d3e4aaf..662f481b 100644
--- a/isl_test.c
+++ b/isl_test.c
@@ -213,6 +213,10 @@ int test_parse(struct isl_ctx *ctx)
 				      "{ [i] : exists a : i = 2 a }") < 0)
 		return -1;
 
+	if (test_parse_map_equal(ctx, "{ [a] -> [b] : a = 5 implies b = 5 }",
+				      "{ [a] -> [b] : a != 5 or b = 5 }") < 0)
+		return -1;
+
 	return 0;
 }