[mypyc] Borrow the callable environment link (#21957)

It's now final, so it can be borrowed more liberally.

This made a nested function microbenchmark ~1.3x faster on 3.14
(free-threaded).

I used coding agent assist.

Author: JukkaL

mypyc/irbuild/env_class.py

@@ -143,7 +143,11 @@ def load_env_registers(builder: IRBuilder, prefix: str = "") -> None:
 
 
 def load_outer_env(
-    builder: IRBuilder, base: Value, outer_env: dict[SymbolNode, SymbolTarget]
+    builder: IRBuilder,
+    base: Value,
+    outer_env: dict[SymbolNode, SymbolTarget],
+    *,
+    borrow: bool = False,
 ) -> Value:
     """Load the environment class for a given base into a register.
 
@@ -156,7 +160,10 @@ def load_outer_env(
 
     Returns the register where the environment class was loaded.
     """
-    env = builder.add(GetAttr(base, ENV_ATTR_NAME, builder.fn_info.fitem.line))
+    if borrow:
+        assert isinstance(base.type, RInstance)
+        assert base.type.class_ir.is_final_attr(ENV_ATTR_NAME)
+    env = builder.add(GetAttr(base, ENV_ATTR_NAME, builder.fn_info.fitem.line, borrow=borrow))
     assert isinstance(env.type, RInstance), f"{env} must be of type RInstance"
 
     for symbol, target in outer_env.items():
@@ -182,7 +189,9 @@ def load_outer_envs(builder: IRBuilder, base: ImplicitClass) -> None:
         if isinstance(base, GeneratorClass):
             base.prev_env_reg = load_outer_env(builder, base.curr_env_reg, outer_env)
         else:
-            base.prev_env_reg = load_outer_env(builder, base.self_reg, outer_env)
+            # The callable stays alive throughout __call__, and its environment link is Final,
+            # so the environment can be borrowed for the duration of the call.
+            base.prev_env_reg = load_outer_env(builder, base.self_reg, outer_env, borrow=True)
         env_reg = base.prev_env_reg
         index -= 1
 

mypyc/test-data/irbuild-basic.test

@@ -2450,7 +2450,7 @@ def g_a_obj.__call__(__mypyc_self__):
     r15 :: object_ptr
     r16 :: object
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = 'Entering'
     r2 = builtins :: module
     r3 = 'print'
@@ -2514,7 +2514,7 @@ def g_b_obj.__call__(__mypyc_self__):
     r15 :: object_ptr
     r16 :: object
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = '---'
     r2 = builtins :: module
     r3 = 'print'
@@ -2571,7 +2571,7 @@ def d_c_obj.__call__(__mypyc_self__):
     r6 :: object_ptr
     r7 :: object
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = 'd'
     r2 = builtins :: module
     r3 = 'print'
@@ -2747,7 +2747,7 @@ def g_a_obj.__call__(__mypyc_self__):
     r15 :: object_ptr
     r16 :: object
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = 'Entering'
     r2 = builtins :: module
     r3 = 'print'
@@ -3575,7 +3575,7 @@ def wrapper_deco_obj.__call__(__mypyc_self__, args):
     r0 :: __main__.deco_env
     r1, r2 :: object
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = r0.fn
     r2 = PyObject_CallObject(r1, args)
     return r2
@@ -3622,7 +3622,7 @@ def wrapper_deco_obj.__call__(__mypyc_self__, args):
     r2 :: tuple
     r3 :: object
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = r0.fn
     r2 = PyList_AsTuple(args)
     r3 = PyObject_CallObject(r1, r2)
@@ -3672,7 +3672,7 @@ def wrapper_deco_obj.__call__(__mypyc_self__, lst, kwargs):
     r3 :: dict
     r4 :: object
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = r0.fn
     r2 = PyList_AsTuple(lst)
     r3 = PyDict_Copy(kwargs)
@@ -3721,7 +3721,7 @@ def wrapper_deco_obj.__call__(__mypyc_self__, args):
     r2 :: tuple
     r3 :: object
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = r0.fn
     r2 = PySequence_Tuple(args)
     r3 = PyObject_CallObject(r1, r2)
@@ -3771,7 +3771,7 @@ def wrapper_deco_obj.__call__(__mypyc_self__, args, kwargs):
     r3 :: dict
     r4 :: object
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = r0.fn
     r2 = PySequence_Tuple(args)
     r3 = PyDict_Copy(kwargs)

mypyc/test-data/irbuild-generics.test

@@ -701,7 +701,7 @@ def inner_deco_obj.__call__(__mypyc_self__, args, kwargs):
     r26 :: object
     r27 :: int
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = var_object_size args
     r2 = PyList_New(r1)
     r3 = 0

mypyc/test-data/irbuild-nested.test

@@ -52,7 +52,7 @@ def inner_a_obj.__call__(__mypyc_self__):
     r0 :: __main__.a_env
     r1 :: object
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = box(None, 1)
     return r1
 def a():
@@ -84,7 +84,7 @@ def second_b_first_obj.__call__(__mypyc_self__):
     r1 :: __main__.b_env
     r2 :: str
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = r0.__mypyc_env__
     r2 = 'b.first.second: nested function'
     return r2
@@ -109,7 +109,7 @@ def first_b_obj.__call__(__mypyc_self__):
     r3 :: __main__.second_b_first_obj
     second :: object
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = first_b_env()
     r1.__mypyc_env__ = r0; r2 = is_error
     r3 = second_b_first_obj()
@@ -145,7 +145,7 @@ def inner_c_obj.__call__(__mypyc_self__, s):
     r0 :: __main__.c_env
     r1, r2 :: str
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = '!'
     r2 = PyUnicode_Concat(s, r1)
     return r2
@@ -179,7 +179,7 @@ def inner_d_obj.__call__(__mypyc_self__, s):
     r0 :: __main__.d_env
     r1, r2 :: str
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = '?'
     r2 = PyUnicode_Concat(s, r1)
     return r2
@@ -278,7 +278,7 @@ def inner_a_obj.__call__(__mypyc_self__):
     r0 :: __main__.a_env
     r1 :: int
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = r0.num
     return r1
 def a(num):
@@ -316,7 +316,7 @@ def inner_b_obj.__call__(__mypyc_self__):
     r1 :: bool
     foo, r2 :: int
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r0.num = 8; r1 = is_error
     foo = 12
     r2 = r0.num
@@ -356,7 +356,7 @@ def inner_c_obj.__call__(__mypyc_self__):
     r0 :: __main__.c_env
     r1 :: str
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = 'f.inner: first definition'
     return r1
 def inner_c_obj_0.__get__(__mypyc_self__, instance, owner):
@@ -377,7 +377,7 @@ def inner_c_obj_0.__call__(__mypyc_self__):
     r0 :: __main__.c_env
     r1 :: str
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = 'f.inner: second definition'
     return r1
 def c(flag):
@@ -435,7 +435,7 @@ def c_a_b_obj.__call__(__mypyc_self__):
     r1 :: __main__.a_env
     r2 :: int
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = r0.__mypyc_env__
     r2 = r1.x
     return r2
@@ -463,7 +463,7 @@ def b_a_obj.__call__(__mypyc_self__):
     c, r7 :: object
     r8 :: int
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = b_a_env()
     r1.__mypyc_env__ = r0; r2 = is_error
     r3 = r0.x
@@ -519,7 +519,7 @@ def inner_f_obj.__call__(__mypyc_self__):
     r0 :: __main__.f_env
     r1 :: str
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = 'f.inner: first definition'
     return r1
 def inner_f_obj_0.__get__(__mypyc_self__, instance, owner):
@@ -540,7 +540,7 @@ def inner_f_obj_0.__call__(__mypyc_self__):
     r0 :: __main__.f_env
     r1 :: str
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = 'f.inner: second definition'
     return r1
 def f(flag):
@@ -604,7 +604,7 @@ def foo_f_obj.__call__(__mypyc_self__):
     r0 :: __main__.f_env
     r1, r2 :: int
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = r0.a
     r2 = CPyTagged_Add(r1, 2)
     return r2
@@ -627,7 +627,7 @@ def bar_f_obj.__call__(__mypyc_self__):
     r1, r2 :: object
     r3 :: int
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = r0.foo
     r2 = PyObject_Vectorcall(r1, 0, 0, 0)
     r3 = unbox(int, r2)
@@ -657,7 +657,7 @@ def baz_f_obj.__call__(__mypyc_self__, n):
     r7 :: object
     r8, r9 :: int
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = int_eq n, 0
     if r1 goto L1 else goto L2 :: bool
 L1:
@@ -742,7 +742,7 @@ def __mypyc_lambda__0_f_obj.__call__(__mypyc_self__, a, b):
     r0 :: __main__.f_env
     r1 :: object
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = PyNumber_Add(a, b)
     return r1
 def __mypyc_lambda__1_f_obj.__get__(__mypyc_self__, instance, owner):
@@ -767,7 +767,7 @@ def __mypyc_lambda__1_f_obj.__call__(__mypyc_self__, a, b):
     r3 :: object_ptr
     r4 :: object
 L0:
-    r0 = __mypyc_self__.__mypyc_env__
+    r0 = borrow __mypyc_self__.__mypyc_env__
     r1 = r0.s
     r2 = [a, b]
     r3 = load_address r2

mypyc/test-data/run-functions.test

@@ -1348,6 +1348,7 @@ def test_nested() -> None:
 
 [case testNestedFunctionEnvironmentIsReadOnly]
 from typing import Any, Callable
+from testutil import assertRaises
 
 def outer(value: str) -> Callable[[], str]:
     def inner() -> str:
@@ -1357,12 +1358,8 @@ def outer(value: str) -> Callable[[], str]:
 def test_environment_link_is_read_only() -> None:
     fn: Any = outer("value")
     environment = fn.__mypyc_env__
-    try:
+    with assertRaises(AttributeError):
         fn.__mypyc_env__ = None
-    except AttributeError:
-        pass
-    else:
-        assert False
     assert fn.__mypyc_env__ is environment
     assert fn() == "value"