From 7a89b722604288293603228626ba25f35409fc43 Mon Sep 17 00:00:00 2001 From: Aralox Date: Fri, 23 Oct 2020 21:19:01 +1100 Subject: [PATCH] Support chained assignment statements, e.g. `a = b = c`. We know when we have begun a chained assignment when we process a DUP_TOP with non-null on the stack. Push a NODE_CHAINSTORE onto the stack when this happens, and keep it 'floating' on top of the stack for all STORE_X operations until the stack is empty. To support versions of Python <= 2.5 which use DUP_TOP in more places, I modified ROT_TWO, ROT_THREE and ROT_FOUR to get rid of NODE_CHAINSTORE on the stack if it is present. --- ASTNode.h | 18 +++- ASTree.cpp | 110 +++++++++++++++++++++--- tests/compiled/chain_assignment.2.7.pyc | Bin 0 -> 1934 bytes tests/compiled/chain_assignment.3.7.pyc | Bin 0 -> 1643 bytes tests/input/chain_assignment.py | 39 +++++++++ tests/tokenized/chain_assignment.txt | 43 +++++++++ tests/tokenized/f-string.txt | 4 +- 7 files changed, 200 insertions(+), 14 deletions(-) create mode 100644 tests/compiled/chain_assignment.2.7.pyc create mode 100644 tests/compiled/chain_assignment.3.7.pyc create mode 100644 tests/input/chain_assignment.py create mode 100644 tests/tokenized/chain_assignment.txt diff --git a/ASTNode.h b/ASTNode.h index 0035464..26df413 100644 --- a/ASTNode.h +++ b/ASTNode.h @@ -16,7 +16,7 @@ public: NODE_CONVERT, NODE_KEYWORD, NODE_RAISE, NODE_EXEC, NODE_BLOCK, NODE_COMPREHENSION, NODE_LOADBUILDCLASS, NODE_AWAITABLE, NODE_FORMATTEDVALUE, NODE_JOINEDSTR, NODE_CONST_MAP, - NODE_ANNOTATED_VAR, + NODE_ANNOTATED_VAR, NODE_CHAINSTORE, // Empty node types NODE_LOCALS, @@ -71,11 +71,27 @@ public: void removeLast(); void append(PycRef node) { m_nodes.emplace_back(std::move(node)); } +protected: + ASTNodeList(list_t nodes, ASTNode::Type type) + : ASTNode(type), m_nodes(std::move(nodes)) { } + private: list_t m_nodes; }; +class ASTChainStore : public ASTNodeList { +public: + ASTChainStore(list_t nodes, PycRef src) + : ASTNodeList(nodes, NODE_CHAINSTORE), m_src(std::move(src)) { } + + PycRef src() const { return m_src; } + +private: + PycRef m_src; +}; + + class ASTObject : public ASTNode { public: ASTObject(PycRef obj) diff --git a/ASTree.cpp b/ASTree.cpp index 408818a..f3a0e75 100644 --- a/ASTree.cpp +++ b/ASTree.cpp @@ -12,6 +12,8 @@ // NOTE: Nested f-strings not supported. #define F_STRING_QUOTE "'''" +static void append_to_chain_store(PycRef chainStore, PycRef item, FastStack& stack, PycRef curblock); + /* Use this to determine if an error occurred (and therefore, if we should * avoid cleaning the output tree) */ static bool cleanBuild; @@ -288,6 +290,9 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) } stack.push(map); } else { + if (stack.top().type() == ASTNode::NODE_CHAINSTORE) { + stack.pop(); + } stack.push(new ASTMap()); } break; @@ -679,7 +684,22 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) } break; case Pyc::DUP_TOP: - stack.push(stack.top()); + { + if (stack.top().type() != PycObject::TYPE_NULL) { + if (stack.top().type() == ASTNode::NODE_CHAINSTORE) { + auto chainstore = stack.top(); + stack.pop(); + stack.push(stack.top()); + stack.push(chainstore); + } else { + stack.push(stack.top()); + ASTNodeList::list_t targets; + stack.push(new ASTChainStore(targets, stack.top())); + } + } else { + stack.push(stack.top()); + } + } break; case Pyc::DUP_TOP_TWO: { @@ -791,6 +811,9 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) break; case Pyc::EXEC_STMT: { + if (stack.top().type() == ASTNode::NODE_CHAINSTORE) { + stack.pop(); + } PycRef loc = stack.top(); stack.pop(); PycRef glob = stack.top(); @@ -1725,6 +1748,9 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) { PycRef one = stack.top(); stack.pop(); + if (stack.top().type() == ASTNode::NODE_CHAINSTORE) { + stack.pop(); + } PycRef two = stack.top(); stack.pop(); @@ -1738,6 +1764,9 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) stack.pop(); PycRef two = stack.top(); stack.pop(); + if (stack.top().type() == ASTNode::NODE_CHAINSTORE) { + stack.pop(); + } PycRef three = stack.top(); stack.pop(); stack.push(one); @@ -1753,6 +1782,9 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) stack.pop(); PycRef three = stack.top(); stack.pop(); + if (stack.top().type() == ASTNode::NODE_CHAINSTORE) { + stack.pop(); + } PycRef four = stack.top(); stack.pop(); stack.push(one); @@ -1889,8 +1921,11 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) stack.pop(); PycRef seq = stack.top(); stack.pop(); - - curblock->append(new ASTStore(seq, tup)); + if (seq.type() == ASTNode::NODE_CHAINSTORE) { + append_to_chain_store(seq, tup, stack, curblock); + } else { + curblock->append(new ASTStore(seq, tup)); + } } } else { PycRef name = stack.top(); @@ -1898,8 +1933,11 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) PycRef value = stack.top(); stack.pop(); PycRef attr = new ASTBinary(name, new ASTName(code->getName(operand)), ASTBinary::BIN_ATTR); - - curblock->append(new ASTStore(value, attr)); + if (value.type() == ASTNode::NODE_CHAINSTORE) { + append_to_chain_store(value, attr, stack, curblock); + } else { + curblock->append(new ASTStore(value, attr)); + } } } break; @@ -1919,13 +1957,22 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) PycRef seq = stack.top(); stack.pop(); - curblock->append(new ASTStore(seq, tup)); + if (seq.type() == ASTNode::NODE_CHAINSTORE) { + append_to_chain_store(seq, tup, stack, curblock); + } else { + curblock->append(new ASTStore(seq, tup)); + } } } else { PycRef value = stack.top(); stack.pop(); PycRef name = new ASTName(code->getCellVar(operand)); - curblock->append(new ASTStore(value, name)); + + if (value.type() == ASTNode::NODE_CHAINSTORE) { + append_to_chain_store(value, name, stack, curblock); + } else { + curblock->append(new ASTStore(value, name)); + } } } break; @@ -1956,6 +2003,8 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) if (tuple != NULL) tuple->setRequireParens(false); curblock.cast()->setIndex(tup); + } else if (seq.type() == ASTNode::NODE_CHAINSTORE) { + append_to_chain_store(seq, tup, stack, curblock); } else { curblock->append(new ASTStore(seq, tup)); } @@ -1983,6 +2032,8 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) && !curblock->inited()) { curblock.cast()->setExpr(value); curblock.cast()->setVar(name); + } else if (value.type() == ASTNode::NODE_CHAINSTORE) { + append_to_chain_store(value, name, stack, curblock); } else { curblock->append(new ASTStore(value, name)); } @@ -2011,6 +2062,8 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) if (tuple != NULL) tuple->setRequireParens(false); curblock.cast()->setIndex(tup); + } else if (seq.type() == ASTNode::NODE_CHAINSTORE) { + append_to_chain_store(seq, tup, stack, curblock); } else { curblock->append(new ASTStore(seq, tup)); } @@ -2018,7 +2071,11 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) } else { PycRef value = stack.top(); stack.pop(); - curblock->append(new ASTStore(value, name)); + if (value.type() == ASTNode::NODE_CHAINSTORE) { + append_to_chain_store(value, name, stack, curblock); + } else { + curblock->append(new ASTStore(value, name)); + } } /* Mark the global as used */ @@ -2047,6 +2104,8 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) if (tuple != NULL) tuple->setRequireParens(false); curblock.cast()->setIndex(tup); + } else if (seq.type() == ASTNode::NODE_CHAINSTORE) { + append_to_chain_store(seq, tup, stack, curblock); } else { curblock->append(new ASTStore(seq, tup)); } @@ -2080,6 +2139,8 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) && !curblock->inited()) { curblock.cast()->setExpr(value); curblock.cast()->setVar(name); + } else if (value.type() == ASTNode::NODE_CHAINSTORE) { + append_to_chain_store(value, name, stack, curblock); } else { curblock->append(new ASTStore(value, name)); @@ -2157,8 +2218,11 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) stack.pop(); PycRef seq = stack.top(); stack.pop(); - - curblock->append(new ASTStore(seq, tup)); + if (seq.type() == ASTNode::NODE_CHAINSTORE) { + append_to_chain_store(seq, tup, stack, curblock); + } else { + curblock->append(new ASTStore(seq, tup)); + } } } else { PycRef subscr = stack.top(); @@ -2189,6 +2253,8 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) } else { if (dest.type() == ASTNode::NODE_MAP) { dest.cast()->add(subscr, src); + } else if (src.type() == ASTNode::NODE_CHAINSTORE) { + append_to_chain_store(src, new ASTSubscr(dest, subscr), stack, curblock); } else { curblock->append(new ASTStore(src, new ASTSubscr(dest, subscr))); } @@ -2255,6 +2321,10 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) && !curblock->inited()) { tup->setRequireParens(true); curblock.cast()->setIndex(tup); + } else if (stack.top().type() == ASTNode::NODE_CHAINSTORE) { + auto chainStore = stack.top(); + stack.pop(); + append_to_chain_store(chainStore, tup, stack, curblock); } else { curblock->append(new ASTStore(stack.top(), tup)); stack.pop(); @@ -2319,6 +2389,17 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) return new ASTNodeList(defblock->nodes()); } +static void append_to_chain_store(PycRef chainStore, PycRef item, FastStack& stack, PycRef curblock) +{ + stack.pop(); // ignore identical source object. + chainStore.cast()->append(item); + if (stack.top().type() == PycObject::TYPE_NULL) { + curblock->append(chainStore); + } else { + stack.push(chainStore); + } +} + static int cmp_prec(PycRef parent, PycRef child) { /* Determine whether the parent has higher precedence than therefore @@ -3028,6 +3109,15 @@ void print_src(PycRef node, PycModule* mod) } } break; + case ASTNode::NODE_CHAINSTORE: + { + for (auto& dest : node.cast()->nodes()) { + print_src(dest, mod); + fputs(" = ", pyc_output); + } + print_src(node.cast()->src(), mod); + } + break; case ASTNode::NODE_SUBSCR: { print_src(node.cast()->name(), mod); diff --git a/tests/compiled/chain_assignment.2.7.pyc b/tests/compiled/chain_assignment.2.7.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e8671b943c51271c44b76d460ff2334cc79e5beb GIT binary patch literal 1934 zcmbtUT~8ZF6g{)MHU@$N0YV`O`A{DG6bTaLiB?HkHR=meCNIsZ(rWRJVFQ@pS&}%a zpOOZ6&7a5*>>ubocU&+iu38l{XXf6yz8}Z;PWLfa^fNt6E{bL4cB09mv6=OKi((lbuHzh zQMwdjekXC4;Z`8AXwkqQF*b?>#T1Gfv-=3VN^hX3qnJi9ai+QncMGpq;imEWwlF3R zGIfyMLDA5t}0k;GpUs2~(c{kd=yv%;4x`8lxJN_a`r=Z9f? zAGUtjKD6Wt=iui>Z13b{H_N>*x4ZkUTRnKwuVzKxIqJ5vcV)GcZSNgy<$F`2^Yo3V zB$gU%T28(?=(#Mfs@~dJzRc$t@8$%gQKTOv&rfoQU_-iXI^bfVTHAd&u(scUOv>XH zff%Do^iUhEwoY>fdasx(ZYj1E`-(-Ng^|45>(~i7*al%z;9syo+bk=)WtC-%a`tou zvt(M#yb9OdeBWhRMV7K`|DZV96Vc9ycpgDd<4Cuhtp!<1a@|xHsa^kn)ZP-u=hVss zjFMYi>>x{GFVn=DE&KBFg{VS_V+$h*E3&+rk~7;RTqXP$hS-+NvW5&PBmDj8lFNTD zAXT1fUnH_6rv|4HXM&QH9pq!+*x)!qN5;qCGcZ8tEH?P2IwlE}I3xa%z=>L&7z_+f zA`HZ>1Q^7!5)5j8OA;i=Bdjt>?M2^GZWY)ns{Ek%M)6Rwr=9gJrpJ-2jTZ*7ll$r( zlV$*zx_+9&W@RS7=xbot!nUq&hrKQrp;O7^rO;!ep7&_sdako7U-b_mR+cNr`S>4v82MX t%|}WW31hh>kO@W8eVoc(i(PYvztLBbrnH$l3Z9ZO!{sK>+9W?*| literal 0 HcmV?d00001 diff --git a/tests/compiled/chain_assignment.3.7.pyc b/tests/compiled/chain_assignment.3.7.pyc new file mode 100644 index 0000000000000000000000000000000000000000..06ffffd86f6b3e9bca3b1eb814a977cc87f22532 GIT binary patch literal 1643 zcmbtUO>Y}F5an=J@dNqQiEo$^)pzWm~1Q2VHvK7m6#9cR$ z5uK94)-`{k=wITsC!c%Esc*QED%;JeT+HkY$>GdA&hXp$`I^8O{P975ltTQ)%4Ad^ zJO|~ENC+z|y(1&;idZNaivA1lki&5HXeq?T$8lwKv z)PE6;^fCfgZP0W8Q1KloUj!>)m?6DPW|0vvpuI+{+${N>`6#B2Y$)oWej#BmgktvSHUbU z$KV=2!5*0i@x;EK7bni@tSH=en@`}uf6goIL&(Hid(+wCwOFMj&Wlw0(J=n{r@1={BOk8;9`PUi_r>*O|Xhx zonyHQOCxqmH)pSb6%xCLZGVHrEJn>)jNA&J#CbXf_49CWoa>+o9IQ{_K>vpt;5x{^ z1}nr^6sUkPYGF)cS>Qz&qsv%67w2>?J5eaEWF$ha(cbhhMxX1IBg>zFU5Jsmpb=fj zkz}Yw5hxy2O1ULjo&ZLk@ z;Y;akvs^aT+0NEZ@1&je@@HHB+B70hKPvh2HmKPk_W);v({K6=?ml%=IPW3*OTYxY zrTE+?>Co}J6{4Lf nzG(;fWT))8XIvgRKMhaiV#FDyGKo{#!1ooc$%V>7v?~7r5^5__ literal 0 HcmV?d00001 diff --git a/tests/input/chain_assignment.py b/tests/input/chain_assignment.py new file mode 100644 index 0000000..9b31b72 --- /dev/null +++ b/tests/input/chain_assignment.py @@ -0,0 +1,39 @@ +a = [y, z] = x = (k1, k2, k3) = [] = c = myfunc(x) + 3 + +x = y = g = {keyA: X} + +global store_global +Gx = Gy = Gz = Gq1 +Gx = [Gy, Gz] = Gq2 +a = b = store_global = c + +def func_with_global(): + global Gx, Gy, Gz, Gq + Gx = Gy = Gz = Gq + +y = store_subscr[0] = x +a[0] = b[x] = c[3] = D[4] +a[0] = (b[x], c[3]) = D[4] +a[0] = Q = [b[x], c[3]] = F = D[4] +q = v = arr[a:b:c] = x + +class store_attr1: + def __init__(self, a,b,c): + self.a = self.b = self.c = x + self.d = y + +class store_attr2: + def __init__(self, a,b,c): self.a = (self.b, self.c) = x + +a.b = c.d = e.f + g.h + +def store_deref(): + a = I + a = b = c = R1 + a = (b, c) = R2 + def store_fast(): + x = a + y = b + z = c + p = q = r = s + p = [q, r] = s diff --git a/tests/tokenized/chain_assignment.txt b/tests/tokenized/chain_assignment.txt new file mode 100644 index 0000000..f116207 --- /dev/null +++ b/tests/tokenized/chain_assignment.txt @@ -0,0 +1,43 @@ +a = ( y , z ) = x = ( k1 , k2 , k3 ) = ( ) = c = myfunc ( x ) + 3 +x = y = g = { keyA : X } +Gx = Gy = Gz = Gq1 +Gx = ( Gy , Gz ) = Gq2 +a = b = store_global = c +def func_with_global ( ) : + +global Gx , Gy , Gz +Gx = Gy = Gz = Gq + +y = store_subscr [ 0 ] = x +a [ 0 ] = b [ x ] = c [ 3 ] = D [ 4 ] +a [ 0 ] = ( b [ x ] , c [ 3 ] ) = D [ 4 ] +a [ 0 ] = Q = ( b [ x ] , c [ 3 ] ) = F = D [ 4 ] +q = v = arr [ a : b : c ] = x +class store_attr1 : + +def __init__ ( self , a , b , c ) : + +self . a = self . b = self . c = x +self . d = y + + +class store_attr2 : + +def __init__ ( self , a , b , c ) : + +self . a = ( self . b , self . c ) = x + + +a . b = c . d = e . f + g . h +def store_deref ( ) : + +a = I +a = b = c = R1 +a = ( b , c ) = R2 +def store_fast ( ) : + +x = a +y = b +z = c +p = q = r = s +p = ( q , r ) = s diff --git a/tests/tokenized/f-string.txt b/tests/tokenized/f-string.txt index 07ac59b..66c1428 100644 --- a/tests/tokenized/f-string.txt +++ b/tests/tokenized/f-string.txt @@ -1,8 +1,6 @@ var1 = 'x' var2 = 'y' -x = 1.23456 -s1 = 1.23456 -var3 = 1.23456 +x = s1 = var3 = 1.23456 a = 15 some_dict = { } some_dict [ 2 ] = 3