optimize 2-arg super also

carljm · Apr 24, 2023 · Apr 10, 2023 · Apr 13, 2023 · Apr 13, 2023 · Apr 13, 2023
commit f229b5bdf08fb21f64b7de91f07eedb1ec879096
diff --git a/Doc/library/dis.rst b/Doc/library/dis.rst
@@ -1036,14 +1036,20 @@ iterations of the loop.
       pushed to the stack before the attribute or unbound method respectively.
 
 
-.. opcode:: LOAD_ZERO_SUPER_ATTR (namei)
-
-   This opcode implements zero-argument :func:`super` (i.e. ``super().method()``
-   and ``super().attr``). It works the same as :opcode:`LOAD_ATTR`, except that
-   instead of expecting a single receiver on the stack, it expects three objects
-   (from top of stack down): ``self`` (the first argument to the current
-   method), ``cls`` (the class within which the current method was defined), and
-   the global ``super``.
+.. opcode:: LOAD_SUPER_ATTR (namei)
+
+   This opcode implements :func:`super` (e.g. ``super().method()`` and
+   ``super().attr``). It works the same as :opcode:`LOAD_ATTR`, except that
+   ``namei`` is shifted left by 2 bits instead of 1, and instead of expecting a
+   single receiver on the stack, it expects three objects (from top of stack
+   down): ``self`` (the first argument to the current method), ``cls`` (the
+   class within which the current method was defined), and the global ``super``.
+
+   The low bit of ``namei`` signals to attempt a method load, as with
+   :opcode:`LOAD_ATTR`.
+
+   The second-low bit of ``namei``, if set, means that this was a zero-argument
+   call to :func:`super`.
 
    .. versionadded:: 3.12
 

diff --git a/Include/internal/pycore_opcode.h b/Include/internal/pycore_opcode.h
diff --git a/Include/opcode.h b/Include/opcode.h
diff --git a/Lib/dis.py b/Lib/dis.py
@@ -41,7 +41,7 @@
 FOR_ITER = opmap['FOR_ITER']
 SEND = opmap['SEND']
 LOAD_ATTR = opmap['LOAD_ATTR']
-LOAD_ZERO_SUPER_ATTR = opmap['LOAD_ZERO_SUPER_ATTR']
+LOAD_SUPER_ATTR = opmap['LOAD_SUPER_ATTR']
 
 CACHE = opmap["CACHE"]
 
