Merge pull request #1675 from Shaikh-Ubaid/wasm_assembler

ubaidsk · web-flow · commit e1c55fd698aa · 2023-04-11T23:17:35.000+05:30
WASM: Implement WASMAssembler Class
diff --git a/src/libasr/codegen/asr_to_wasm.cpp b/src/libasr/codegen/asr_to_wasm.cpp
diff --git a/src/libasr/codegen/wasm_assembler.h b/src/libasr/codegen/wasm_assembler.h
diff --git a/src/libasr/codegen/wasm_decoder.h b/src/libasr/codegen/wasm_decoder.h
@@ -4,8 +4,6 @@
 #include <fstream>
 
 #include <libasr/assert.h>
-#include <libasr/alloc.h>
-#include <libasr/containers.h>
 #include <libasr/codegen/wasm_utils.h>
 
 // #define WAT_DEBUG
@@ -48,9 +46,6 @@ class WASMDecoder {
     Struct &self() { return static_cast<Struct &>(*this); }
 
    public:
-    std::unordered_map<uint8_t, std::string> var_type_to_string;
-    std::unordered_map<uint8_t, std::string> kind_to_string;
-
     Allocator &al;
     diag::Diagnostics &diag;
     Vec<uint8_t> wasm_bytes;
@@ -67,10 +62,6 @@ class WASMDecoder {
 
     WASMDecoder(Allocator &al, diag::Diagnostics &diagonostics)
         : al(al), diag(diagonostics) {
-        var_type_to_string = {
-            {0x7F, "i32"}, {0x7E, "i64"}, {0x7D, "f32"}, {0x7C, "f64"}};
-        kind_to_string = {
-            {0x00, "func"}, {0x01, "table"}, {0x02, "memory"}, {0x03, "global"}};
 
         PREAMBLE_SIZE = 8 /* BYTES */;
         // wasm_bytes.reserve(al, 1024 * 128);
diff --git a/src/libasr/codegen/wasm_to_wat.cpp b/src/libasr/codegen/wasm_to_wat.cpp
@@ -313,13 +313,12 @@ class WATVisitor : public WASMDecoder<WATVisitor>,
             result +=
                 indent + "(type (;" + std::to_string(i) + ";) (func (param";
             for (uint32_t j = 0; j < func_types[i].param_types.size(); j++) {
-                result +=
-                    " " + var_type_to_string[func_types[i].param_types.p[j]];
+                result += " " + vt2s(func_types[i].param_types.p[j]);
             }
             result += ") (result";
             for (uint32_t j = 0; j < func_types[i].result_types.size(); j++) {
                 result +=
-                    " " + var_type_to_string[func_types[i].result_types.p[j]];
+                    " " + vt2s(func_types[i].result_types.p[j]);
             }
             result += ")))";
         }
@@ -350,8 +349,8 @@ class WATVisitor : public WASMDecoder<WATVisitor>,
                 decode_instructions();
                 global_initialization_insts = this->src;
             }
-            std::string global_type = ((globals[i].mut == 0x00) ? var_type_to_string[globals[i].type]:
-            "(mut " + var_type_to_string[globals[i].type] + ")" );
+            std::string global_type = ((globals[i].mut == 0x00) ? vt2s(globals[i].type):
+            "(mut " + vt2s(globals[i].type) + ")" );
             result += indent + "(global $" + std::to_string(i);
             result += " " + global_type;
             result += " (" + global_initialization_insts + "))";
@@ -365,20 +364,20 @@ class WATVisitor : public WASMDecoder<WATVisitor>,
                  j++) {
                 result +=
                     " " +
-                    var_type_to_string[func_types[func_index].param_types.p[j]];
+                    vt2s(func_types[func_index].param_types.p[j]);
             }
             result += ") (result";
             for (uint32_t j = 0; j < func_types[func_index].result_types.size();
                  j++) {
-                result += " " + var_type_to_string[func_types[func_index]
-                                                       .result_types.p[j]];
+                result += " " + vt2s(func_types[func_index]
+                                                       .result_types.p[j]);
             }
             result += ")";
             result += indent + "    (local";
             for (uint32_t j = 0; j < codes.p[i].locals.size(); j++) {
                 for (uint32_t k = 0; k < codes.p[i].locals.p[j].count; k++) {
                     result +=
-                        " " + var_type_to_string[codes.p[i].locals.p[j].type];
+                        " " + vt2s(codes.p[i].locals.p[j].type);
                 }
             }
             result += ")";
