1 files changed, 50 insertions, 10 deletions

diff --git a/src/optimize.c b/src/optimize.c
index ea067f80..180a7d55 100644
--- a/src/optimize.c
+++ b/src/optimize.c
@@ -550,12 +550,58 @@ static void merge_stmts(const struct optimize_ctx *ctx,
 	}
 }
 
+static void __merge_concat_stmts(const struct optimize_ctx *ctx, uint32_t i,
+				 const struct merge *merge, struct expr *set)
+{
+	struct expr *concat, *next, *expr, *concat_clone, *clone, *elem;
+	LIST_HEAD(pending_list);
+	LIST_HEAD(concat_list);
+	struct stmt *stmt_a;
+	uint32_t k;
+
+	concat = concat_expr_alloc(&internal_location);
+	list_add(&concat->list, &concat_list);
+
+	for (k = 0; k < merge->num_stmts; k++) {
+		list_for_each_entry_safe(concat, next, &concat_list, list) {
+			stmt_a = ctx->stmt_matrix[i][merge->stmt[k]];
+			switch (stmt_a->expr->right->etype) {
+			case EXPR_SET:
+				list_for_each_entry(expr, &stmt_a->expr->right->expressions, list) {
+					concat_clone = expr_clone(concat);
+					clone = expr_clone(expr->key);
+					compound_expr_add(concat_clone, clone);
+					list_add_tail(&concat_clone->list, &pending_list);
+				}
+				list_del(&concat->list);
+				expr_free(concat);
+				break;
+			case EXPR_SYMBOL:
+			case EXPR_VALUE:
+				clone = expr_clone(stmt_a->expr->right);
+				compound_expr_add(concat, clone);
+				break;
+			default:
+				assert(0);
+				break;
+			}
+		}
+		list_splice_init(&pending_list, &concat_list);
+	}
+
+	list_for_each_entry_safe(concat, next, &concat_list, list) {
+		list_del(&concat->list);
+		elem = set_elem_expr_alloc(&internal_location, concat);
+		compound_expr_add(set, elem);
+	}
+}
+
 static void merge_concat_stmts(const struct optimize_ctx *ctx,
 			       uint32_t from, uint32_t to,
 			       const struct merge *merge)
 {
-	struct expr *concat, *elem, *set;
 	struct stmt *stmt, *stmt_a;
+	struct expr *concat, *set;
 	uint32_t i, k;
 
 	stmt = ctx->stmt_matrix[from][merge->stmt[0]];
@@ -573,15 +619,9 @@ static void merge_concat_stmts(const struct optimize_ctx *ctx,
 	set = set_expr_alloc(&internal_location, NULL);
 	set->set_flags |= NFT_SET_ANONYMOUS;
 
-	for (i = from; i <= to; i++) {
-		concat = concat_expr_alloc(&internal_location);
-		for (k = 0; k < merge->num_stmts; k++) {
-			stmt_a = ctx->stmt_matrix[i][merge->stmt[k]];
-			compound_expr_add(concat, expr_get(stmt_a->expr->right));
-		}
-		elem = set_elem_expr_alloc(&internal_location, concat);
-		compound_expr_add(set, elem);
-	}
+	for (i = from; i <= to; i++)
+		__merge_concat_stmts(ctx, i, merge, set);
+
 	expr_free(stmt->expr->right);
 	stmt->expr->right = set;