1 files changed, 93 insertions, 12 deletions

diff --git a/src/optimize.c b/src/optimize.c
index c52966a8..9a93e3b8 100644
--- a/src/optimize.c
+++ b/src/optimize.c
@@ -101,6 +101,15 @@ static bool __stmt_type_eq(const struct stmt *stmt_a, const struct stmt *stmt_b)
 	case STMT_NOTRACK:
 		break;
 	case STMT_VERDICT:
+		expr_a = stmt_a->expr;
+		expr_b = stmt_b->expr;
+
+		if (expr_a->etype != expr_b->etype)
+			return false;
+
+		if (expr_a->etype == EXPR_MAP &&
+		    expr_b->etype == EXPR_MAP)
+			return __expr_cmp(expr_a->map, expr_b->map);
 		break;
 	case STMT_LIMIT:
 		if (stmt_a->limit.rate != stmt_b->limit.rate ||
@@ -139,14 +148,8 @@ static bool __stmt_type_eq(const struct stmt *stmt_a, const struct stmt *stmt_b)
 	return true;
 }
 
-static bool stmt_verdict_eq(const struct stmt *stmt_a, const struct stmt *stmt_b)
+static bool expr_verdict_eq(const struct expr *expr_a, const struct expr *expr_b)
 {
-	struct expr *expr_a, *expr_b;
-
-	assert (stmt_a->ops->type == STMT_VERDICT);
-
-	expr_a = stmt_a->expr;
-	expr_b = stmt_b->expr;
 	if (expr_a->verdict != expr_b->verdict)
 		return false;
 	if (expr_a->chain && expr_b->chain) {
@@ -162,6 +165,25 @@ static bool stmt_verdict_eq(const struct stmt *stmt_a, const struct stmt *stmt_b
 	return true;
 }
 
+static bool stmt_verdict_eq(const struct stmt *stmt_a, const struct stmt *stmt_b)
+{
+	struct expr *expr_a, *expr_b;
+
+	assert (stmt_a->ops->type == STMT_VERDICT);
+
+	expr_a = stmt_a->expr;
+	expr_b = stmt_b->expr;
+	if (expr_a->etype == EXPR_VERDICT &&
+	    expr_b->etype == EXPR_VERDICT)
+		return expr_verdict_eq(expr_a, expr_b);
+
+	if (expr_a->etype == EXPR_MAP &&
+	    expr_b->etype == EXPR_MAP)
+		return __expr_cmp(expr_a->map, expr_b->map);
+
+	return false;
+}
+
 static bool stmt_type_eq(const struct stmt *stmt_a, const struct stmt *stmt_b)
 {
 	if (!stmt_a && !stmt_b)
@@ -194,6 +216,10 @@ static int rule_collect_stmts(struct optimize_ctx *ctx, struct rule *rule)
 		if (stmt_type_find(ctx, stmt))
 			continue;
 
+		if (stmt->ops->type == STMT_VERDICT &&
+		    stmt->expr->etype == EXPR_MAP)
+			continue;
+
 		/* No refcounter available in statement objects, clone it to
 		 * to store in the array of selectors.
 		 */
@@ -273,16 +299,15 @@ struct merge {
 	uint32_t	num_stmts;
 };
 
-static void merge_stmts(const struct optimize_ctx *ctx,
-			uint32_t from, uint32_t to, const struct merge *merge)
+static void merge_expr_stmts(const struct optimize_ctx *ctx,
+			     uint32_t from, uint32_t to,
+			     const struct merge *merge,
+			     struct stmt *stmt_a)
 {
-	struct stmt *stmt_a = ctx->stmt_matrix[from][merge->stmt[0]];
 	struct expr *expr_a, *expr_b, *set, *elem;
 	struct stmt *stmt_b;
 	uint32_t i;
 
-	assert (stmt_a->ops->type == STMT_EXPRESSION);
-
 	set = set_expr_alloc(&internal_location, NULL);
 	set->set_flags |= NFT_SET_ANONYMOUS;
 
@@ -301,6 +326,62 @@ static void merge_stmts(const struct optimize_ctx *ctx,
 	stmt_a->expr->right = set;
 }
 
+static void merge_vmap(const struct optimize_ctx *ctx,
+		       struct stmt *stmt_a, const struct stmt *stmt_b)
+{
+	struct expr *mappings, *mapping, *expr;
+
+	mappings = stmt_b->expr->mappings;
+	list_for_each_entry(expr, &mappings->expressions, list) {
+		mapping = expr_clone(expr);
+		compound_expr_add(stmt_a->expr->mappings, mapping);
+	}
+}
+
+static void merge_verdict_stmts(const struct optimize_ctx *ctx,
+				uint32_t from, uint32_t to,
+				const struct merge *merge,
+				struct stmt *stmt_a)
+{
+	struct stmt *stmt_b;
+	uint32_t i;
+
+	for (i = from + 1; i <= to; i++) {
+		stmt_b = ctx->stmt_matrix[i][merge->stmt[0]];
+		switch (stmt_b->ops->type) {
+		case STMT_VERDICT:
+			switch (stmt_b->expr->etype) {
+			case EXPR_MAP:
+				merge_vmap(ctx, stmt_a, stmt_b);
+				break;
+			default:
+				assert(1);
+			}
+			break;
+		default:
+			assert(1);
+			break;
+		}
+	}
+}
+
+static void merge_stmts(const struct optimize_ctx *ctx,
+			uint32_t from, uint32_t to, const struct merge *merge)
+{
+	struct stmt *stmt_a = ctx->stmt_matrix[from][merge->stmt[0]];
+
+	switch (stmt_a->ops->type) {
+	case STMT_EXPRESSION:
+		merge_expr_stmts(ctx, from, to, merge, stmt_a);
+		break;
+	case STMT_VERDICT:
+		merge_verdict_stmts(ctx, from, to, merge, stmt_a);
+		break;
+	default:
+		assert(1);
+	}
+}
+
 static void merge_concat_stmts(const struct optimize_ctx *ctx,
 			       uint32_t from, uint32_t to,
 			       const struct merge *merge)