@@ -403,7 +402,7 @@ class WATVisitor : public WASMDecoder<WATVisitor>,
 
         for (uint32_t i = 0; i < exports.size(); i++) {
             result += indent + "(export \"" + exports.p[i].name + "\" (" +
-                      kind_to_string[exports.p[i].kind] + " " +
+                      k2s(exports.p[i].kind) + " " +
                       std::to_string(exports.p[i].index) + "))";
         }
 
diff --git a/src/libasr/codegen/wasm_to_x64.cpp b/src/libasr/codegen/wasm_to_x64.cpp
@@ -212,7 +212,7 @@ class X64Visitor : public WASMDecoder<X64Visitor>,
 
     void visit_GlobalGet(uint32_t globalidx) {
         std::string loc = "global_" + std::to_string(globalidx);
-        std::string var_type = var_type_to_string[globals[globalidx].type];
+        std::string var_type = vt2s(globals[globalidx].type);
 
         X64Reg base = X64Reg::rbx;
         m_a.asm_mov_r64_label(X64Reg::rbx, loc);
@@ -235,7 +235,7 @@ class X64Visitor : public WASMDecoder<X64Visitor>,
         }
 
         std::string loc = "global_" + std::to_string(globalidx);
-        std::string var_type = var_type_to_string[globals[globalidx].type];
+        std::string var_type = vt2s(globals[globalidx].type);
 
         X64Reg base = X64Reg::rbx;
         m_a.asm_mov_r64_label(X64Reg::rbx, loc);
