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 0000000..e8671b9 Binary files /dev/null and b/tests/compiled/chain_assignment.2.7.pyc differ diff --git a/tests/compiled/chain_assignment.3.7.pyc b/tests/compiled/chain_assignment.3.7.pyc new file mode 100644 index 0000000..06ffffd Binary files /dev/null and b/tests/compiled/chain_assignment.3.7.pyc differ 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