support find ins unit by bytecode frame

src/PikaVM.c

@@ -1733,33 +1733,58 @@ static Arg* VM_instruction_handler_RET(PikaObj* self,
     return NULL;
 }
 
+static InstructUnit* _find_ins_unit_up(ByteCodeFrame* bcframe,
+                                       int32_t pc_start,
+                                       enum InstructIndex index,
+                                       int32_t* p_offset) {
+    /* find super class */
+    int instructArray_size = instructArray_getSize(&(bcframe->instruct_array));
+    while (1) {
+        *p_offset -= instructUnit_getSize();
+        if (pc_start + *p_offset >= instructArray_size) {
+            return 0;
+        }
+        InstructUnit* unit = instructArray_getByOffset(
+            &(bcframe->instruct_array), pc_start + *p_offset);
+        if (CLS == instructUnit_getInstructIndex(unit)) {
+            return unit;
+        }
+    }
+    return NULL;
+}
+
+static InstructUnit* _find_ins_unit_down(ByteCodeFrame* bcframe,
+                                         int32_t pc_start,
+                                         enum InstructIndex index,
+                                         int32_t* p_offset) {
+    /* find super class */
+    int instructArray_size = instructArray_getSize(&(bcframe->instruct_array));
+    while (1) {
+        *p_offset += instructUnit_getSize();
+        if (pc_start + *p_offset >= instructArray_size) {
+            return 0;
+        }
+        InstructUnit* unit = instructArray_getByOffset(
+            &(bcframe->instruct_array), pc_start + *p_offset);
+        if (index == instructUnit_getInstructIndex(unit)) {
+            return unit;
+        }
+    }
+    return NULL;
+}
+
 #if !PIKA_NANO_ENABLE
-static char* _find_super_class_name(VMState* vm) {
+static char* _find_super_class_name(ByteCodeFrame* bcframe, int32_t pc_start) {
     /* find super class */
     int offset = 0;
     char* super_class_name = NULL;
-    while (1) {
-        offset -= instructUnit_getSize();
-        if (vm->pc + offset >= (int)VMState_getInstructArraySize(vm)) {
-            return 0;
-        }
-        if ((CLS == VMstate_getInstructWithOffset(vm, offset))) {
-            break;
-        }
-    }
-
-    while (1) {
-        offset += instructUnit_getSize();
-        if (vm->pc + offset >= (int)VMState_getInstructArraySize(vm)) {
-            return 0;
-        }
-        if ((RUN == instructUnit_getInstructIndex(
-                        VMState_getInstructUnitWithOffset(vm, offset)))) {
-            super_class_name = VMState_getConstWithOffset(vm, offset);
+    _find_ins_unit_up(bcframe, pc_start, CLS, &offset);
+    InstructUnit* unit_run =
+        _find_ins_unit_down(bcframe, pc_start, RUN, &offset);
+    super_class_name = constPool_getByOffset(
+        &(bcframe->const_pool), instructUnit_getConstPoolIndex(unit_run));
     return super_class_name;
 }
-    }
-}
 #endif
 
 #if !PIKA_NANO_ENABLE
@@ -1841,7 +1866,7 @@ static Arg* VM_instruction_handler_RUN(PikaObj* self,
 #if !PIKA_NANO_ENABLE
     /* support for super() */
     if (strEqu(sRunPath, "super")) {
-        sRunPath = _find_super_class_name(vm);
+        sRunPath = _find_super_class_name(vm->bytecode_frame, vm->pc);
         sArgName = strPointToLastToken(sRunPath, '.');
         vm->in_super = PIKA_TRUE;
         vm->super_invoke_deepth = VMState_getInvokeDeepthNow(vm);

test/except-test.cpp

@@ -177,8 +177,7 @@ TEST(except, except_break) {
             "        print('in excepton')\n"
             "        break\n"
             "print(sum)\n"
-            "\n"
-            );
+            "\n");
     /* collect */
     /* assert */
     EXPECT_EQ(obj_getInt(pikaMain, "sum"), 6);
@@ -204,6 +203,10 @@ TEST(except, while_try_while) {
     EXPECT_EQ(pikaMemNow(), 0);
 }
 
+TEST_RUN_SINGLE_FILE(except,
+                     try_while_return,
+                     "test/python/except/try_while_return.py");
+
 #endif
 
 TEST_END

test/python/except/try_while_return.py

@@ -0,0 +1,85 @@
+def try_while_return():
+    try:
+        i = 0
+        while i < 5:
+            i += 1
+            if i == 3:
+                return i
+    except Exception as e:
+        return "Error occurred: %s" % str(e)
+
+
+def try_finally_while_return():
+    try:
+        i = 0
+        while i < 5:
+            i += 1
+            if i == 3:
+                return i
+    finally:
+        pass
+
+
+def while_try_return():
+    i = 0
+    while i < 5:
+        try:
+            if i == 3:
+                return i
+            i += 1
+        except Exception as e:
+            return "Error occurred: %s" % str(e)
+
+
+def while_try_finally_return():
+    i = 0
+    while i < 5:
+        try:
+            if i == 3:
+                return i
+            i += 1
+        finally:
+            pass
+
+
+def try_except_return():
+    try:
+        raise ValueError("An error occurred")
+    except ValueError as e:
+        return "Caught an error: %s" % str(e)
+
+def try_except_return_sinple():
+    try:
+        raise
+    except:
+        return "Caught an error"
+
+
+def while_try_except_return():
+    i = 0
+    while i < 5:
+        try:
+            if i == 3:
+                raise ValueError("An error occurred at 3")
+            i += 1
+        except ValueError as e:
+            return "Caught an error: %s" % str(e)
+
+def while_try_except_return_simple():
+    i = 0
+    while i < 5:
+        try:
+            if i == 3:
+                raise
+            i += 1
+        except:
+            return "Caught an error"
+
+assert try_while_return() == 3
+assert while_try_return() == 3
+assert try_except_return_sinple() == "Caught an error"
+assert while_try_except_return_simple() == "Caught an error"
+# assert while_try_except_return() == "Caught an error: An error occurred at 3"
+# assert try_except_return() == "Caught an error: An error occurred"
+# assert try_finally_while_return() == 3
+# assert while_try_finally_return() == 3