@@ -472,10 +472,14 @@ def _get_instructions_bytes(code, varname_from_oparg=None,
                     argval, argrepr = _get_name_info(arg//2, get_name)
                     if (arg & 1) and argrepr:
                         argrepr = "NULL + " + argrepr
-                elif deop == LOAD_ATTR or deop == LOAD_ZERO_SUPER_ATTR:
+                elif deop == LOAD_ATTR:
                     argval, argrepr = _get_name_info(arg//2, get_name)
                     if (arg & 1) and argrepr:
                         argrepr = "NULL|self + " + argrepr
+                elif deop == LOAD_SUPER_ATTR:
+                    argval, argrepr = _get_name_info(arg//4, get_name)
+                    if (arg & 1) and argrepr:
+                        argrepr = "NULL|self + " + argrepr
                 else:
                     argval, argrepr = _get_name_info(arg, get_name)
             elif deop in hasjabs:

diff --git a/Lib/opcode.py b/Lib/opcode.py
@@ -196,7 +196,7 @@ def pseudo_op(name, op, real_ops):
 def_op('DELETE_DEREF', 139)
 hasfree.append(139)
 jrel_op('JUMP_BACKWARD', 140)    # Number of words to skip (backwards)
-name_op('LOAD_ZERO_SUPER_ATTR', 141)
+name_op('LOAD_SUPER_ATTR', 141)
 def_op('CALL_FUNCTION_EX', 142)  # Flags
 
 def_op('EXTENDED_ARG', 144)
@@ -264,7 +264,9 @@ def pseudo_op(name, op, real_ops):
 pseudo_op('JUMP_NO_INTERRUPT', 261, ['JUMP_FORWARD', 'JUMP_BACKWARD_NO_INTERRUPT'])
 
 pseudo_op('LOAD_METHOD', 262, ['LOAD_ATTR'])
-pseudo_op('LOAD_ZERO_SUPER_METHOD', 263, ['LOAD_ZERO_SUPER_ATTR'])
+pseudo_op('LOAD_SUPER_METHOD', 263, ['LOAD_SUPER_ATTR'])
+pseudo_op('LOAD_ZERO_SUPER_METHOD', 264, ['LOAD_SUPER_ATTR'])
+pseudo_op('LOAD_ZERO_SUPER_ATTR', 265, ['LOAD_SUPER_ATTR'])
 
 MAX_PSEUDO_OPCODE = MIN_PSEUDO_OPCODE + len(_pseudo_ops) - 1
 

diff --git a/Lib/test/test_super.py b/Lib/test/test_super.py
@@ -362,6 +362,21 @@ def method(self):
         with patch("test.test_super.super", MySuper) as m:
             self.assertEqual(C().method(), "super super")
 
+    def test_shadowed_dynamic_two_arg(self):
+        call_args = []
+        class MySuper:
+            def __init__(self, *args):
+                call_args.append(args)
+            msg = "super super"
+
+        class C:
+            def method(self):
+                return super(1, 2).msg
+
+        with patch("test.test_super.super", MySuper) as m:
+            self.assertEqual(C().method(), "super super")
+            self.assertEqual(call_args, [(1, 2)])
+
     def test_attribute_error(self):
         class C:
             def method(self):

@@ -1551,8 +1551,8 @@ dummy_func(
             PREDICT(JUMP_BACKWARD);
         }
 
-        inst(LOAD_ZERO_SUPER_ATTR, (global_super, class, self -- res2 if (oparg & 1), res)) {
-            PyObject *name = GETITEM(frame->f_code->co_names, oparg >> 1);
+        inst(LOAD_SUPER_ATTR, (global_super, class, self -- res2 if (oparg & 1), res)) {
+            PyObject *name = GETITEM(frame->f_code->co_names, oparg >> 2);
             if (global_super == (PyObject *)&PySuper_Type) {
                 int meth_found = 0;
                 Py_DECREF(global_super);
@@ -1571,7 +1571,13 @@ dummy_func(
                     Py_DECREF(self);
                 }
             } else {
-                PyObject *super = PyObject_Vectorcall(global_super, NULL, 0, NULL);
+                PyObject *super;
+                if (oparg & 2) {
+                    super = PyObject_Vectorcall(global_super, NULL, 0, NULL);
+                } else {
+                    PyObject *stack[] = {class, self};
+                    super = PyObject_Vectorcall(global_super, stack, 2, NULL);
+                }
                 DECREF_INPUTS();
                 ERROR_IF(super == NULL, error);
                 res = PyObject_GetAttr(super, name);

@@ -829,7 +829,9 @@ stack_effect(int opcode, int oparg, int jump)
 
         case LOAD_METHOD:
             return 1;
+        case LOAD_SUPER_METHOD:
         case LOAD_ZERO_SUPER_METHOD:
+        case LOAD_ZERO_SUPER_ATTR:
             return -1;
         default:
             return PY_INVALID_STACK_EFFECT;
@@ -1041,19 +1043,32 @@ compiler_addop_name(struct compiler_unit *u, location loc,
     if (arg < 0) {
         return ERROR;
     }
-    if (opcode == LOAD_ATTR || opcode == LOAD_ZERO_SUPER_ATTR) {
+    if (opcode == LOAD_ATTR) {
         arg <<= 1;
     }
     if (opcode == LOAD_METHOD) {
         opcode = LOAD_ATTR;
         arg <<= 1;
         arg |= 1;
     }
-    if (opcode == LOAD_ZERO_SUPER_METHOD) {
-        opcode = LOAD_ZERO_SUPER_ATTR;
-        arg <<= 1;
+    if (opcode == LOAD_SUPER_ATTR) {
+        arg <<= 2;
+    }
+    if (opcode == LOAD_SUPER_METHOD) {
+        opcode = LOAD_SUPER_ATTR;
+        arg <<= 2;
         arg |= 1;
     }
+    if (opcode == LOAD_ZERO_SUPER_ATTR) {
+        opcode = LOAD_SUPER_ATTR;
+        arg <<= 2;
+        arg |= 2;
+    }
+    if (opcode == LOAD_ZERO_SUPER_METHOD) {
+        opcode = LOAD_SUPER_ATTR;
+        arg <<= 2;
+        arg |= 3;
+    }
     return codegen_addop_i(&u->u_instr_sequence, opcode, arg, loc);
 }
 
@@ -4222,15 +4237,15 @@ is_import_originated(struct compiler *c, expr_ty e)
 }
 
 static int
-is_zero_arg_super_call(struct compiler *c, expr_ty e)
+can_optimize_super_call(struct compiler *c, expr_ty e)
 {
     if (e->kind != Call_kind ||
         e->v.Call.func->kind != Name_kind ||
         !_PyUnicode_EqualToASCIIString(e->v.Call.func->v.Name.id, "super") ||
-        asdl_seq_LEN(e->v.Call.args) != 0 ||
         asdl_seq_LEN(e->v.Call.keywords) != 0) {
         return 0;
     }
+    Py_ssize_t num_args = asdl_seq_LEN(e->v.Call.args);
 
     PyObject *super_name = e->v.Call.func->v.Name.id;
     // try to detect statically-visible shadowing of 'super' name
@@ -4242,6 +4257,24 @@ is_zero_arg_super_call(struct compiler *c, expr_ty e)
     if (scope != 0) {
         return 0;
     }
+
+    if (num_args == 2) {
+        for (Py_ssize_t i = 0; i < num_args; i++) {
+            expr_ty elt = asdl_seq_GET(e->v.Call.args, i);
+            if (elt->kind == Starred_kind) {
+                return 0;
+            }
+        }
+        // exactly two non-starred args; we can just load
+        // the provided args
+        return 1;
+    }
+
+    if (num_args != 0) {
+        return 0;
+    }
+    // we need the following for zero-arg super():
+
     // enclosing function should have at least one argument
     if (c->u->u_metadata.u_argcount == 0 &&
         c->u->u_metadata.u_posonlyargcount == 0) {
@@ -4255,13 +4288,19 @@ is_zero_arg_super_call(struct compiler *c, expr_ty e)
 }
 
 static int
-load_args_for_zero_super(struct compiler *c, expr_ty e) {
+load_args_for_super(struct compiler *c, expr_ty e) {
     location loc = LOC(e);
 
     // load super() global
-    PyObject *super_name = e->v.Attribute.value->v.Call.func->v.Name.id;
+    PyObject *super_name = e->v.Call.func->v.Name.id;
     RETURN_IF_ERROR(compiler_nameop(c, loc, super_name, Load));
 
+    if (asdl_seq_LEN(e->v.Call.args) == 2) {
+        VISIT(c, expr, asdl_seq_GET(e->v.Call.args, 0));
+        VISIT(c, expr, asdl_seq_GET(e->v.Call.args, 1));
+        return SUCCESS;
+    }
+
     // load __class__ cell
     PyObject *name = &_Py_ID(__class__);
     assert(get_ref_type(c, name) == FREE);
@@ -4349,9 +4388,11 @@ maybe_optimize_method_call(struct compiler *c, expr_ty e)
     /* Alright, we can optimize the code. */
     location loc = LOC(meth);
 
-    if (is_zero_arg_super_call(c, meth->v.Attribute.value)) {
-        RETURN_IF_ERROR(load_args_for_zero_super(c, meth));
-        ADDOP_NAME(c, loc, LOAD_ZERO_SUPER_METHOD, meth->v.Attribute.attr, names);
+    if (can_optimize_super_call(c, meth->v.Attribute.value)) {
+        RETURN_IF_ERROR(load_args_for_super(c, meth->v.Attribute.value));
+        int opcode = asdl_seq_LEN(meth->v.Attribute.value->v.Call.args) ?
+            LOAD_SUPER_METHOD : LOAD_ZERO_SUPER_METHOD;
+        ADDOP_NAME(c, loc, opcode, meth->v.Attribute.attr, names);
     } else {
         VISIT(c, expr, meth->v.Attribute.value);
         loc = update_start_location_to_match_attr(c, loc, meth);
@@ -5365,9 +5406,11 @@ compiler_visit_expr1(struct compiler *c, expr_ty e)
         return compiler_formatted_value(c, e);
     /* The following exprs can be assignment targets. */
     case Attribute_kind:
-        if (e->v.Attribute.ctx == Load && is_zero_arg_super_call(c, e->v.Attribute.value)) {
-            RETURN_IF_ERROR(load_args_for_zero_super(c, e));
-            ADDOP_NAME(c, loc, LOAD_ZERO_SUPER_ATTR, e->v.Attribute.attr, names);
+        if (e->v.Attribute.ctx == Load && can_optimize_super_call(c, e->v.Attribute.value)) {
+            RETURN_IF_ERROR(load_args_for_super(c, e->v.Attribute.value));
+            int opcode = asdl_seq_LEN(e->v.Attribute.value->v.Call.args) ?
+                LOAD_SUPER_ATTR : LOAD_ZERO_SUPER_ATTR;
+            ADDOP_NAME(c, loc, opcode, e->v.Attribute.attr, names);
             return SUCCESS;
         }
         VISIT(c, expr, e->v.Attribute.value);