Parse EXTENDED_ARG in Runtime::collectAttributes

Add and modify tests.

Fix #428

Author: tekknolagi

runtime/runtime-test.cpp

@@ -1438,9 +1438,33 @@ TEST_F(RuntimeTest, CollectAttributes) {
   result = dictAtByStr(thread_, attributes, baz);
   ASSERT_TRUE(result.isStr());
   EXPECT_TRUE(Str::cast(*result).equals(*baz));
+
+  // Bytecode for the snippet:
+  //
+  //   def __init__(self):
+  //       self.foo = 100
+  //
+  // Manually add duplicate and useless EXTENDED_ARG intentionally to ensure
+  // that collectAttributes can handle it.
+  const byte bytecode2[] = {LOAD_CONST,   0, EXTENDED_ARG, 0, EXTENDED_ARG, 0,
+                            EXTENDED_ARG, 0, LOAD_FAST,    0, STORE_ATTR,   0,
+                            RETURN_VALUE, 0};
+  Code code2(&scope, newCodeWithBytesConstsNamesLocals(bytecode2, consts, names,
+                                                       &locals));
+
+  attributes = runtime_->newDict();
+  runtime_->collectAttributes(code2, attributes);
+
+  // We should have collected a single attribute: 'foo'
+  EXPECT_EQ(attributes.numItems(), 1);
+
+  // Check that we collected 'foo'
+  result = dictAtByStr(thread_, attributes, foo);
+  ASSERT_TRUE(result.isStr());
+  EXPECT_TRUE(Str::cast(*result).equals(*foo));
 }
 
-TEST_F(RuntimeTest, CollectAttributesWithExtendedArg) {
+TEST_F(RuntimeTest, CollectAttributesWithExtendedArgAccumulates) {
   HandleScope scope(thread_);
 
   Str foo(&scope, runtime_->newStrFromCStr("foo"));
@@ -1461,22 +1485,16 @@ TEST_F(RuntimeTest, CollectAttributesWithExtendedArg) {
   //
   // There is an additional LOAD_FAST that is preceded by an EXTENDED_ARG
   // that must be skipped.
-  const byte bytecode[] = {LOAD_CONST, 0, EXTENDED_ARG, 1, LOAD_FAST, 0,
-                           STORE_ATTR, 1, LOAD_CONST,   0, LOAD_FAST, 0,
-                           STORE_ATTR, 0, RETURN_VALUE, 0};
+  const byte bytecode[] = {LOAD_CONST, 0, EXTENDED_ARG, 1, LOAD_FAST,    0,
+                           STORE_ATTR, 1, LOAD_CONST,   0, RETURN_VALUE, 0};
   Code code(&scope, newCodeWithBytesConstsNamesLocals(bytecode, consts, names,
                                                       &locals));
 
   Dict attributes(&scope, runtime_->newDict());
   runtime_->collectAttributes(code, attributes);
 
-  // We should have collected a single attribute: 'foo'
-  EXPECT_EQ(attributes.numItems(), 1);
-
-  // Check that we collected 'foo'
-  Object result(&scope, dictAtByStr(thread_, attributes, foo));
-  ASSERT_TRUE(result.isStr());
-  EXPECT_TRUE(Str::cast(*result).equals(*foo));
+  // We should have collected no attributes because EXTENDED_ARG makes 0x0100
+  EXPECT_EQ(attributes.numItems(), 0);
 }
 
 TEST_F(RuntimeTest, GetTypeConstructor) {

runtime/runtime.cpp

@@ -2879,15 +2879,17 @@ void Runtime::collectAttributes(const Code& code, const Dict& attributes) {
   Tuple names(&scope, code.names());
 
   word len = bc.length();
-  for (word i = 0; i < len - 3; i += 2) {
-    // If the current instruction is EXTENDED_ARG we must skip it and the next
-    // instruction.
-    if (bc.byteAt(i) == Bytecode::EXTENDED_ARG) {
-      i += 2;
-      continue;
+  for (word i = 0; i < len - (kCompilerCodeUnitSize + 1);
+       i += kCompilerCodeUnitSize) {
+    byte op = bc.byteAt(i);
+    int32_t arg = bc.byteAt(i + 1);
+    while (op == Bytecode::EXTENDED_ARG) {
+      i += kCompilerCodeUnitSize;
+      op = bc.byteAt(i);
+      arg = (arg << kBitsPerByte) | bc.byteAt(i + 1);
     }
     // Check for LOAD_FAST 0 (self)
-    if (bc.byteAt(i) != Bytecode::LOAD_FAST || bc.byteAt(i + 1) != 0) {
+    if (op != Bytecode::LOAD_FAST || arg != 0) {
       continue;
     }
     // Followed by a STORE_ATTR