Preserve tuple elements and unboxed dynamic indices

git-bug: ac72b510da295835d7420e58584ceafbf3b49b8b88ca64d204ba81ed45a6a3d2

Author: darmie

crates/zyntax_builtins/src/dynamic.rs

@@ -1456,6 +1456,83 @@ pub(crate) fn declarations(policy: &Policy, list_type: TypeId) -> Vec<Decl> {
             )),
         ],
     ));
+    // Positional indexing already has an unboxed integer in the Python
+    // frontend (notably while unpacking nested tuples). Keep that integer
+    // through the dispatch instead of allocating a box only to read it
+    // back in every list branch. Dicts still need a boxed key.
+    let position = local("position", i64());
+    d.push(define(
+        "zb_any_getitem_i64",
+        &[&x, &position],
+        any(),
+        vec![
+            cat.decl(category(x.e())),
+            when(
+                is(&cat, STR),
+                vec![ret(box_str(call(
+                    "zb_str_get",
+                    vec![get_str(x.e()), position.e()],
+                    string(),
+                )))],
+            ),
+            when(
+                ne(cat.e(), int(CUSTOM)),
+                vec![type_error(add(
+                    quoted(type_name(x.e())),
+                    text(" object is not subscriptable"),
+                ))],
+            ),
+            when(
+                is_dict(x.e()),
+                vec![ret(call(
+                    "zb_dict_get",
+                    vec![raw_any(x.e()), box_i64(position.e())],
+                    any(),
+                ))],
+            ),
+            when(
+                kind_is(Kind::Int, x.e()),
+                vec![ret(box_i64(call(
+                    "zb_list_get_i64",
+                    vec![unbox(Kind::Int, x.e()), position.e()],
+                    i64(),
+                )))],
+            ),
+            when(
+                kind_is(Kind::Float, x.e()),
+                vec![ret(box_f64(call(
+                    "zb_list_get_f64",
+                    vec![unbox(Kind::Float, x.e()), position.e()],
+                    f64(),
+                )))],
+            ),
+            when(
+                kind_is(Kind::Str, x.e()),
+                vec![ret(box_str(call(
+                    "zb_list_get_str",
+                    vec![unbox(Kind::Str, x.e()), position.e()],
+                    string(),
+                )))],
+            ),
+            when(
+                kind_is(Kind::Ptr, x.e()),
+                vec![ret(call(
+                    "zb_hook_box_instance",
+                    vec![call(
+                        "zb_list_get_ptr",
+                        vec![unbox(Kind::Ptr, x.e()), position.e()],
+                        Kind::Ptr.ty(),
+                    )],
+                    any(),
+                ))],
+            ),
+            ret(call(
+                "zb_list_get_any",
+                vec![iter(x.e()), position.e()],
+                any(),
+            )),
+        ],
+    ));
     let mutable_list = |x: Expr| {
         and(
             eq(category(x.clone()), int(CUSTOM)),
@@ -1567,6 +1644,99 @@ pub(crate) fn declarations(policy: &Policy, list_type: TypeId) -> Vec<Decl> {
             ret_void(),
         ],
     ));
