perf(python): xs[a:b] = xs[c:d:-1] reverses in place, and slice assignment copies its source only when it aliases the target

The frontend lowers the reversed self-slice to one call that normalises
both ranges and swaps in place when they name the same elements, taking
the slice first otherwise. assign_slice compares the two storages and
copies only the list assigned from itself.

git-bug: 8db796da28a7bb96bca3865e11f325e52d9b961fc263a44c1852a6d3defdc55f

Author: darmie

crates/zyntax_builtins/src/lists.rs

@@ -552,17 +552,129 @@ fn kind_declarations(k: &KindOps) -> Vec<Decl> {
             ret(out.e()),
         ],
     ));
-    // Replace selected elements in the same header, including when the source aliases it.
-    let replacement = local("replacement", k.list.clone());
+    // Replace selected elements in the same header. The source is read
+    // as it is, unless it is the target's own storage, which is copied
+    // first so the elements moved are the ones the slice named.
+    let replacement = borrowed("replacement", k.list.clone());
     let selected = local("selected", i64());
     let width = local("width", i64());
     let common = local("common", i64());
+    let data = |xs: &Local| fld(xs.e(), "data", i64());
+    let assign_from = |source: Expr, step: Expr| {
+        expr(call(
+            &name("assign_slice_from"),
+            vec![source, xs.e(), start.e(), stop.e(), step, mask.e()],
+            unit(),
+        ))
+    };
     d.push(define(
         &name("assign_slice"),
         &[&ys, &xs, &start, &stop, &step, &mask],
         unit(),
         vec![
-            replacement.decl(call(&name("copy"), vec![ys.e()], k.list.clone())),
+            when(
+                eq(data(&ys), data(&xs)),
+                vec![
+                    replacement.decl(call(&name("copy"), vec![ys.e()], k.list.clone())),
+                    assign_from(replacement.e(), step.e()),
+                    ret_void(),
+                ],
+            ),
+            assign_from(ys.e(), step.e()),
+            ret_void(),
+        ],
+    ));
+    // xs[start:stop] = xs[rstart:rstop:-1]. The two ranges naming the
+    // same elements is a reversal in place; anything else is the slice
+    // taken first and assigned as any other source.
+    let rstart = local("rstart", i64());
+    let rstop = local("rstop", i64());
+    let rmask = local("rmask", i64());
+    let rlo = local("rlo", i64());
+    let rhi = local("rhi", i64());
+    let taken = local("taken", k.list.clone());
+    d.push(define(
+        &name("assign_reversed_slice"),
+        &[&xs, &start, &stop, &mask, &rstart, &rstop, &rmask],
+        unit(),
+        vec![
+            n.decl(len(xs.e())),
+            lo.decl(int(0)),
+            hi.decl(n.e()),
+            when(
+                ne(bitand(mask.e(), int(1)), int(0)),
+                vec![lo.set(call(
+                    "zb_slice_bound",
+                    vec![start.e(), n.e(), int(1)],
+                    i64(),
+                ))],
+            ),
+            when(
+                ne(bitand(mask.e(), int(2)), int(0)),
+                vec![hi.set(call("zb_slice_bound", vec![stop.e(), n.e(), int(1)], i64()))],
+            ),
+            when(lt(hi.e(), lo.e()), vec![hi.set(lo.e())]),
+            rlo.decl(sub(n.e(), int(1))),
+            rhi.decl(int(-1)),
+            when(
+                ne(bitand(rmask.e(), int(1)), int(0)),
+                vec![rlo.set(call(
+                    "zb_slice_bound",
+                    vec![rstart.e(), n.e(), int(-1)],
+                    i64(),
+                ))],
+            ),
+            when(
+                ne(bitand(rmask.e(), int(2)), int(0)),
+                vec![rhi.set(call(
+                    "zb_slice_bound",
+                    vec![rstop.e(), n.e(), int(-1)],
+                    i64(),
+                ))],
+            ),
+            if_(
+                and(
+                    eq(rlo.e(), sub(hi.e(), int(1))),
+                    eq(rhi.e(), sub(lo.e(), int(1))),
+                ),
+                vec![
+                    i.decl(lo.e()),
+                    j.decl(sub(hi.e(), int(1))),
+                    while_(
+                        lt(i.e(), j.e()),
+                        vec![
+                            a.decl(el(&xs, i.e())),
+                            b.decl(el(&xs, j.e())),
+                            set_idx(xs.e(), i.e(), b.e()),
+                            set_idx(xs.e(), j.e(), a.e()),
+                            i.add_assign(int(1)),
+                            j.set(sub(j.e(), int(1))),
+                        ],
+                    ),
+                ],
+                vec![
+                    taken.decl(call(
+                        &name("slice"),
+                        vec![
+                            xs.e(),
+                            rstart.e(),
+                            rstop.e(),
+                            int(-1),
+                            bitor(rmask.e(), int(4)),
+                        ],
+                        k.list.clone(),
+                    )),
+                    assign_from(taken.e(), int(0)),
+                ],
+            ),
+            ret_void(),
+        ],
+    ));
+    d.push(define(
+        &name("assign_slice_from"),
+        &[&replacement, &xs, &start, &stop, &step, &mask],
+        unit(),
+        vec![
             n.decl(len(xs.e())),
             st.decl(int(1)),
             when(ne(bitand(mask.e(), int(4)), int(0)), vec![st.set(step.e())]),

crates/zyntax_python/conformance/list/reversed_slice_assignment.expected

@@ -0,0 +1,16 @@
+prefix [2, 1, 0, 3, 4, 5]
+k [4, 3, 0, 1, 2, 5]
+middle [4, 3, 2, 1, 0, 5]
+whole [5, 0, 1, 2, 3, 4]
+same length elsewhere [4, 3, 1, 2, 3, 4]
+longer [4, 3, 2, 1, 1, 2, 3, 4]
+shorter [4, 1, 2, 3, 4]
+inverted target [4, 1, 2, 4, 3, 3, 4]
+negative [4, 1, 2, 4, 4, 3, 3]
+out of range [3, 3, 4, 4, 2, 1, 4]
+floats [2.5, 1.5, 3.5]
+strings ['a', 'd', 'c', 'b']
+mixed [None, 2.5, 'x', 1]
+empty [4, 1, 2, 4, 4, 3, 3, 3, 3, 4, 4, 2, 1, 4]
+itself [0, 0, 1, 2]
+own slice [0, 1, 0, 1, 2]

crates/zyntax_python/conformance/list/reversed_slice_assignment.py

@@ -0,0 +1,46 @@
+# Every shape of xs[a:b] = xs[c:d:-1]: coinciding ranges, ranges of
+# the same length elsewhere, different lengths, empty, out of range,
+# negative bounds, and a second name for the same list.
+
+def show(label, xs):
+    print(label, xs)
+
+xs = [0, 1, 2, 3, 4, 5]
+xs[:3] = xs[2::-1]
+show("prefix", xs)
+k = 4
+xs[:k + 1] = xs[k::-1]
+show("k", xs)
+xs[2:5] = xs[4:1:-1]
+show("middle", xs)
+xs[:] = xs[::-1]
+show("whole", xs)
+xs[0:2] = xs[5:3:-1]
+show("same length elsewhere", xs)
+xs[1:2] = xs[4:1:-1]
+show("longer", xs)
+xs[1:4] = xs[0:0:-1]
+show("shorter", xs)
+xs[3:1] = xs[5:2:-1]
+show("inverted target", xs)
+xs[-3:] = xs[-1:-4:-1]
+show("negative", xs)
+xs[:100] = xs[100::-1]
+show("out of range", xs)
+ys = [1.5, 2.5, 3.5]
+ys[:2] = ys[1::-1]
+show("floats", ys)
+ws = ["a", "b", "c", "d"]
+ws[1:] = ws[:0:-1]
+show("strings", ws)
+zs = [1, "x", 2.5, None]
+zs[:] = zs[::-1]
+show("mixed", zs)
+n = 0
+xs[:n] = xs[n - 1::-1]
+show("empty", xs)
+xs = [0, 1, 2]
+xs[1:3] = xs
+show("itself", xs)
+xs[0:1] = xs[1:3]
+show("own slice", xs)

crates/zyntax_python/src/lower.rs

@@ -2009,6 +2009,11 @@ impl<'m> Lowerer<'m> {
                 push(out, TypedStatement::Expression(Box::new(v.node)));
             }
             py::Stmt::Assign(a) => {
+                if a.targets.len() == 1
+                    && self.reversed_slice_assign(&a.targets[0], &a.value, span, out)?
+                {
+                    return Ok(());
+                }
                 // `a = b = v` evaluates `v` once and binds each target
                 // to it, left to right.
                 let value = self.expr(&a.value)?;
@@ -2273,6 +2278,91 @@ impl<'m> Lowerer<'m> {
         Ok(())
     }
 
+    /// `xs[a:b] = xs[c:d:-1]` over one list name, with bounds that
+    /// cannot rebind the name: lowered as one call that reverses the
+    /// elements in place when the two ranges name the same ones, so the
+    /// flip of a permutation copies nothing. Anything else is left to
+    /// the general slice assignment; says whether it was taken.
+    fn reversed_slice_assign(
+        &mut self,
+        target: &py::Expr,
+        value: &py::Expr,
+        span: Span,
+        out: &mut Vec<Stmt>,
+    ) -> Result<bool> {
+        let (py::Expr::Subscript(t), py::Expr::Subscript(v)) = (target, value) else {
+            return Ok(false);
+        };
+        let (py::Expr::Name(tn), py::Expr::Name(vn)) = (&*t.value, &*v.value) else {
+            return Ok(false);
+        };
+        let (py::Expr::Slice(ts), py::Expr::Slice(vs)) = (&*t.slice, &*v.slice) else {
+            return Ok(false);
+        };
+        let minus_one = |e: &Option<Box<py::Expr>>| {
+            matches!(e.as_deref(), Some(py::Expr::UnaryOp(u))
+                if u.op == py::UnaryOp::USub
+                    && matches!(&*u.operand, py::Expr::NumberLiteral(n)
+                        if matches!(&n.value, py::Number::Int(i) if i.as_u64() == Some(1))))
+        };
+        fn pure_bound(e: &Option<Box<py::Expr>>) -> bool {
+            fn pure(e: &py::Expr) -> bool {
+                match e {
+                    py::Expr::Name(_) | py::Expr::NumberLiteral(_) => true,
+                    py::Expr::BinOp(b) => pure(&b.left) && pure(&b.right),
+                    py::Expr::UnaryOp(u) => pure(&u.operand),
+                    _ => false,
+                }
+            }
+            e.as_deref().is_none_or(pure)
+        }
+        if tn.id != vn.id
+            || ts.step.is_some()
+            || !minus_one(&vs.step)
+            || !pure_bound(&ts.lower)
+            || !pure_bound(&ts.upper)
+            || !pure_bound(&vs.lower)
+            || !pure_bound(&vs.upper)
+        {
+            return Ok(false);
+        }
+        let seq = self.expr(&t.value)?;
+        let Ty::List(e) = seq.ty else {
+            return Ok(false);
+        };
+        let bounds = |this: &mut Self,
+                      lower: &Option<Box<py::Expr>>,
+                      upper: &Option<Box<py::Expr>>|
+         -> Result<(Node, Node, Node)> {
+            let mut mask = 0;
+            let mut bound = |this: &mut Self, e: &Option<Box<py::Expr>>, bit: i64| match e {
+                Some(e) => {
+                    mask |= bit;
+                    this.expr_as(e, Ty::Int)
+                }
+                None => Ok(int_lit(0, span)),
+            };
+            let lo = bound(this, lower, 1)?;
+            let hi = bound(this, upper, 2)?;
+            Ok((lo, hi, int_lit(mask, span)))
+        };
+        let (start, stop, mask) = bounds(self, &ts.lower, &ts.upper)?;
+        let (rstart, rstop, rmask) = bounds(self, &vs.lower, &vs.upper)?;
+        let call = call(
+            &list_fn("assign_reversed_slice", e),
+            vec![seq.node, start, stop, mask, rstart, rstop, rmask],
+            Ty::None,
+            span,
+        );
+        out.push(TypedNode::new(
+            TypedStatement::Expression(Box::new(call)),
+            Type::Unknown,
+            span,
+        ));
+        out.push(self.pending_check(span));
+        Ok(true)
+    }
+
     /// `target = value`: a `let` the first time a name is seen in this
     /// function, an assignment afterwards; the value converted to the
     /// type inference gave the name.