Implement operation cloning and replacement methods for OptBuilder

compiler/include/compiler/backend/optree/optimizer/opt_builder.hpp

@@ -16,17 +16,28 @@ class OptBuilder : public Builder {
         Callback onInsert;
         Callback onUpdate;
         Callback onErase;
+
+        Notifier() = default;
+        Notifier(const Notifier &) = default;
+        Notifier(Notifier &&) = default;
+        ~Notifier() = default;
+
+        Notifier(const Callback &onInsert, const Callback &onUpdate, const Callback &onErase)
+            : onInsert(onInsert), onUpdate(onUpdate), onErase(onErase){};
     };
 
     OptBuilder(const Notifier &notifier) : Builder(), notifier(notifier){};
     OptBuilder(const OptBuilder &) = delete;
     OptBuilder(OptBuilder &&) = default;
     ~OptBuilder() override = default;
 
+    using Builder::insert;
+
     void insert(const Operation::Ptr &op) override;
     Operation::Ptr clone(const Operation::Ptr &op);
     void erase(const Operation::Ptr &op);
     void update(const Operation::Ptr &op, const std::function<void()> &actor);
+    void replace(const Operation::Ptr &op, const Operation::Ptr &newOp);
 
   private:
     const Notifier &notifier;

compiler/include/compiler/optree/adaptors.hpp

@@ -2,6 +2,7 @@
 
 #include <cstdint>
 #include <string>
+#include <type_traits>
 #include <vector>
 
 #include "compiler/optree/base_adaptor.hpp"
@@ -76,6 +77,18 @@ struct ConstantOp : Adaptor {
     void init(const Type::Ptr &type, double value);
     void init(const Type::Ptr &type, const std::string &value);
 
+    template <typename T>
+    void init(const Type::Ptr &type, const T &value) {
+        if constexpr (std::is_same_v<std::remove_cvref_t<T>, bool>)
+            init(type, value);
+        else if constexpr (std::is_integral_v<T>)
+            init(type, static_cast<int64_t>(value));
+        else if constexpr (std::is_floating_point_v<T>)
+            init(type, static_cast<double>(value));
+        else
+            init(type, value);
+    }
+
     OPTREE_ADAPTOR_ATTRIBUTE_OPAQUE(value, 0)
     OPTREE_ADAPTOR_RESULT(result, 0)
 };

compiler/include/compiler/utils/helpers.hpp

@@ -49,6 +49,84 @@
     }
 };
 
+template <typename Range>
+class ReversedRange {
+    Range &&range;
+
+  public:
+    ReversedRange(Range &&range) : range(range){};
+    ~ReversedRange() = default;
+
+    auto begin() const {
+        return std::rbegin(range);
+    }
+
+    auto end() const {
+        return std::rend(range);
+    }
+};
+
+template <typename... Ranges>
+class ZippedRanges {
+    template <typename Range>
+    using NativeIterator = decltype(std::begin(std::declval<Range>()));
+    template <typename Range>
+    using NativeValueType = decltype(*std::declval<NativeIterator<Range>>());
+
+    std::tuple<Ranges &&...> ranges;
+
+    class Iterator {
+        std::tuple<NativeIterator<Ranges>...> iterators;
+
+      public:
+        Iterator() = delete;
+        Iterator(const Iterator &) = default;
+        Iterator(Iterator &&) = default;
+        ~Iterator() = default;
+
+        Iterator(NativeIterator<Ranges> &&...iterators) : iterators(iterators...){};
+
+        Iterator &operator++() {
+            std::apply([](auto &&...args) { ((++args), ...); }, iterators);
+            return *this;
+        }
+
+        Iterator operator++(int) {
+            auto temp = *this;
+            ++*this;
+            return temp;
+        }
+
+        bool operator==(const Iterator &other) const {
+            return iterators == other.iterators;
+        }
+
+        bool operator!=(const Iterator &other) const {
+            return !(*this == other);
+        }
+
+        auto operator*() {
+            return std::apply([](auto &&...args) { return std::make_tuple((*args)...); }, iterators);
+        }
+    };
+
+  public:
+    ZippedRanges() = delete;
+    ZippedRanges(const ZippedRanges &) = delete;
+    ZippedRanges(ZippedRanges &&) = default;
+    ~ZippedRanges() = default;
+
+    ZippedRanges(Ranges &&...ranges) : ranges(ranges...){};
+
+    Iterator begin() const {
+        return std::apply([](auto &&...args) { return Iterator(std::begin(args)...); }, ranges);
+    }
+
+    Iterator end() const {
+        return std::apply([](auto &&...args) { return Iterator(std::end(args)...); }, ranges);
+    }
+};
+
 } // namespace detail
 
 template <typename Range, typename UnaryPred, typename NullaryPred>
@@ -73,6 +151,16 @@
     return detail::AdvanceEarlyRange<Iterator>(begin, end);
 }
 
