From a504452b7b33025aef01da5a8766b4459ba6318e Mon Sep 17 00:00:00 2001
From: John Richards <john@pushpoplabs.com>
Date: Thu, 7 Oct 2021 01:51:53 -0400
Subject: [PATCH] Adds support for LIST_EXTEND opcode

Addresses #199
---
 ASTree.cpp                         |  27 +++++++++++++++++++++++++++
 tests/compiled/list_extend.3.9.pyc | Bin 0 -> 273 bytes
 tests/input/list_extend.py         |   8 ++++++++
 tests/tokenized/list_extend.txt    |   8 ++++++++
 4 files changed, 43 insertions(+)
 create mode 100644 tests/compiled/list_extend.3.9.pyc
 create mode 100644 tests/input/list_extend.py
 create mode 100644 tests/tokenized/list_extend.txt

diff --git a/ASTree.cpp b/ASTree.cpp
index ba31382..c92480e 100644
--- a/ASTree.cpp
+++ b/ASTree.cpp
@@ -1455,6 +1455,33 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                 }
             }
             break;
+        case Pyc::LIST_EXTEND_A:
+            {
+                PycRef<ASTNode> rhs = stack.top();
+                stack.pop();
+                PycRef<ASTList> lhs = stack.top().cast<ASTList>();
+                stack.pop();
+
+                if (rhs.type() != ASTNode::NODE_OBJECT) {
+                    fprintf(stderr, "Unsupported argument found for LIST_EXTEND\n");
+                    break;
+                }
+
+                // I've only ever seen this be a SMALL_TUPLE, but let's be careful...
+                PycRef<PycObject> obj = rhs.cast<ASTObject>()->object();
+                if (obj->type() != PycObject::TYPE_TUPLE && obj->type() != PycObject::TYPE_SMALL_TUPLE) {
+                    fprintf(stderr, "Unsupported argument type found for LIST_EXTEND\n");
+                    break;
+                }
+
+                ASTList::value_t result = lhs->values();
+                for (const auto& it : obj.cast<PycTuple>()->values()) {
+                    result.push_back(new ASTObject(it));
+                }
+
+                stack.push(new ASTList(result));
+            }
+            break;
         case Pyc::LOAD_ATTR_A:
             {
                 PycRef<ASTNode> name = stack.top();
diff --git a/tests/compiled/list_extend.3.9.pyc b/tests/compiled/list_extend.3.9.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..19fcbe6f902b6a42ce3b0186aced1dd8a75777c8
GIT binary patch
literal 273
zcmYk0I}XAy42I*TFCJy+EwXfh1*t-upbJ|EP{AQWDi5_q0*Mo_asxK5l$9x0V8Uj?
zR=z*}^~-XVBz-~kxG$?4>6?fD@@SaTYjc7TD4tNV0x6(}(s-qZzJXF=K=ifHKwe-(
z8B#9fJo*cvA|)Nr62~i@Bq(WN;sPE+M)dPs%>N#PPb)Z8Sy)@xC|MMRikboW&5)50
zET9E#VQay>M3b`5_@wo&?N<8I=>ultY?_^}HqQ4)JlA^0r?ewN4BR1g3_g4TGvhT0

literal 0
HcmV?d00001

diff --git a/tests/input/list_extend.py b/tests/input/list_extend.py
new file mode 100644
index 0000000..416a3c7
--- /dev/null
+++ b/tests/input/list_extend.py
@@ -0,0 +1,8 @@
+a = [1, 2, 3]
+b = ['4', 5, 6]
+c = [7, (8, 9, 10)]
+d = [None]
+e = [('a', 'b', 'c')]
+f = [a, b]
+g = []
+a = [None, None, None]
diff --git a/tests/tokenized/list_extend.txt b/tests/tokenized/list_extend.txt
new file mode 100644
index 0000000..9bb99d4
--- /dev/null
+++ b/tests/tokenized/list_extend.txt
@@ -0,0 +1,8 @@
+a = [ 1 , 2 , 3 ] <EOL>
+b = [ '4' , 5 , 6 ] <EOL>
+c = [ 7 , ( 8 , 9 , 10 ) ] <EOL>
+d = [ None ] <EOL>
+e = [ ( 'a' , 'b' , 'c' ) ] <EOL>
+f = [ a , b ] <EOL>
+g = [ ] <EOL>
+a = [ None , None , None ] <EOL>