assignment target sequences flattened when possible

Grammar/python.gram

@@ -1162,7 +1162,7 @@ kwarg_or_double_starred[KeywordOrStarred*]:
 
 # NOTE: star_targets may contain *bitwise_or, targets may not.
 star_targets[expr_ty]:
-    | a=star_target !',' { a }
+    | !'*' a=star_target !',' { a }
     | a=star_target b=(',' c=star_target { c })* [','] {
         _PyAST_Tuple(CHECK(asdl_expr_seq*, _PyPegen_seq_insert_in_front(p, a, b)), Store, EXTRA) }
 

Python/bltinmodule.c

@@ -2587,9 +2587,11 @@ sum as builtin_sum
     /
     start: object(c_default="NULL") = 0
 
-Return the sum of a 'start' value (default: 0) plus an iterable of numbers
+Return the sum of a 'start' value (default: 0) plus an iterable of numbers,
+or when given a number as "iterable", the sum of the first number integers
+from start.
 
-When the iterable is empty, return the start value.
+When the iterable is empty or 0, return the start value.
 This function is intended specifically for use with numeric values and may
 reject non-numeric types.
 [clinic start generated code]*/
@@ -2601,12 +2603,66 @@ builtin_sum_impl(PyObject *module, PyObject *iterable, PyObject *start)
     PyObject *result = start;
     PyObject *temp, *item, *iter;
 
