Fix issues of compiling control related opcodes (#1063)

core/iwasm/fast-jit/fe/jit_emit_control.c

@@ -56,7 +56,7 @@ static JitBlock *
 get_target_block(JitCompContext *cc, uint32 br_depth)
 {
     uint32 i = br_depth;
-    JitBlock *block = cc->block_stack.block_list_end;
+    JitBlock *block = jit_block_stack_top(&cc->block_stack);
 
     while (i-- > 0 && block) {
         block = block->prev;
@@ -180,7 +180,7 @@ push_jit_block_to_stack_and_pass_params(JitCompContext *cc, JitBlock *block,
            the new block's value stack */
         for (i = 0; i < block->param_count; i++) {
             jit_value = jit_value_stack_pop(
-                &cc->block_stack.block_list_end->value_stack);
+                &jit_block_stack_top(&cc->block_stack)->value_stack);
             if (!value_list_head) {
                 value_list_head = value_list_end = jit_value;
                 jit_value->prev = jit_value->next = NULL;
@@ -386,8 +386,13 @@ handle_func_return(JitCompContext *cc, JitBlock *block)
     GEN_INSN(RETURNBC, NEW_CONST(I32, JIT_INTERP_ACTION_NORMAL), 0, 0);
 }
 
+/**
+ * is_block_polymorphic: whether current block's stack is in polymorphic state,
+ * if the opcode is one of unreachable/br/br_table/return, stack is marked
+ * to polymorphic state until the block's 'end' opcode is processed
+ */
 static bool
-handle_op_end(JitCompContext *cc, uint8 **p_frame_ip, bool from_same_block)
+handle_op_end(JitCompContext *cc, uint8 **p_frame_ip, bool is_block_polymorphic)
 {
     JitFrame *jit_frame = cc->jit_frame;
     JitBlock *block, *block_prev;
@@ -395,7 +400,7 @@ handle_op_end(JitCompContext *cc, uint8 **p_frame_ip, bool from_same_block)
     JitInsn *insn;
 
     /* Check block stack */
-    if (!(block = cc->block_stack.block_list_end)) {
+    if (!(block = jit_block_stack_top(&cc->block_stack))) {
         jit_set_last_error(cc, "WASM block stack underflow");
         return false;
     }
@@ -407,6 +412,7 @@ handle_op_end(JitCompContext *cc, uint8 **p_frame_ip, bool from_same_block)
         if (block->label_type == LABEL_TYPE_FUNCTION) {
             handle_func_return(cc, block);
             SET_BB_END_BCIP(cc->cur_basic_block, *p_frame_ip - 1);
+            clear_values(jit_frame);
         }
         else if (block->result_count > 0) {
             JitValue *value_list_head = NULL, *value_list_end = NULL;
@@ -460,9 +466,13 @@ handle_op_end(JitCompContext *cc, uint8 **p_frame_ip, bool from_same_block)
         CREATE_BASIC_BLOCK(block->basic_block_end);
         SET_BB_END_BCIP(cc->cur_basic_block, *p_frame_ip - 1);
         SET_BB_BEGIN_BCIP(block->basic_block_end, *p_frame_ip);
-        if (from_same_block)
+        /* No need to create 'JMP' insn if block is in stack polymorphic
+           state, as previous br/br_table opcode has created 'JMP' insn
+           to this end basic block */
+        if (!is_block_polymorphic) {
             /* Jump to the end basic block */
             BUILD_BR(block->basic_block_end);
+        }
 
         /* Patch the INSNs which jump to this basic block */
         incoming_insn = block->incoming_insns_for_end_bb;
@@ -490,6 +500,7 @@ handle_op_end(JitCompContext *cc, uint8 **p_frame_ip, bool from_same_block)
         if (block->label_type == LABEL_TYPE_FUNCTION) {
             handle_func_return(cc, block);
             SET_BB_END_BCIP(cc->cur_basic_block, *p_frame_ip - 1);
+            clear_values(jit_frame);
         }
         else {
             if (!load_block_results(cc, block)) {
@@ -506,10 +517,16 @@ fail:
     return false;
 }
 
+/**
+ * is_block_polymorphic: whether current block's stack is in polymorphic state,
+ * if the opcode is one of unreachable/br/br_table/return, stack is marked
+ * to polymorphic state until the block's 'end' opcode is processed
+ */
 static bool
-handle_op_else(JitCompContext *cc, uint8 **p_frame_ip)
+handle_op_else(JitCompContext *cc, uint8 **p_frame_ip,
+               bool is_block_polymorphic)
 {
-    JitBlock *block = cc->block_stack.block_list_end;
+    JitBlock *block = jit_block_stack_top(&cc->block_stack);
     JitFrame *jit_frame = cc->jit_frame;
     JitInsn *insn;
 
@@ -527,7 +544,7 @@ handle_op_else(JitCompContext *cc, uint8 **p_frame_ip)
         /* The if branch is handled like OP_BLOCK (cond is const and != 0),
            just skip the else branch and handle OP_END */
         *p_frame_ip = block->wasm_code_end + 1;
-        return handle_op_end(cc, p_frame_ip, true);
+        return handle_op_end(cc, p_frame_ip, false);
     }
     else {
         /* Has else branch and need to translate else branch */
@@ -537,14 +554,19 @@ handle_op_else(JitCompContext *cc, uint8 **p_frame_ip)
         /* Clear frame values */
         clear_values(jit_frame);
 
-        /* Jump to end basic block */
-        if (!(insn = GEN_INSN(JMP, 0))) {
-            jit_set_last_error(cc, "generate jmp insn failed");
-            return false;
-        }
-        if (!jit_block_add_incoming_insn(block, insn)) {
-            jit_set_last_error(cc, "add incoming insn failed");
-            return false;
+        /* No need to create 'JMP' insn if block is in stack polymorphic
+           state, as previous br/br_table opcode has created 'JMP' insn
+           to this end basic block */
+        if (!is_block_polymorphic) {
+            /* Jump to end basic block */
+            if (!(insn = GEN_INSN(JMP, 0))) {
+                jit_set_last_error(cc, "generate jmp insn failed");
+                return false;
+            }
+            if (!jit_block_add_incoming_insn(block, insn)) {
+                jit_set_last_error(cc, "add incoming insn failed");
+                return false;
+            }
         }
 
         /* Clear value stack, restore param values and
@@ -578,29 +600,34 @@ fail:
 static bool
 handle_next_reachable_block(JitCompContext *cc, uint8 **p_frame_ip)
 {
-    JitBlock *block = cc->block_stack.block_list_end, *block_prev;
+    JitBlock *block = jit_block_stack_top(&cc->block_stack);
 
     bh_assert(block);
 
     do {
-        block_prev = block->prev;
-
         if (block->label_type == LABEL_TYPE_IF
             && block->incoming_insn_for_else_bb
             && *p_frame_ip <= block->wasm_code_else) {
             /* Else branch hasn't been translated,
                start to translate the else branch */
             *p_frame_ip = block->wasm_code_else + 1;
-            return handle_op_else(cc, p_frame_ip);
+            /* Restore jit frame's sp to block's sp begin */
+            cc->jit_frame->sp = block->frame_sp_begin;
+            return handle_op_else(cc, p_frame_ip, true);
         }
         else if (block->incoming_insns_for_end_bb) {
             *p_frame_ip = block->wasm_code_end + 1;
-            return handle_op_end(cc, p_frame_ip, false);
+            /* Restore jit frame's sp to block's sp end  */
+            cc->jit_frame->sp =
+                block->frame_sp_begin
+                + wasm_get_cell_num(block->result_types, block->result_count);
+            return handle_op_end(cc, p_frame_ip, true);
         }
         else {
+            *p_frame_ip = block->wasm_code_end + 1;
             jit_block_stack_pop(&cc->block_stack);
             jit_block_destroy(block);
-            block = block_prev;
+            block = jit_block_stack_top(&cc->block_stack);
         }
     } while (block != NULL);
 
@@ -619,7 +646,7 @@ jit_compile_op_block(JitCompContext *cc, uint8 **p_frame_ip,
     uint8 *else_addr, *end_addr;
 
     /* Check block stack */
-    if (!cc->block_stack.block_list_end) {
+    if (!jit_block_stack_top(&cc->block_stack)) {
         jit_set_last_error(cc, "WASM block stack underflow");
         return false;
     }
@@ -736,13 +763,13 @@ fail:
 bool
 jit_compile_op_else(JitCompContext *cc, uint8 **p_frame_ip)
 {
-    return handle_op_else(cc, p_frame_ip);
+    return handle_op_else(cc, p_frame_ip, false);
 }
 
 bool
 jit_compile_op_end(JitCompContext *cc, uint8 **p_frame_ip)
 {
-    return handle_op_end(cc, p_frame_ip, true);
+    return handle_op_end(cc, p_frame_ip, false);
 }
 
 static bool
@@ -757,7 +784,7 @@ handle_op_br(JitCompContext *cc, uint32 br_depth, uint8 **p_frame_ip)
     uint32 offset;
 
     /* Check block stack */
-    if (!(block = cc->block_stack.block_list_end)) {
+    if (!(block = jit_block_stack_top(&cc->block_stack))) {
         jit_set_last_error(cc, "WASM block stack underflow");
         return false;
     }
@@ -795,9 +822,14 @@ handle_op_br(JitCompContext *cc, uint32 br_depth, uint8 **p_frame_ip)
 #else
         GEN_INSN(ADD, frame_sp_dst, cc->fp_reg, NEW_CONST(I32, offset));
 #endif
-    }
 
-    gen_commit_values(jit_frame, jit_frame->lp, block->frame_sp_begin);
+        /* No need to commit results as they will be copied to dest block */
+        gen_commit_values(jit_frame, jit_frame->lp, block->frame_sp_begin);
+    }
+    else {
+        /* Commit all including results as they won't be copied */
+        gen_commit_values(jit_frame, jit_frame->lp, jit_frame->sp);
+    }
 
     if (block_dst->label_type == LABEL_TYPE_LOOP) {
         if (copy_arities) {
@@ -949,6 +981,7 @@ jit_compile_op_return(JitCompContext *cc, uint8 **p_frame_ip)
 
     handle_func_return(cc, block_func);
     SET_BB_END_BCIP(cc->cur_basic_block, *p_frame_ip - 1);
+    clear_values(cc->jit_frame);
 
     return handle_next_reachable_block(cc, p_frame_ip);
 }

core/iwasm/fast-jit/fe/jit_emit_parametric.c

@@ -13,17 +13,17 @@ pop_value_from_wasm_stack(JitCompContext *cc, bool is_32bit, JitReg *p_value,
     JitValue *jit_value;
     uint8 type;
 
-    if (!cc->block_stack.block_list_end) {
+    if (!jit_block_stack_top(&cc->block_stack)) {
         jit_set_last_error(cc, "WASM block stack underflow.");
         return false;
     }
-    if (!cc->block_stack.block_list_end->value_stack.value_list_end) {
+    if (!jit_block_stack_top(&cc->block_stack)->value_stack.value_list_end) {
         jit_set_last_error(cc, "WASM data stack underflow.");
         return false;
     }
 
-    jit_value =
-        jit_value_stack_pop(&cc->block_stack.block_list_end->value_stack);
+    jit_value = jit_value_stack_pop(
+        &jit_block_stack_top(&cc->block_stack)->value_stack);
     type = jit_value->type;
 
     if (p_type != NULL) {

core/iwasm/fast-jit/jit_ir.c

@@ -1347,6 +1347,12 @@ jit_block_stack_push(JitBlockStack *stack, JitBlock *block)
     }
 }
 
+JitBlock *
+jit_block_stack_top(JitBlockStack *stack)
+{
+    return stack->block_list_end;
+}
+
 JitBlock *
 jit_block_stack_pop(JitBlockStack *stack)
 {
@@ -1432,17 +1438,17 @@ jit_cc_pop_value(JitCompContext *cc, uint8 type, JitReg *p_value)
     JitValue *jit_value = NULL;
     JitReg value = 0;
 
-    if (!cc->block_stack.block_list_end) {
+    if (!jit_block_stack_top(&cc->block_stack)) {
         jit_set_last_error(cc, "WASM block stack underflow");
         return false;
     }
-    if (!cc->block_stack.block_list_end->value_stack.value_list_end) {
+    if (!jit_block_stack_top(&cc->block_stack)->value_stack.value_list_end) {
         jit_set_last_error(cc, "WASM data stack underflow");
         return false;
     }
 
-    jit_value =
-        jit_value_stack_pop(&cc->block_stack.block_list_end->value_stack);
+    jit_value = jit_value_stack_pop(
+        &jit_block_stack_top(&cc->block_stack)->value_stack);
     bh_assert(jit_value);
 
     if (jit_value->type != to_stack_value_type(type)) {
@@ -1480,7 +1486,7 @@ jit_cc_push_value(JitCompContext *cc, uint8 type, JitReg value)
 {
     JitValue *jit_value;
 
-    if (!cc->block_stack.block_list_end) {
+    if (!jit_block_stack_top(&cc->block_stack)) {
         jit_set_last_error(cc, "WASM block stack underflow");
         return false;
     }
@@ -1492,7 +1498,7 @@ jit_cc_push_value(JitCompContext *cc, uint8 type, JitReg value)
 
     jit_value->type = to_stack_value_type(type);
     jit_value->value = value;
-    jit_value_stack_push(&cc->block_stack.block_list_end->value_stack,
+    jit_value_stack_push(&jit_block_stack_top(&cc->block_stack)->value_stack,
                          jit_value);
 
     switch (jit_value->type) {

core/iwasm/fast-jit/jit_ir.h

@@ -1766,6 +1766,9 @@ jit_value_stack_pop(JitValueStack *stack);
 void
 jit_value_stack_destroy(JitValueStack *stack);
 
+JitBlock *
+jit_block_stack_top(JitBlockStack *stack);
+
 void
 jit_block_stack_push(JitBlockStack *stack, JitBlock *block);