+    let list_set_i64 = |k: Kind, x: Expr, i: Expr, v: Expr| {
+        expr(call(
+            &format!("zb_list_set_{}", k.suffix()),
+            vec![unbox(k, x), i, v],
+            unit(),
+        ))
+    };
+    d.push(define(
+        "zb_any_setitem_i64",
+        &[&x, &position, &v],
+        unit(),
+        vec![
+            when(
+                boxed_dict(x.e()),
+                vec![
+                    expr(call(
+                        "zb_dict_set",
+                        vec![raw_any(x.e()), box_i64(position.e()), v.e()],
+                        unit(),
+                    )),
+                    ret_void(),
+                ],
+            ),
+            when(
+                not(mutable_list(x.e())),
+                vec![type_error(add(
+                    quoted(type_name(x.e())),
+                    text(" object does not support item assignment"),
+                ))],
+            ),
+            when(
+                kind_is(Kind::Int, x.e()),
+                vec![
+                    list_set_i64(
+                        Kind::Int,
+                        x.e(),
+                        position.e(),
+                        call("zb_any_as_i64", vec![v.e()], i64()),
+                    ),
+                    ret_void(),
+                ],
+            ),
+            when(
+                kind_is(Kind::Float, x.e()),
+                vec![
+                    list_set_i64(
+                        Kind::Float,
+                        x.e(),
+                        position.e(),
+                        call("zb_any_as_f64", vec![v.e()], f64()),
+                    ),
+                    ret_void(),
+                ],
+            ),
+            when(
+                kind_is(Kind::Str, x.e()),
+                vec![
+                    list_set_i64(
+                        Kind::Str,
+                        x.e(),
+                        position.e(),
+                        call("zb_any_as_str", vec![v.e()], string()),
+                    ),
+                    ret_void(),
+                ],
+            ),
+            when(
+                kind_is(Kind::Ptr, x.e()),
+                vec![
+                    when(
+                        not(is_instance(v.e())),
+                        vec![type_error(add(
+                            text("expected an instance, got "),
+                            quoted(type_name(v.e())),
+                        ))],
+                    ),
+                    list_set_i64(
+                        Kind::Ptr,
+                        x.e(),
+                        position.e(),
+                        cast(call("zb_unbox_instance_raw", vec![v.e()], i64()), usize()),
+                    ),
+                    ret_void(),
+                ],
+            ),
+            expr(call(
+                "zb_list_set_any",
+                vec![unbox(Kind::Any, x.e()), position.e(), v.e()],
+                unit(),
+            )),
+            ret_void(),
+        ],
+    ));
     d.push(define(
         "zb_any_delitem",
         &[&x, &i],

crates/zyntax_python/conformance/tuple/literal_unpack.expected

@@ -0,0 +1,12 @@
+1 2 [1, 2]
+3 4.5 five
+20
+20.5
+a
+two
+7
+7.5
+x
+zero
+[9.5, 2.5]
+{0: 9.5}

crates/zyntax_python/conformance/tuple/literal_unpack.py

@@ -0,0 +1,47 @@
+def record(events, value):
+    events.append(value)
+    return value
+
+
+events = []
+left, right = (record(events, 1), record(events, 2))
+print(left, right, events)
+
+first, (second, third) = (record(events, 3), (4.5, "five"))
+print(first, second, third)
+
+def get_first(value):
+    return value[0]
+
+def last(value):
+    return value[-1]
+
+def get_second(value):
+    return value[1]
+
+
+def dict_item(value):
+    return value[2]
+
+
+print(last((10, 20)))
+print(get_second([10.5, 20.5]))
+print(get_first("abc"))
+print(dict_item({2: "two"}))
+
+def dynamic_zero(value):
+    return value[0]
+
+
+print(dynamic_zero((7, 8)))
+print(dynamic_zero([7.5, 8.5]))
+print(dynamic_zero("xyz"))
+print(dynamic_zero({0: "zero"}))
+
+def replace_zero(value):
+    value[0] = 9.5
+    return value
+
+
+print(replace_zero([1.5, 2.5]))
+print(replace_zero({0: 1.5}))

crates/zyntax_python/src/lower.rs

@@ -2107,6 +2107,34 @@ impl<'m> Lowerer<'m> {
                 {
                     return Ok(());
                 }
+                // A tuple literal consumed immediately by an unpacking
+                // assignment has no observable tuple identity. Evaluate
+                // every RHS expression first, then bind the targets in
+                // order, retaining each element's inferred type.
+                if let ([target], py::Expr::Tuple(source)) = (a.targets.as_slice(), &*a.value) {
+                    let targets = match target {
+                        py::Expr::Tuple(t) => Some(&t.elts),
+                        py::Expr::List(l) => Some(&l.elts),
+                        _ => None,
+                    };
+                    if let Some(targets) = targets.filter(|t| {
+                        t.len() == source.elts.len()
+                            && t.iter().all(|target| {
+                                !matches!(target, py::Expr::Tuple(_) | py::Expr::List(_))
+                            })
+                    }) {
+                        let mut values = Vec::with_capacity(source.elts.len());
+                        for expr in &source.elts {
+                            let value = self.expr(expr)?;
+                            out.append(&mut self.hoisted);
+                            values.push(self.hold(value, out, span));
+                        }
+                        for (target, value) in targets.iter().zip(values) {
+                            self.bind(target, value, span, out)?;
+                        }
+                        return Ok(());
+                    }
+                }
                 // A literal that says nothing of its elements is built
                 // as the list its name or field holds, which inference
                 // typed by what the program puts in it.
@@ -2695,9 +2723,19 @@ impl<'m> Lowerer<'m> {
                         )
                     }
                     Ty::Object => {
-                        let i = self.expr_as(&sub.slice, Ty::Object)?;
+                        let i = self.expr(&sub.slice)?;
                         let v = self.coerce(value, Ty::Object);
-                        call("zb_any_setitem", vec![seq.node, i, v], Ty::None, span)
+                        if i.ty == Ty::Int {
+                            call(
+                                "zb_any_setitem_i64",
+                                vec![seq.node, i.node, v],
+                                Ty::None,
+                                span,
+                            )
+                        } else {
+                            let i = self.coerce(i, Ty::Object);
+                            call("zb_any_setitem", vec![seq.node, i, v], Ty::None, span)
+                        }
                     }
                     _ => {
                         return unsupported(
@@ -2915,16 +2953,12 @@ impl<'m> Lowerer<'m> {
             Ty::List(e) => elem_call("get", e, vec![seq.node, index], span),
             Ty::Tuple => call("zb_list_get_any", vec![seq.node, index], Ty::Object, span),
             Ty::Str => call("zb_str_get", vec![seq.node, index], Ty::Str, span),
-            _ => {
-                let i = self.coerce(
-                    Val {
-                        node: index,
-                        ty: Ty::Int,
-                    },
-                    Ty::Object,
-                );
-                call("zb_any_getitem", vec![seq.node, i], Ty::Object, span)
-            }
+            _ => call(
+                "zb_any_getitem_i64",
+                vec![seq.node, index],
+                Ty::Object,
+                span,
+            ),
         };
         Val { node, ty: elem_ty }
     }
@@ -4922,10 +4956,16 @@ impl<'m> Lowerer<'m> {
                 })
             }
             _ => {
-                let key = self.expr_as(&sub.slice, Ty::Object)?;
+                let key = self.expr(&sub.slice)?;
                 let o = self.coerce(seq, Ty::Object);
+                let node = if key.ty == Ty::Int {
+                    call("zb_any_getitem_i64", vec![o, key.node], Ty::Object, span)
+                } else {
+                    let key = self.coerce(key, Ty::Object);
+                    call("zb_any_getitem", vec![o, key], Ty::Object, span)
+                };
                 Ok(Val {
-                    node: call("zb_any_getitem", vec![o, key], Ty::Object, span),
+                    node,
                     ty: Ty::Object,
                 })
             }

crates/zyntax_python/src/types.rs

@@ -3358,6 +3358,34 @@ impl Walker<'_> {
         }
     }
 
