From 9dcd1247583cc4f2810b96b165dd45216d4dd686 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=81=95=E5=B1=B1?= <lihaokun.lhk@antgroup.com>
Date: Thu, 9 Jan 2025 14:54:11 +0800
Subject: [PATCH] add inst combine rule for inst_aggr

---
 .../godel-frontend/src/ir/inst_combine.cpp         |  9 +++++++++
 godel-script/godel-frontend/src/ir/inst_combine.h  |  1 +
 godel-script/godel-frontend/src/ir/lir.cpp         | 14 +++++++++++++-
 godel-script/godel-frontend/src/ir/lir.h           |  4 ++++
 4 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/godel-script/godel-frontend/src/ir/inst_combine.cpp b/godel-script/godel-frontend/src/ir/inst_combine.cpp
index 2e8b47ea..eabdc67e 100644
--- a/godel-script/godel-frontend/src/ir/inst_combine.cpp
+++ b/godel-script/godel-frontend/src/ir/inst_combine.cpp
@@ -241,6 +241,15 @@ void combine_worker::visit_binary(lir::binary* node) {
     }
 }
 
+void combine_worker::visit_aggregator(lir::aggregator* node) {
+    const auto& tgt = node->get_target();
+    if (is_single_ref_ssa_temp(tgt.content)) {
+        const auto& ref = get_single_ref(tgt.content);
+        node->get_mutable_target().content = ref.first;
+        ref.second->set_flag_eliminated(true);
+    }
+}
+
 void combine_worker::mark(souffle_rule_impl* b) {
     b->get_block()->accept(this);
 }
diff --git a/godel-script/godel-frontend/src/ir/inst_combine.h b/godel-script/godel-frontend/src/ir/inst_combine.h
index e14abc7c..b6b0c253 100644
--- a/godel-script/godel-frontend/src/ir/inst_combine.h
+++ b/godel-script/godel-frontend/src/ir/inst_combine.h
@@ -53,6 +53,7 @@ class combine_worker: public lir::inst_visitor {
     void visit_record(lir::record*) override;
     void visit_unary(lir::unary*) override;
     void visit_binary(lir::binary*) override;
+    void visit_aggregator(lir::aggregator*) override;
 
 public:
     combine_worker(const inst_combine_pass::ref_graph& g): vg(g) {}
diff --git a/godel-script/godel-frontend/src/ir/lir.cpp b/godel-script/godel-frontend/src/ir/lir.cpp
index 6c76ea7f..5e43efa6 100644
--- a/godel-script/godel-frontend/src/ir/lir.cpp
+++ b/godel-script/godel-frontend/src/ir/lir.cpp
@@ -359,7 +359,19 @@ void not_operand::dump(std::ostream& os, const std::string& indent) const {
     // only one statement in the block
     if (body->get_content().size()==1) {
         os << indent << "!(";
-        body->get_content()[0]->dump(os, "");
+        switch(body->get_content()[0]->get_kind()) {
+            case inst_kind::inst_not:
+            case inst_kind::inst_and:
+            case inst_kind::inst_or:
+            case inst_kind::inst_aggr:
+                os << "\n";
+                body->get_content()[0]->dump(os, indent + "  ");
+                os << "\n" << indent;
+                break;
+            default:
+                body->get_content()[0]->dump(os, "");
+                break;
+        }
         os << ")";
         return;
     }
diff --git a/godel-script/godel-frontend/src/ir/lir.h b/godel-script/godel-frontend/src/ir/lir.h
index e69ec945..8cb2e8c5 100644
--- a/godel-script/godel-frontend/src/ir/lir.h
+++ b/godel-script/godel-frontend/src/ir/lir.h
@@ -666,6 +666,10 @@ class aggregator: public inst {
     void accept(inst_visitor* v) override {
         v->visit_aggregator(this);
     }
+
+public:
+    const auto& get_target() const { return destination; }
+    auto& get_mutable_target() { return destination; }
 };
 
 }