expression copying utility

2 files changed, 149 insertions, 7 deletions

diff --git a/src/ast_utils.h b/src/ast_utils.h
index 3d4d57b77..c2a3db435 100644
--- a/src/ast_utils.h
+++ b/src/ast_utils.h
@@ -20,6 +20,7 @@
 #include "support/hash.h"
 #include "wasm.h"
 #include "wasm-traversal.h"
+#include "wasm-builder.h"
 
 namespace wasm {
 
@@ -210,6 +211,117 @@ struct ExpressionManipulator {
     new (output) OutputType(allocator);
     return output;
   }
+
+  template<typename T>
+  static Expression* flexibleCopy(Expression* original, Module& wasm, T& custom) {
+    struct Copier : public Visitor<Copier, Expression*> {
+      Module& wasm;
+      T& custom;
+
+      Builder builder;
+
+      Copier(Module& wasm, T& custom) : wasm(wasm), custom(custom), builder(wasm) {}
+
+      Expression* copy(Expression* curr) {
+        if (!curr) return nullptr;
+        auto* ret = custom.copy(curr);
+        if (ret) return ret;
+        return Visitor<Copier, Expression*>::visit(curr);
+      }
+
+      Expression* visitBlock(Block *curr) {
+        auto* ret = builder.makeBlock();
+        for (Index i = 0; i < curr->list.size(); i++) {
+          ret->list.push_back(copy(curr->list[i]));
+        }
+        ret->name = curr->name;
+        ret->finalize(curr->type);
+        return ret;
+      }
+      Expression* visitIf(If *curr) {
+        return builder.makeIf(copy(curr->condition), copy(curr->ifTrue), copy(curr->ifFalse));
+      }
+      Expression* visitLoop(Loop *curr) {
+        return builder.makeLoop(curr->out, curr->in, copy(curr->body));
+      }
+      Expression* visitBreak(Break *curr) {
+        return builder.makeBreak(curr->name, copy(curr->value), copy(curr->condition));
+      }
+      Expression* visitSwitch(Switch *curr) {
+        return builder.makeSwitch(curr->targets, curr->default_, copy(curr->condition), copy(curr->value));
+      }
+      Expression* visitCall(Call *curr) {
+        auto* ret = builder.makeCall(curr->target, {}, curr->type);
+        for (Index i = 0; i < curr->operands.size(); i++) {
+          ret->operands.push_back(copy(curr->operands[i]));
+        }
+        return ret;
+      }
+      Expression* visitCallImport(CallImport *curr) {
+        auto* ret = builder.makeCallImport(curr->target, {});
+        for (Index i = 0; i < curr->operands.size(); i++) {
+          ret->operands.push_back(copy(curr->operands[i]));
+        }
+        return ret;
+      }
+      Expression* visitCallIndirect(CallIndirect *curr) {
+        auto* ret = builder.makeCallIndirect(curr->fullType, curr->target, {}, curr->type);
+        for (Index i = 0; i < curr->operands.size(); i++) {
+          ret->operands.push_back(copy(curr->operands[i]));
+        }
+        return ret;
+      }
+      Expression* visitGetLocal(GetLocal *curr) {
+        return builder.makeGetLocal(curr->index, curr->type);
+      }
+      Expression* visitSetLocal(SetLocal *curr) {
+        return builder.makeSetLocal(curr->index, copy(curr->value));
+      }
+      Expression* visitLoad(Load *curr) {
+        return builder.makeLoad(curr->bytes, curr->signed_, curr->offset, curr->align, copy(curr->ptr), curr->type);
+      }
+      Expression* visitStore(Store *curr) {
+        return builder.makeStore(curr->bytes, curr->offset, curr->align, copy(curr->ptr), copy(curr->value));
+      }
+      Expression* visitConst(Const *curr) {
+        return builder.makeConst(curr->value);
+      }
+      Expression* visitUnary(Unary *curr) {
+        return builder.makeUnary(curr->op, copy(curr->value));
+      }
+      Expression* visitBinary(Binary *curr) {
+        return builder.makeBinary(curr->op, copy(curr->left), copy(curr->right));
+      }
+      Expression* visitSelect(Select *curr) {
+        return builder.makeSelect(copy(curr->condition), copy(curr->ifTrue), copy(curr->ifFalse));
+      }
+      Expression* visitReturn(Return *curr) {
+        return builder.makeReturn(copy(curr->value));
+      }
+      Expression* visitHost(Host *curr) {
+        assert(curr->operands.size() == 0);
+        return builder.makeHost(curr->op, curr->nameOperand, {});
+      }
+      Expression* visitNop(Nop *curr) {
+        return builder.makeNop();
+      }
+      Expression* visitUnreachable(Unreachable *curr) {
+        return builder.makeUnreachable();
+      }
+    };
+
+    Copier copier(wasm, custom);
+    return copier.copy(original);
+  }
+
+  static Expression* copy(Expression* original, Module& wasm) {
+    struct Copier {
+      Expression* copy(Expression* curr) {
+        return nullptr;
+      }
+    } copier;
+    return flexibleCopy(original, wasm, copier);
+  }
 };
 
 struct ExpressionAnalyzer {
diff --git a/src/wasm-builder.h b/src/wasm-builder.h
index 2eb051251..e58d66795 100644
--- a/src/wasm-builder.h
+++ b/src/wasm-builder.h
@@ -93,10 +93,20 @@ public:
     ret->finalize();
     return ret;
   }
