[mypyc] Allow defining a single-item free "list" for a native class (#19785)

It's quite common to have a class where we almost always have at most a
single allocated instance (per thread). Now these instances can be
allocated more quickly, by reusing the memory that was used for the most
recently freed instance (a separate memory block is reused for each
thread on free-threaded builds).

It's used like this (only the value 1 is supported for now):
```
from mypy_extensions import mypyc_attr

@mypyc_attr(free_list=1)
class Foo:
    ...
```

This makes a microbenchmark that only allocates and immediately frees
simple objects repeatedly around 3.8x faster.

It's probably worth extending this to support larger free lists in the
future.

We can later look into enabling this automatically for certain native
classes based on profile information.

Author: JukkaL

mypyc/irbuild/prepare.py

@@ -351,13 +351,23 @@ def prepare_class_def(
     ir = mapper.type_to_ir[cdef.info]
     info = cdef.info
 
-    attrs = get_mypyc_attrs(cdef)
+    attrs, attrs_lines = get_mypyc_attrs(cdef)
     if attrs.get("allow_interpreted_subclasses") is True:
         ir.allow_interpreted_subclasses = True
     if attrs.get("serializable") is True:
         # Supports copy.copy and pickle (including subclasses)
         ir._serializable = True
 
+    free_list_len = attrs.get("free_list_len")
+    if free_list_len is not None:
+        line = attrs_lines["free_list_len"]
+        if ir.is_trait:
+            errors.error('"free_list_len" can\'t be used with traits', path, line)
+        if free_list_len == 1:
+            ir.reuse_freed_instance = True
+        else:
+            errors.error(f'Unsupported value for "free_list_len": {free_list_len}', path, line)
+
     # Check for subclassing from builtin types
     for cls in info.mro:
         # Special case exceptions and dicts

mypyc/irbuild/util.py

@@ -96,6 +96,8 @@ def get_mypyc_attr_literal(e: Expression) -> Any:
         return False
     elif isinstance(e, RefExpr) and e.fullname == "builtins.None":
         return None
+    elif isinstance(e, IntExpr):
+        return e.value
     return NotImplemented
 
 
@@ -110,20 +112,23 @@ def get_mypyc_attr_call(d: Expression) -> CallExpr | None:
     return None
 
 
-def get_mypyc_attrs(stmt: ClassDef | Decorator) -> dict[str, Any]:
+def get_mypyc_attrs(stmt: ClassDef | Decorator) -> tuple[dict[str, Any], dict[str, int]]:
     """Collect all the mypyc_attr attributes on a class definition or a function."""
     attrs: dict[str, Any] = {}
+    lines: dict[str, int] = {}
     for dec in stmt.decorators:
         d = get_mypyc_attr_call(dec)
         if d:
             for name, arg in zip(d.arg_names, d.args):
                 if name is None:
                     if isinstance(arg, StrExpr):
                         attrs[arg.value] = True
+                        lines[arg.value] = d.line
                 else:
                     attrs[name] = get_mypyc_attr_literal(arg)
+                    lines[name] = d.line
 
-    return attrs
+    return attrs, lines
 
 
 def is_extension_class(path: str, cdef: ClassDef, errors: Errors) -> bool:

mypyc/test-data/irbuild-classes.test

@@ -1734,3 +1734,23 @@ class NonNative:
 class InheritsPython(dict):
     def __new__(cls) -> InheritsPython:
         return super().__new__(cls)  # E: super().__new__() not supported for classes inheriting from non-native classes
+
+[case testClassWithFreeList]
+from mypy_extensions import mypyc_attr, trait
+
+@mypyc_attr(free_list_len=1)
+class UsesFreeList:
+    pass
+
+@mypyc_attr(free_list_len=None)
+class NoFreeList:
+    pass
+
+@mypyc_attr(free_list_len=2)  # E: Unsupported value for "free_list_len": 2
+class FreeListError:
+    pass
+
+@trait
+@mypyc_attr(free_list_len=1)  # E: "free_list_len" can't be used with traits
+class NonNative:
+    pass

mypyc/test-data/run-classes.test

@@ -3794,3 +3794,39 @@ assert t.native == 43
 assert t.generic == "{}"
 assert t.bitfield == 0x0C
 assert t.default == 10
+
+[case testPerTypeFreeList]
+from __future__ import annotations
+
+from mypy_extensions import mypyc_attr
+
+a = []
+
+@mypyc_attr(free_list=1)
+class Foo:
+    def __init__(self, x: int) -> None:
+        self.x = x
+        a.append(x)
+
+def test_alloc() -> None:
+    x: Foo | None
+    y: Foo | None
+
+    x = Foo(1)
+    assert x.x == 1
+    x = None
+
+    x = Foo(2)
+    assert x.x == 2
+    y = Foo(3)
+    assert x.x == 2
+    assert y.x == 3
+    x = None
+    y = None
+    assert a == [1, 2, 3]
+
+    x = Foo(4)
+    assert x.x == 4
+    y = Foo(5)
+    assert x.x == 4
+    assert y.x == 5