+    if (PyLong_CheckExact(iterable) && (!start || PyLong_CheckExact(start))) {
+        /* Range sum path. */
+        size_t num_n_bits = _PyLong_NumBits(iterable);
+        if (num_n_bits == 0) {
+            return _PyLong_GetZero();
+        }
+
+        if (_PyLong_DigitCount((PyLongObject *)iterable) > 1 ||
+            start && _PyLong_DigitCount((PyLongObject *)start) > 1)
+        {
+            goto object_compute;
+        }
+
+        size_t num_start_bits = start ? _PyLong_NumBits(start) : 0;
+        size_t num_max_bits = Py_MAX(num_n_bits, num_start_bits*2);
+
+        if (num_max_bits + num_n_bits + 1 > 8*SIZEOF_LONG_LONG) {
+            goto object_compute;
+        }
+
+        unsigned long long n = PyLong_AsUnsignedLongLong(iterable);
+        if (n == -1 && PyErr_Occurred()) {
+            return NULL;
+        }
+
+        unsigned long long res = 0;
+        if (start) {
+            res = PyLong_AsUnsignedLongLong(start);
+            if (res == -1 && PyErr_Occurred()) {
+                return NULL;
+            }
+        }
+
+        return PyLong_FromUnsignedLongLong(n * (res + res+n - 1) / 2);
+
+    object_compute:
+        PyObject *temp;
+
+#define DO_OPERATION(OP, RHS) \
+        temp = _PyLong_ ## OP ((PyLongObject *)start, (PyLongObject *)(RHS)); \
+        if (temp == NULL) { \
+            return NULL; \
+        } \
+        start = temp; \
+
+        DO_OPERATION(Add, start);
+        DO_OPERATION(Add, iterable);
+        DO_OPERATION(Subtract, _PyLong_GetOne());
+        DO_OPERATION(Multiply, iterable);
+
+        return _PyLong_Rshift(start, 2);
+#undef DO_OPERATION
+    }
+
     iter = PyObject_GetIter(iterable);
     if (iter == NULL)
         return NULL;
 
     if (result == NULL) {
-        result = PyLong_FromLong(0);
+        result = _PyLong_GetZero();
         if (result == NULL) {
             Py_DECREF(iter);
             return NULL;

Python/compile.c

@@ -4545,13 +4545,101 @@ starunpack_helper(struct compiler *c, location loc,
 }
 
 static int
-unpack_helper(struct compiler *c, location loc, asdl_expr_seq *elts)
+flatten_seq_len(struct compiler *c, location loc, asdl_expr_seq *elts,
+                Py_ssize_t *n)
 {
-    Py_ssize_t n = asdl_seq_LEN(elts);
+    /* Compute the overall length of the flattened sequence and
+       also check for duplicate starred expressions.
+       Returns -1 on error, 0 if the sequence is already flattened,
+       1 if flattening was done.
+    */
+    assert(n != NULL);
+    *n = asdl_seq_LEN(elts);
+
+    int has_change = 0;
     int seen_star = 0;
-    for (Py_ssize_t i = 0; i < n; i++) {
+    for (Py_ssize_t i = 0; i < asdl_seq_LEN(elts); i++) {
         expr_ty elt = asdl_seq_GET(elts, i);
-        if (elt->kind == Starred_kind && !seen_star) {
+        if (elt->kind == Starred_kind) {
+            expr_ty val = elt->v.Starred.value;
+            if (val->kind == Tuple_kind) {
+                has_change = 1;
+                Py_ssize_t m;
+                if (flatten_seq_len(c, loc, val->v.Tuple.elts, &m) < 0) {
+                    return -1;
+                }
+                *n += m - 1;
+            }
+            else if (val->kind == List_kind) {
+                has_change = 1;
+                Py_ssize_t m;
+                if (flatten_seq_len(c, loc, val->v.List.elts, &m) < 0) {
+                    return -1;
+                }
+                *n += m - 1;
+            }
+            else if (!seen_star) {
+                seen_star = 1;
+            }
+            else {
+                compiler_error(c, loc,
+                    "multiple starred expressions in assignment");
+                return -1;
+            }
+        }
+    }
+
+    return has_change;
+}
+
+static void
+flatextend_seq(asdl_expr_seq *old, asdl_expr_seq *new, Py_ssize_t *idx)
+{
+    /* Flatten `old` and put the new elements into `new`. */
+    assert(idx != NULL);
+
+    for (Py_ssize_t i = 0; i < asdl_seq_LEN(old); i++) {
+        expr_ty elt = asdl_seq_GET(old, i);
+        if (elt->kind == Starred_kind) {
+            expr_ty val = elt->v.Starred.value;
+            if (val->kind == Tuple_kind) {
+                flatextend_seq(val->v.Tuple.elts, new, idx);
+                continue;
+            }
+            else if (val->kind == List_kind) {
+                flatextend_seq(val->v.List.elts, new, idx);
+                continue;
+            }
+        }
+        asdl_seq_SET(new, (*idx)++, elt);
+    }
+}
+
+static int
+unpack_helper(struct compiler *c, location loc, asdl_expr_seq **elts)
+{
+    Py_ssize_t n;
+    int changed = flatten_seq_len(c, loc, *elts, &n);
+    if (changed < 0) {
+        return ERROR;
+    }
+    if (changed) {
+        asdl_expr_seq *temp = _Py_asdl_expr_seq_new(n, c->c_arena);
+        if (!temp) {
+            return compiler_error(c, loc,
+                "failed to flatten starred sequence in assignment");
+        }
+
+        Py_ssize_t i = 0;
+        flatextend_seq(*elts, temp, &i);
+        assert(i == n);
+        *elts = temp;
+    }
+
+    int seen_star = 0;
+    for (Py_ssize_t i = 0; i < n; i++) {
+        expr_ty elt = asdl_seq_GET(*elts, i);
+        if (elt->kind == Starred_kind) {
             if ((i >= (1 << 8)) ||
                 (n-i-1 >= (INT_MAX >> 8))) {
                 return compiler_error(c, loc,
@@ -4561,10 +4649,6 @@ unpack_helper(struct compiler *c, location loc, asdl_expr_seq *elts)
             ADDOP_I(c, loc, UNPACK_EX, (i + ((n-i-1) << 8)));
             seen_star = 1;
         }
-        else if (elt->kind == Starred_kind) {
-            return compiler_error(c, loc,
-                "multiple starred expressions in assignment");
-        }
     }
     if (!seen_star) {
         ADDOP_I(c, loc, UNPACK_SEQUENCE, n);
@@ -4575,8 +4659,8 @@ unpack_helper(struct compiler *c, location loc, asdl_expr_seq *elts)
 static int
 assignment_helper(struct compiler *c, location loc, asdl_expr_seq *elts)
 {
+    RETURN_IF_ERROR(unpack_helper(c, loc, &elts));
     Py_ssize_t n = asdl_seq_LEN(elts);
-    RETURN_IF_ERROR(unpack_helper(c, loc, elts));
     for (Py_ssize_t i = 0; i < n; i++) {
         expr_ty elt = asdl_seq_GET(elts, i);
         VISIT(c, expr, elt->kind != Starred_kind ? elt : elt->v.Starred.value);