@@ -257,7 +257,7 @@ class X64Visitor : public WASMDecoder<X64Visitor>,
         auto cur_func_param_type = func_types[type_indices[cur_func_idx]];
         int no_of_params = (int)cur_func_param_type.param_types.size();
         if ((int)localidx < no_of_params) {
-            std::string var_type = var_type_to_string[cur_func_param_type.param_types[localidx]];
+            std::string var_type = vt2s(cur_func_param_type.param_types[localidx]);
             if (var_type == "i32" || var_type == "i64") {
                 m_a.asm_mov_r64_m64(X64Reg::rax, &base, nullptr, 1, 8 * (2 + no_of_params - (int)localidx - 1));
                 m_a.asm_push_r64(X64Reg::rax);
@@ -271,7 +271,7 @@ class X64Visitor : public WASMDecoder<X64Visitor>,
             }
         } else {
             localidx -= no_of_params;
-            std::string var_type = var_type_to_string[codes[cur_func_idx].locals[localidx].type];
+            std::string var_type = vt2s(codes[cur_func_idx].locals[localidx].type);
             if (var_type == "i32" || var_type == "i64") {
                 m_a.asm_mov_r64_m64(X64Reg::rax, &base, nullptr, 1, -8 * (1 + (int)localidx));
                 m_a.asm_push_r64(X64Reg::rax);
@@ -291,7 +291,7 @@ class X64Visitor : public WASMDecoder<X64Visitor>,
         auto cur_func_param_type = func_types[type_indices[cur_func_idx]];
         int no_of_params = (int)cur_func_param_type.param_types.size();
         if ((int)localidx < no_of_params) {
-            std::string var_type = var_type_to_string[cur_func_param_type.param_types[localidx]];
+            std::string var_type = vt2s(cur_func_param_type.param_types[localidx]);
             if (var_type == "i32" || var_type == "i64") {
                 m_a.asm_pop_r64(X64Reg::rax);
                 m_a.asm_mov_m64_r64(&base, nullptr, 1, 8 * (2 + no_of_params - (int)localidx - 1), X64Reg::rax);
@@ -305,7 +305,7 @@ class X64Visitor : public WASMDecoder<X64Visitor>,
             }
         } else {
             localidx -= no_of_params;
-            std::string var_type = var_type_to_string[codes[cur_func_idx].locals[localidx].type];
+            std::string var_type = vt2s(codes[cur_func_idx].locals[localidx].type);
             if (var_type == "i32" || var_type == "i64") {
                 m_a.asm_pop_r64(X64Reg::rax);
                 m_a.asm_mov_m64_r64(&base, nullptr, 1, -8 * (1 + (int)localidx), X64Reg::rax);
diff --git a/src/libasr/codegen/wasm_to_x86.cpp b/src/libasr/codegen/wasm_to_x86.cpp
@@ -209,7 +209,7 @@ class X86Visitor : public WASMDecoder<X86Visitor>,
         auto cur_func_param_type = func_types[type_indices[cur_func_idx]];
         int no_of_params = (int)cur_func_param_type.param_types.size();
         if ((int)localidx < no_of_params) {
-            std::string var_type = var_type_to_string[cur_func_param_type.param_types[localidx]];
+            std::string var_type = vt2s(cur_func_param_type.param_types[localidx]);
             if (var_type == "i32") {
                 m_a.asm_mov_r32_m32(X86Reg::eax, &base, nullptr, 1, 4 * (2 + no_of_params - (int)localidx - 1));
                 m_a.asm_push_r32(X86Reg::eax);
@@ -227,7 +227,7 @@ class X86Visitor : public WASMDecoder<X86Visitor>,
 
         } else {
             localidx -= no_of_params;
-            std::string var_type = var_type_to_string[codes[cur_func_idx].locals[localidx].type];
+            std::string var_type = vt2s(codes[cur_func_idx].locals[localidx].type);
             if (var_type == "i32") {
                 m_a.asm_mov_r32_m32(X86Reg::eax, &base, nullptr, 1, -4 * (1 + localidx));
                 m_a.asm_push_r32(X86Reg::eax);
@@ -249,7 +249,7 @@ class X86Visitor : public WASMDecoder<X86Visitor>,
         auto cur_func_param_type = func_types[type_indices[cur_func_idx]];
         int no_of_params = (int)cur_func_param_type.param_types.size();
         if ((int)localidx < no_of_params) {
-            std::string var_type = var_type_to_string[cur_func_param_type.param_types[localidx]];
+            std::string var_type = vt2s(cur_func_param_type.param_types[localidx]);
             if (var_type == "i32") {
                 m_a.asm_pop_r32(X86Reg::eax);
                 m_a.asm_mov_m32_r32(&base, nullptr, 1, 4 * (2 + no_of_params - (int)localidx - 1), X86Reg::eax);
@@ -267,7 +267,7 @@ class X86Visitor : public WASMDecoder<X86Visitor>,
 
         } else {
             localidx -= no_of_params;
-            std::string var_type = var_type_to_string[codes[cur_func_idx].locals[localidx].type];
+            std::string var_type = vt2s(codes[cur_func_idx].locals[localidx].type);
             if (var_type == "i32") {
                 m_a.asm_pop_r32(X86Reg::eax);
                 m_a.asm_mov_m32_r32(&base, nullptr, 1, -4 * (1 + (int)localidx), X86Reg::eax);
diff --git a/src/libasr/codegen/wasm_utils.h b/src/libasr/codegen/wasm_utils.h
@@ -11,6 +11,56 @@ namespace LCompilers {
 
 namespace wasm {
 
+enum var_type: uint8_t {
+    i32 = 0x7F,
+    i64 = 0x7E,
+    f32 = 0x7D,
+    f64 = 0x7C
+};
+
+enum mem_align : uint8_t {
+    b8 = 0,
+    b16 = 1,
+    b32 = 2,
+    b64 = 3
+};
+
+enum wasm_kind: uint8_t {
+    func = 0x00,
+    table = 0x01,
+    memory = 0x02,
+    global = 0x03
+};
+
+template <typename T>
+std::string vt2s(T vt) {
+    switch(vt) {
+        case var_type::i32: return "i32";
+        case var_type::i64: return "i64";
+        case var_type::f32: return "f32";
+        case var_type::f64: return "f64";
+        default:
+            std::cerr << "Unsupported wasm var_type" << std::endl;
+            LCOMPILERS_ASSERT(false);
+            return "";
+
+    }
+}
+
+template <typename T>
+std::string k2s(T k) {
+    switch(k) {
+        case wasm_kind::func: return "func";
+        case wasm_kind::table: return "table";
+        case wasm_kind::memory: return "memory";
+        case wasm_kind::global: return "global";
+        default:
+            std::cerr << "Unsupported wasm kind" << std::endl;
+            LCOMPILERS_ASSERT(false);
+            return "";
+    }
+}
+
 struct FuncType {
     Vec<uint8_t> param_types;
     Vec<uint8_t> result_types;
diff --git a/src/libasr/wasm_instructions_visitor.py b/src/libasr/wasm_instructions_visitor.py
@@ -105,10 +105,12 @@ def visit_BaseWASMVisitor(self, mod, *args):
         self.emit("};\n", 0)
 
     def visit_WASMInstsAssembler(self, mod):
+        self.emit("template <class Struct>", 0)
         self.emit("class WASMInstsAssembler {", 0)
         self.emit("private:", 0)
         self.emit(    "Allocator &m_al;", 1)
         self.emit(    "Vec<uint8_t> &m_code;\n", 1)
+        self.emit(    "Struct &self() { return static_cast<Struct &>(*this); }", 1)
         self.emit("public:", 0)
         self.emit(    "WASMInstsAssembler(Allocator &al, Vec<uint8_t> &code): m_al(al), m_code(code) {}\n", 1)
 
