Major fixes and new features

venv/lib/python3.12/site-packages/mypyc/ir/class_ir.py

@@ -0,0 +1,499 @@
+"""Intermediate representation of classes."""
+
+from __future__ import annotations
+
+from typing import List, NamedTuple
+
+from mypyc.common import PROPSET_PREFIX, JsonDict
+from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature
+from mypyc.ir.ops import DeserMaps, Value
+from mypyc.ir.rtypes import RInstance, RType, deserialize_type
+from mypyc.namegen import NameGenerator, exported_name
+
+# Some notes on the vtable layout: Each concrete class has a vtable
+# that contains function pointers for its methods. So that subclasses
+# may be efficiently used when their parent class is expected, the
+# layout of child vtables must be an extension of their base class's
+# vtable.
+#
+# This makes multiple inheritance tricky, since obviously we cannot be
+# an extension of multiple parent classes. We solve this by requiring
+# all but one parent to be "traits", which we can operate on in a
+# somewhat less efficient way. For each trait implemented by a class,
+# we generate a separate vtable for the methods in that trait.
+# We then store an array of (trait type, trait vtable) pointers alongside
+# a class's main vtable. When we want to call a trait method, we
+# (at runtime!) search the array of trait vtables to find the correct one,
+# then call through it.
+# Trait vtables additionally need entries for attribute getters and setters,
+# since they can't always be in the same location.
+#
+# To keep down the number of indirections necessary, we store the
+# array of trait vtables in the memory *before* the class vtable, and
+# search it backwards.  (This is a trick we can only do once---there
+# are only two directions to store data in---but I don't think we'll
+# need it again.)
+# There are some tricks we could try in the future to store the trait
+# vtables inline in the trait table (which would cut down one indirection),
+# but this seems good enough for now.
+#
+# As an example:
+# Imagine that we have a class B that inherits from a concrete class A
+# and traits T1 and T2, and that A has methods foo() and
+# bar() and B overrides bar() with a more specific type.
+# Then B's vtable will look something like:
+#
+#      T1 type object
+#      ptr to B's T1 trait vtable
+#      T2 type object
+#      ptr to B's T2 trait vtable
+# -> | A.foo
+#    | Glue function that converts between A.bar's type and B.bar
+#      B.bar
+#      B.baz
+#
+# The arrow points to the "start" of the vtable (what vtable pointers
+# point to) and the bars indicate which parts correspond to the parent
+# class A's vtable layout.
+#
+# Classes that allow interpreted code to subclass them also have a
+# "shadow vtable" that contains implementations that delegate to
+# making a pycall, so that overridden methods in interpreted children
+# will be called. (A better strategy could dynamically generate these
+# vtables based on which methods are overridden in the children.)
+
+# Descriptions of method and attribute entries in class vtables.
+# The 'cls' field is the class that the method/attr was defined in,
+# which might be a parent class.
+# The 'shadow_method', if present, contains the method that should be
+# placed in the class's shadow vtable (if it has one).
+
+
+class VTableMethod(NamedTuple):
+    cls: "ClassIR"
+    name: str
+    method: FuncIR
+    shadow_method: FuncIR | None
+
+
+VTableEntries = List[VTableMethod]
+
+
+class ClassIR:
+    """Intermediate representation of a class.
+
+    This also describes the runtime structure of native instances.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        module_name: str,
+        is_trait: bool = False,
+        is_generated: bool = False,
+        is_abstract: bool = False,
+        is_ext_class: bool = True,
+    ) -> None:
+        self.name = name
+        self.module_name = module_name
+        self.is_trait = is_trait
+        self.is_generated = is_generated
+        self.is_abstract = is_abstract
+        self.is_ext_class = is_ext_class
+        # An augmented class has additional methods separate from what mypyc generates.
+        # Right now the only one is dataclasses.
+        self.is_augmented = False
+        # Does this inherit from a Python class?
+        self.inherits_python = False
+        # Do instances of this class have __dict__?
+        self.has_dict = False
+        # Do we allow interpreted subclasses? Derived from a mypyc_attr.
+        self.allow_interpreted_subclasses = False
+        # Does this class need getseters to be generated for its attributes? (getseters are also
+        # added if is_generated is False)
+        self.needs_getseters = False
+        # Is this class declared as serializable (supports copy.copy
+        # and pickle) using @mypyc_attr(serializable=True)?
+        #
+        # Additionally, any class with this attribute False but with
+        # an __init__ that can be called without any arguments is
+        # *implicitly serializable*. In this case __init__ will be
+        # called during deserialization without arguments. If this is
+        # True, we match Python semantics and __init__ won't be called
+        # during deserialization.
+        #
+        # This impacts also all subclasses. Use is_serializable() to
+        # also consider base classes.
+        self._serializable = False
+        # If this a subclass of some built-in python class, the name
+        # of the object for that class. We currently only support this
+        # in a few ad-hoc cases.
+        self.builtin_base: str | None = None
+        # Default empty constructor
+        self.ctor = FuncDecl(name, None, module_name, FuncSignature([], RInstance(self)))
+        # Attributes defined in the class (not inherited)
+        self.attributes: dict[str, RType] = {}
+        # Deletable attributes
+        self.deletable: list[str] = []
+        # We populate method_types with the signatures of every method before
+        # we generate methods, and we rely on this information being present.
+        self.method_decls: dict[str, FuncDecl] = {}
+        # Map of methods that are actually present in an extension class
+        self.methods: dict[str, FuncIR] = {}
+        # Glue methods for boxing/unboxing when a class changes the type
+        # while overriding a method. Maps from (parent class overridden, method)
+        # to IR of glue method.
+        self.glue_methods: dict[tuple[ClassIR, str], FuncIR] = {}
+
+        # Properties are accessed like attributes, but have behavior like method calls.
+        # They don't belong in the methods dictionary, since we don't want to expose them to
+        # Python's method API. But we want to put them into our own vtable as methods, so that
+        # they are properly handled and overridden. The property dictionary values are a tuple
+        # containing a property getter and an optional property setter.
+        self.properties: dict[str, tuple[FuncIR, FuncIR | None]] = {}
+        # We generate these in prepare_class_def so that we have access to them when generating
+        # other methods and properties that rely on these types.
+        self.property_types: dict[str, RType] = {}
+
+        self.vtable: dict[str, int] | None = None
+        self.vtable_entries: VTableEntries = []
+        self.trait_vtables: dict[ClassIR, VTableEntries] = {}
+        # N.B: base might not actually quite be the direct base.
+        # It is the nearest concrete base, but we allow a trait in between.
+        self.base: ClassIR | None = None
+        self.traits: list[ClassIR] = []
+        # Supply a working mro for most generated classes. Real classes will need to
+        # fix it up.
+        self.mro: list[ClassIR] = [self]
+        # base_mro is the chain of concrete (non-trait) ancestors
+        self.base_mro: list[ClassIR] = [self]
+
+        # Direct subclasses of this class (use subclasses() to also include non-direct ones)
+        # None if separate compilation prevents this from working.
+        #
+        # Often it's better to use has_no_subclasses() or subclasses() instead.
+        self.children: list[ClassIR] | None = []
+
+        # Instance attributes that are initialized in the class body.
+        self.attrs_with_defaults: set[str] = set()
+
+        # Attributes that are always initialized in __init__ or class body
+        # (inferred in mypyc.analysis.attrdefined using interprocedural analysis)
+        self._always_initialized_attrs: set[str] = set()
+
+        # Attributes that are sometimes initialized in __init__
+        self._sometimes_initialized_attrs: set[str] = set()
+
+        # If True, __init__ can make 'self' visible to unanalyzed/arbitrary code
+        self.init_self_leak = False
+
+        # Definedness of these attributes is backed by a bitmap. Index in the list
+        # indicates the bit number. Includes inherited attributes. We need the
+        # bitmap for types such as native ints that can't have a dedicated error
+        # value that doesn't overlap a valid value. The bitmap is used if the
+        # value of an attribute is the same as the error value.
+        self.bitmap_attrs: list[str] = []
+
+    def __repr__(self) -> str:
+        return (
+            "ClassIR("
+            "name={self.name}, module_name={self.module_name}, "
+            "is_trait={self.is_trait}, is_generated={self.is_generated}, "
+            "is_abstract={self.is_abstract}, is_ext_class={self.is_ext_class}"
+            ")".format(self=self)
+        )
+
+    @property
+    def fullname(self) -> str:
+        return f"{self.module_name}.{self.name}"
+
+    def real_base(self) -> ClassIR | None:
+        """Return the actual concrete base class, if there is one."""
+        if len(self.mro) > 1 and not self.mro[1].is_trait:
+            return self.mro[1]
+        return None
+
+    def vtable_entry(self, name: str) -> int:
+        assert self.vtable is not None, "vtable not computed yet"
+        assert name in self.vtable, f"{self.name!r} has no attribute {name!r}"
+        return self.vtable[name]
+
+    def attr_details(self, name: str) -> tuple[RType, ClassIR]:
+        for ir in self.mro:
+            if name in ir.attributes:
+                return ir.attributes[name], ir
+            if name in ir.property_types:
+                return ir.property_types[name], ir
+        raise KeyError(f"{self.name!r} has no attribute {name!r}")
+
+    def attr_type(self, name: str) -> RType:
+        return self.attr_details(name)[0]
+
+    def method_decl(self, name: str) -> FuncDecl:
+        for ir in self.mro:
+            if name in ir.method_decls:
+                return ir.method_decls[name]
+        raise KeyError(f"{self.name!r} has no attribute {name!r}")
+
+    def method_sig(self, name: str) -> FuncSignature:
+        return self.method_decl(name).sig
+
+    def has_method(self, name: str) -> bool:
+        try:
+            self.method_decl(name)
+        except KeyError:
+            return False
+        return True
+
+    def is_method_final(self, name: str) -> bool:
+        subs = self.subclasses()
+        if subs is None:
+            # TODO: Look at the final attribute!
+            return False
+
+        if self.has_method(name):
+            method_decl = self.method_decl(name)
+            for subc in subs:
+                if subc.method_decl(name) != method_decl:
+                    return False
+            return True
+        else:
+            return not any(subc.has_method(name) for subc in subs)
+
+    def has_attr(self, name: str) -> bool:
+        try:
+            self.attr_type(name)
+        except KeyError:
+            return False
+        return True
+
+    def is_deletable(self, name: str) -> bool:
+        return any(name in ir.deletable for ir in self.mro)
+
+    def is_always_defined(self, name: str) -> bool:
+        if self.is_deletable(name):
+            return False
+        return name in self._always_initialized_attrs
+
+    def name_prefix(self, names: NameGenerator) -> str:
+        return names.private_name(self.module_name, self.name)
+
+    def struct_name(self, names: NameGenerator) -> str:
+        return f"{exported_name(self.fullname)}Object"
+
+    def get_method_and_class(
+        self, name: str, *, prefer_method: bool = False
+    ) -> tuple[FuncIR, ClassIR] | None:
+        for ir in self.mro:
+            if name in ir.methods:
+                func_ir = ir.methods[name]
+                if not prefer_method and func_ir.decl.implicit:
+                    # This is an implicit accessor, so there is also an attribute definition
+                    # which the caller prefers. This happens if an attribute overrides a
+                    # property.
+                    return None
+                return func_ir, ir
+
+        return None
+
+    def get_method(self, name: str, *, prefer_method: bool = False) -> FuncIR | None:
+        res = self.get_method_and_class(name, prefer_method=prefer_method)
+        return res[0] if res else None
+
+    def has_method_decl(self, name: str) -> bool:
+        return any(name in ir.method_decls for ir in self.mro)
+
+    def has_no_subclasses(self) -> bool:
+        return self.children == [] and not self.allow_interpreted_subclasses
+
+    def subclasses(self) -> set[ClassIR] | None:
+        """Return all subclasses of this class, both direct and indirect.
+
+        Return None if it is impossible to identify all subclasses, for example
+        because we are performing separate compilation.
+        """
+        if self.children is None or self.allow_interpreted_subclasses:
+            return None
+        result = set(self.children)
+        for child in self.children:
+            if child.children:
+                child_subs = child.subclasses()
+                if child_subs is None:
+                    return None
+                result.update(child_subs)
+        return result
+
+    def concrete_subclasses(self) -> list[ClassIR] | None:
+        """Return all concrete (i.e. non-trait and non-abstract) subclasses.
+
+        Include both direct and indirect subclasses. Place classes with no children first.
+        """
+        subs = self.subclasses()
+        if subs is None:
+            return None
+        concrete = {c for c in subs if not (c.is_trait or c.is_abstract)}
+        # We place classes with no children first because they are more likely
+        # to appear in various isinstance() checks. We then sort leaves by name
+        # to get stable order.
+        return sorted(concrete, key=lambda c: (len(c.children or []), c.name))
+
+    def is_serializable(self) -> bool:
+        return any(ci._serializable for ci in self.mro)
+
+    def serialize(self) -> JsonDict:
+        return {
+            "name": self.name,
+            "module_name": self.module_name,
+            "is_trait": self.is_trait,
+            "is_ext_class": self.is_ext_class,
+            "is_abstract": self.is_abstract,
+            "is_generated": self.is_generated,
+            "is_augmented": self.is_augmented,
+            "inherits_python": self.inherits_python,
+            "has_dict": self.has_dict,
+            "allow_interpreted_subclasses": self.allow_interpreted_subclasses,
+            "needs_getseters": self.needs_getseters,
+            "_serializable": self._serializable,
+            "builtin_base": self.builtin_base,
+            "ctor": self.ctor.serialize(),
+            # We serialize dicts as lists to ensure order is preserved
+            "attributes": [(k, t.serialize()) for k, t in self.attributes.items()],
+            # We try to serialize a name reference, but if the decl isn't in methods
+            # then we can't be sure that will work so we serialize the whole decl.
+            "method_decls": [
+                (k, d.id if k in self.methods else d.serialize())
+                for k, d in self.method_decls.items()
+            ],
+            # We serialize method fullnames out and put methods in a separate dict
+            "methods": [(k, m.id) for k, m in self.methods.items()],
+            "glue_methods": [
+                ((cir.fullname, k), m.id) for (cir, k), m in self.glue_methods.items()
+            ],
+            # We serialize properties and property_types separately out of an
+            # abundance of caution about preserving dict ordering...
+            "property_types": [(k, t.serialize()) for k, t in self.property_types.items()],
+            "properties": list(self.properties),
+            "vtable": self.vtable,
+            "vtable_entries": serialize_vtable(self.vtable_entries),
+            "trait_vtables": [
+                (cir.fullname, serialize_vtable(v)) for cir, v in self.trait_vtables.items()
+            ],
+            # References to class IRs are all just names
+            "base": self.base.fullname if self.base else None,
+            "traits": [cir.fullname for cir in self.traits],
+            "mro": [cir.fullname for cir in self.mro],
+            "base_mro": [cir.fullname for cir in self.base_mro],
+            "children": [cir.fullname for cir in self.children]
+            if self.children is not None
+            else None,
+            "deletable": self.deletable,
+            "attrs_with_defaults": sorted(self.attrs_with_defaults),
+            "_always_initialized_attrs": sorted(self._always_initialized_attrs),
+            "_sometimes_initialized_attrs": sorted(self._sometimes_initialized_attrs),
+            "init_self_leak": self.init_self_leak,
+        }
+
+    @classmethod
+    def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ClassIR:
+        fullname = data["module_name"] + "." + data["name"]
+        assert fullname in ctx.classes, "Class %s not in deser class map" % fullname
+        ir = ctx.classes[fullname]
+
+        ir.is_trait = data["is_trait"]
+        ir.is_generated = data["is_generated"]
+        ir.is_abstract = data["is_abstract"]
+        ir.is_ext_class = data["is_ext_class"]
+        ir.is_augmented = data["is_augmented"]
+        ir.inherits_python = data["inherits_python"]
+        ir.has_dict = data["has_dict"]
+        ir.allow_interpreted_subclasses = data["allow_interpreted_subclasses"]
+        ir.needs_getseters = data["needs_getseters"]
+        ir._serializable = data["_serializable"]
+        ir.builtin_base = data["builtin_base"]
+        ir.ctor = FuncDecl.deserialize(data["ctor"], ctx)
+        ir.attributes = {k: deserialize_type(t, ctx) for k, t in data["attributes"]}
+        ir.method_decls = {
+            k: ctx.functions[v].decl if isinstance(v, str) else FuncDecl.deserialize(v, ctx)
+            for k, v in data["method_decls"]
+        }
+        ir.methods = {k: ctx.functions[v] for k, v in data["methods"]}
+        ir.glue_methods = {
+            (ctx.classes[c], k): ctx.functions[v] for (c, k), v in data["glue_methods"]
+        }
+        ir.property_types = {k: deserialize_type(t, ctx) for k, t in data["property_types"]}
+        ir.properties = {
+            k: (ir.methods[k], ir.methods.get(PROPSET_PREFIX + k)) for k in data["properties"]
+        }
+
+        ir.vtable = data["vtable"]
+        ir.vtable_entries = deserialize_vtable(data["vtable_entries"], ctx)
+        ir.trait_vtables = {
+            ctx.classes[k]: deserialize_vtable(v, ctx) for k, v in data["trait_vtables"]
+        }
+
+        base = data["base"]
+        ir.base = ctx.classes[base] if base else None
+        ir.traits = [ctx.classes[s] for s in data["traits"]]
+        ir.mro = [ctx.classes[s] for s in data["mro"]]
+        ir.base_mro = [ctx.classes[s] for s in data["base_mro"]]
+        ir.children = data["children"] and [ctx.classes[s] for s in data["children"]]
+        ir.deletable = data["deletable"]
+        ir.attrs_with_defaults = set(data["attrs_with_defaults"])
+        ir._always_initialized_attrs = set(data["_always_initialized_attrs"])
+        ir._sometimes_initialized_attrs = set(data["_sometimes_initialized_attrs"])
+        ir.init_self_leak = data["init_self_leak"]
+
+        return ir
+
+
+class NonExtClassInfo:
+    """Information needed to construct a non-extension class (Python class).
+
+    Includes the class dictionary, a tuple of base classes,
+    the class annotations dictionary, and the metaclass.
+    """
+
+    def __init__(self, dict: Value, bases: Value, anns: Value, metaclass: Value) -> None:
+        self.dict = dict
+        self.bases = bases
+        self.anns = anns
+        self.metaclass = metaclass
+
+
+def serialize_vtable_entry(entry: VTableMethod) -> JsonDict:
+    return {
+        ".class": "VTableMethod",
+        "cls": entry.cls.fullname,
+        "name": entry.name,
+        "method": entry.method.decl.id,
+        "shadow_method": entry.shadow_method.decl.id if entry.shadow_method else None,
+    }
+
+
+def serialize_vtable(vtable: VTableEntries) -> list[JsonDict]:
+    return [serialize_vtable_entry(v) for v in vtable]
+
+
+def deserialize_vtable_entry(data: JsonDict, ctx: DeserMaps) -> VTableMethod:
+    if data[".class"] == "VTableMethod":
+        return VTableMethod(
+            ctx.classes[data["cls"]],
+            data["name"],
+            ctx.functions[data["method"]],
+            ctx.functions[data["shadow_method"]] if data["shadow_method"] else None,
+        )
+    assert False, "Bogus vtable .class: %s" % data[".class"]
+
+
+def deserialize_vtable(data: list[JsonDict], ctx: DeserMaps) -> VTableEntries:
+    return [deserialize_vtable_entry(x, ctx) for x in data]
+
+
+def all_concrete_classes(class_ir: ClassIR) -> list[ClassIR] | None:
+    """Return all concrete classes among the class itself and its subclasses."""
+    concrete = class_ir.concrete_subclasses()
+    if concrete is None:
+        return None
+    if not (class_ir.is_abstract or class_ir.is_trait):
+        concrete.append(class_ir)
+    return concrete