+    /// A tuple literal gives each unpacked name the type of its own
+    /// expression. No tuple escapes this assignment, so the lowering can
+    /// bind those values directly without putting them through `Any`.
+    fn target_tuple_literal(&mut self, target: &py::Expr, value: &py::Expr) -> bool {
+        let py::Expr::Tuple(source) = value else {
+            return false;
+        };
+        let targets = match target {
+            py::Expr::Tuple(t) => &t.elts,
+            py::Expr::List(l) => &l.elts,
+            _ => return false,
+        };
+        if targets.len() != source.elts.len() {
+            return false;
+        }
+        if targets
+            .iter()
+            .any(|target| matches!(target, py::Expr::Tuple(_) | py::Expr::List(_)))
+        {
+            return false;
+        }
+        for (target, value) in targets.iter().zip(&source.elts) {
+            let ty = self.expr(value);
+            self.target(target, ty);
+        }
+        true
+    }
+
     fn stmts(&mut self, stmts: &[py::Stmt]) {
         for s in stmts {
             self.stmt(s);
@@ -3369,6 +3397,9 @@ impl Walker<'_> {
             py::Stmt::Assign(a) => {
                 let ty = self.expr(&a.value);
                 for t in &a.targets {
+                    if a.targets.len() == 1 && self.target_tuple_literal(t, &a.value) {
+                        continue;
+                    }
                     // An unkinded literal is the list its writes make
                     // it, decided afresh from the types known now.
                     if let py::Expr::Name(n) = t {