1 files changed, 20 insertions, 139 deletions

diff --git a/src/ir/iteration.h b/src/ir/iteration.h
index 6cf149894..c6e97d392 100644
--- a/src/ir/iteration.h
+++ b/src/ir/iteration.h
@@ -18,6 +18,7 @@
 #define wasm_ir_iteration_h
 
 #include "wasm.h"
+#include "wasm-traversal.h"
 
 namespace wasm {
 
@@ -56,147 +57,27 @@ class ChildIterator {
 public:
   std::vector<Expression*> children;
 
-  ChildIterator(Expression* expr) {
-    switch (expr->_id) {
-      case Expression::Id::BlockId: {
-        auto& list = expr->cast<Block>()->list;
-        for (auto* child : list) {
-          children.push_back(child);
-        }
-        break;
-      }
-      case Expression::Id::IfId: {
-        auto* iff = expr->cast<If>();
-        children.push_back(iff->condition);
-        children.push_back(iff->ifTrue);
-        if (iff->ifFalse) children.push_back(iff->ifFalse);
-        break;
-      }
-      case Expression::Id::LoopId: {
-        children.push_back(expr->cast<Loop>()->body);
-        break;
-      }
-      case Expression::Id::BreakId: {
-        auto* br = expr->cast<Break>();
-        if (br->value) children.push_back(br->value);
-        if (br->condition) children.push_back(br->condition);
-        break;
-      }
-      case Expression::Id::SwitchId: {
-        auto* br = expr->cast<Switch>();
-        if (br->value) children.push_back(br->value);
-        children.push_back(br->condition);
-        break;
-      }
-      case Expression::Id::CallId: {
-        auto& operands = expr->cast<Call>()->operands;
-        for (auto* child : operands) {
-          children.push_back(child);
-        }
-        break;
-      }
-      case Expression::Id::CallImportId: {
-        auto& operands = expr->cast<CallImport>()->operands;
-        for (auto* child : operands) {
-          children.push_back(child);
-        }
-        break;
-      }
-      case Expression::Id::CallIndirectId: {
-        auto* call = expr->cast<CallIndirect>();
-        auto& operands = call->operands;
-        for (auto* child : operands) {
-          children.push_back(child);
-        }
-        children.push_back(call->target);
-        break;
-      }
-      case Expression::Id::SetLocalId: {
-        children.push_back(expr->cast<SetLocal>()->value);
-        break;
-      }
-      case Expression::Id::SetGlobalId: {
-        children.push_back(expr->cast<SetGlobal>()->value);
-        break;
-      }
-      case Expression::Id::LoadId: {
-        children.push_back(expr->cast<Load>()->ptr);
-        break;
-      }
-      case Expression::Id::StoreId: {
-        auto* store = expr->cast<Store>();
-        children.push_back(store->ptr);
-        children.push_back(store->value);
-        break;
-      }
-      case Expression::Id::UnaryId: {
-        children.push_back(expr->cast<Unary>()->value);
-        break;
-      }
-      case Expression::Id::BinaryId: {
-        auto* binary = expr->cast<Binary>();
-        children.push_back(binary->left);
-        children.push_back(binary->right);
-        break;
-      }
-      case Expression::Id::SelectId: {
-        auto* select = expr->cast<Select>();
-        children.push_back(select->ifTrue);
-        children.push_back(select->ifFalse);
-        children.push_back(select->condition);
-        break;
-      }
-      case Expression::Id::DropId: {
-        children.push_back(expr->cast<Drop>()->value);
-        break;
-      }
-      case Expression::Id::ReturnId: {
-        auto* ret = expr->dynCast<Return>();
-        if (ret->value) children.push_back(ret->value);
-        break;
-      }
-      case Expression::Id::HostId: {
-        auto& operands = expr->cast<Host>()->operands;
-        for (auto* child : operands) {
-          children.push_back(child);
+  ChildIterator(Expression* parent) {
+    struct Traverser : public PostWalker<Traverser> {
+      Expression* parent;
+      std::vector<Expression*>* children;
+
+      // We need to scan subchildren exactly once - just the parent.
+      bool scanned = false;
+
+      static void scan(Traverser* self, Expression** currp) {
+        if (!self->scanned) {
+          self->scanned = true;
+          PostWalker<Traverser, UnifiedExpressionVisitor<Traverser>>::scan(self, currp);
+        } else {
+          // This is one of the children. Do not scan further, just note it.
+          self->children->push_back(*currp);
         }
-        break;
-      }
-      case Expression::Id::AtomicRMWId: {
-        auto* atomic = expr->cast<AtomicRMW>();
-        children.push_back(atomic->ptr);
-        children.push_back(atomic->value);
-        break;
-      }
-      case Expression::Id::AtomicCmpxchgId: {
-        auto* atomic = expr->cast<AtomicCmpxchg>();
-        children.push_back(atomic->ptr);
-        children.push_back(atomic->expected);
-        children.push_back(atomic->replacement);
-        break;
-      }
-      case Expression::Id::AtomicWaitId: {
-        auto* atomic = expr->cast<AtomicWait>();
-        children.push_back(atomic->ptr);
-        children.push_back(atomic->expected);
-        children.push_back(atomic->timeout);
-        break;
       }
-      case Expression::Id::AtomicWakeId: {
-        auto* atomic = expr->cast<AtomicWake>();
-        children.push_back(atomic->ptr);
-        children.push_back(atomic->wakeCount);
-        break;
-      }
-      case Expression::Id::GetLocalId:
-      case Expression::Id::GetGlobalId:
-      case Expression::Id::ConstId:
-      case Expression::Id::NopId:
-      case Expression::Id::UnreachableId: {
-        break; // no children
-      }
-      default: WASM_UNREACHABLE();
-    }
+    } traverser;
+    traverser.parent = parent;
+    traverser.children = &children;
+    traverser.walk(parent);
   }
 
   Iterator begin() const {