+template <typename Range>
+auto reversed(Range &&range) {
+    return detail::ReversedRange<Range>(range);
+}
+
+template <typename... Ranges>
+auto zip(Ranges &&...ranges) {
+    return detail::ZippedRanges<Ranges...>(ranges...);
+}
+
 template <typename... Types, typename Type>
 inline bool isAny(const std::shared_ptr<Type> &object) {
     return (object->template is<Types>() || ...);

compiler/lib/backend/optree/optimizer/opt_builder.cpp

@@ -1,29 +1,43 @@
 #include "optimizer/opt_builder.hpp"
 
+#include <algorithm>
 #include <functional>
 
 #include "compiler/optree/builder.hpp"
 #include "compiler/optree/operation.hpp"
+#include "compiler/utils/helpers.hpp"
 
 using namespace optree;
 using namespace optree::optimizer;
 
+namespace {
+
+void notifyInsertRecursively(const Operation::Ptr &op, const OptBuilder::Notifier &notifier) {
+    for (const auto &nestedOp : op->body) {
+        notifyInsertRecursively(nestedOp, notifier);
+        notifier.onInsert(nestedOp);
+    }
+}
+
+} // namespace
+
 void OptBuilder::insert(const Operation::Ptr &op) {
     Builder::insert(op);
     notifier.onInsert(op);
 }
 
 Operation::Ptr OptBuilder::clone(const Operation::Ptr &op) {
-    // TODO
-    notifier.onInsert(op);
+    auto newOp = op->clone();
+    notifyInsertRecursively(op, notifier);
+    insert(op);
     return op;
 }
 
 void OptBuilder::erase(const Operation::Ptr &op) {
     if (op->parent)
         setInsertPointAfter(op);
-    for (const auto &nestedOp : op->body)
-        erase(nestedOp);
+    for (auto &op : utils::reversed(op->body))
+        erase(op);
     op->erase();
     notifier.onErase(op);
 }
@@ -32,3 +46,14 @@
     actor();
     notifier.onUpdate(op);
 }
+
+void OptBuilder::replace(const Operation::Ptr &op, const Operation::Ptr &newOp) {
+    for (const auto &[oldResult, newResult] : utils::zip(op->results, newOp->results)) {
+        for (const auto &use : oldResult->uses) {
+            auto user = use.lock();
+            update(user, [&] { user->operand(use.operandNumber) = newResult; });
+        }
+        newResult->uses.splice_after(newResult->uses.before_begin(), oldResult->uses);
+    }
+    erase(op);
+}

compiler/lib/backend/optree/optimizer/optimizer.cpp

@@ -13,6 +13,8 @@
 using namespace optree;
 using namespace optree::optimizer;
 
+namespace {
+
 class OperationSet {
     std::vector<Operation::Ptr> data;
     std::unordered_map<const Operation *, size_t> positions;
@@ -63,44 +65,82 @@ class OperationSet {
     }
 };
 
-namespace {
+class MutationTracker {
+    Operation *const trackedOp;
+    bool updatedTag;
+    bool erasedTag;
+
+  public:
+    MutationTracker(const MutationTracker &) = delete;
+    MutationTracker(MutationTracker &&) = delete;
+    ~MutationTracker() = default;
+
+    explicit MutationTracker(const Operation::Ptr &trackedOp)
+        : trackedOp(trackedOp.get()), updatedTag(false), erasedTag(false){};
+
+    bool updated() const {
+        return updatedTag;
+    }
+    bool erased() const {
+        return erasedTag;
+    }
+
+    void raiseUpdated(const Operation::Ptr &op) {
+        if (op.get() == trackedOp)
+            updatedTag = true;
+    }
+    void raiseErased(const Operation::Ptr &op) {
+        if (op.get() == trackedOp)
+            erasedTag = true;
+    }
+};
 
 void pushToSet(const Operation::Ptr &root, OperationSet &ops) {
     for (const auto &op : root->body)
         pushToSet(op, ops);
     ops.push(root);
 }
 
-} // namespace
-
-Optimizer::Optimizer() : iterLimit(100U) {
-    transforms.emplace_back(createEraseUnusedOps());
-}
-
-void Optimizer::process(Program &program) const {
-    OperationSet ops;
-    bool mutated = false;
+OptBuilder::Notifier makeNotifier(OperationSet &ops, bool &mutated, MutationTracker &tracker) {
     OptBuilder::Notifier notifier;
     notifier.onInsert = [&ops, &mutated](const Operation::Ptr &op) {
         ops.push(op);
         mutated = true;
     };
-    notifier.onUpdate = [&ops, &mutated](const Operation::Ptr &op) {
+    notifier.onUpdate = [&ops, &mutated, &tracker](const Operation::Ptr &op) {
         ops.push(op);
         mutated = true;
+        tracker.raiseUpdated(op);
     };
-    notifier.onErase = [&ops, &mutated](const Operation::Ptr &op) {
+    notifier.onErase = [&ops, &mutated, &tracker](const Operation::Ptr &op) {
         ops.erase(op);
         mutated = true;
+        tracker.raiseErased(op);
     };
+    return notifier;
+}
+
+} // namespace
+
+Optimizer::Optimizer() : iterLimit(100U) {
+    transforms.emplace_back(createEraseUnusedOps());
+}
+
+void Optimizer::process(Program &program) const {
+    OperationSet ops;
+    bool mutated = false;
     size_t iter = 0;
     do {
         mutated = false;
         ops.clear();
         pushToSet(program.root, ops);
         while (!ops.empty()) {
             Operation::Ptr op = ops.pop();
+            MutationTracker tracker(op);
+            auto notifier = makeNotifier(ops, mutated, tracker);
             for (const auto &transform : transforms) {
+                if (tracker.erased())
+                    break;
                 if (!transform->canRun(op))
                     continue;
                 OptBuilder builder(notifier);