-  // Switch
-  // CallBase
-  // Call
-  // Also do a version which takes a sig?
+  template<typename T>
+  Switch* makeSwitch(T& list, Name default_, Expression* condition, Expression* value = nullptr) {
+    auto* ret = wasm.allocator.alloc<Switch>();
+    ret->targets.set(list);
+    ret->default_ = default_; ret->value = value; ret->condition = condition;
+    return ret;
+  }
+  Call* makeCall(Name target, const std::vector<Expression*>& args, WasmType type) {
+    auto* call = wasm.allocator.alloc<Call>();
+    call->type = type; // not all functions may exist yet, so type must be provided
+    call->target = target;
+    call->operands.set(args);
+    return call;
+  }
   CallImport* makeCallImport(Name target, const std::vector<Expression*>& args) {
     auto* call = wasm.allocator.alloc<CallImport>();
     call->type = wasm.getImport(target)->type->result;
@@ -112,6 +122,14 @@ public:
     call->operands.set(args);
     return call;
   }
+  CallIndirect* makeCallIndirect(Name fullType, Expression* target, const std::vector<Expression*>& args, WasmType type) {
+    auto* call = wasm.allocator.alloc<CallIndirect>();
+    call->fullType = fullType;
+    call->type = type;
+    call->target = target;
+    call->operands.set(args);
+    return call;
+  }
   // FunctionType
   GetLocal* makeGetLocal(Index index, WasmType type) {
     auto* ret = wasm.allocator.alloc<GetLocal>();
@@ -126,7 +144,12 @@ public:
     ret->type = value->type;
     return ret;
   }
-  // Load
+  Load* makeLoad(unsigned bytes, bool signed_, uint32_t offset, unsigned align, Expression *ptr, WasmType type) {
+    auto* ret = wasm.allocator.alloc<Load>();
+    ret->bytes = bytes; ret->signed_ = signed_; ret->offset = offset; ret->align = align; ret->ptr = ptr;
+    ret->type = type;
+    return ret;
+  }
   Store* makeStore(unsigned bytes, uint32_t offset, unsigned align, Expression *ptr, Expression *value) {
     auto* ret = wasm.allocator.alloc<Store>();
     ret->bytes = bytes; ret->offset = offset; ret->align = align; ret->ptr = ptr; ret->value = value;
@@ -152,7 +175,12 @@ public:
     ret->finalize();
     return ret;
   }
-  // Select
+  Select* makeSelect(Expression* condition, Expression *ifTrue, Expression *ifFalse) {
+    auto* ret = wasm.allocator.alloc<Select>();
+    ret->condition = condition; ret->ifTrue = ifTrue; ret->ifFalse = ifFalse;
+    ret->finalize();
+    return ret;
+  }
   Return* makeReturn(Expression *value) {
     auto* ret = wasm.allocator.alloc<Return>();
     ret->value = value;
@@ -165,7 +193,9 @@ public:
     ret->operands.set(operands);
     return ret;
   }
-  // Unreachable
+  Unreachable* makeUnreachable() {
+    return wasm.allocator.alloc<Unreachable>();
+  }
 
   // Additional utility functions for building on top of nodes