[mypyc] Add irbuild support for vec types (python#20724)

Add basic irbuild support for `vec[t]`. The runtime representation of
`vec[t]` is a C struct. This is similar to how fixed-length tuples are
represented. Multiple different structs are used, depending on the item
type (`VecI32` for `vec[i32`] and so on).

The C extension `librt.vec` that defines the `vec` type was added in
python#20653 and python#20656. These PRs also explain the implementation in more
detail.

Add RType subclass RVec that is used for vecs. We need a new RType
class, since primitives types can't be generic and they can't be struct
types.

This is based on an old branch, so it mostly uses old-style primitives.
I am planning to modernize some of the primitives in follow-up PRs.

This doesn't include codegen support, and irbuild test cases are only
included for `vec[i64]`. I will create follow-up PRs that add the
remaining irbuild tests, codegen support and run tests. All these tests
are are passing on my local full branch.

Related issue: mypyc/mypyc#840

Author: JukkaL

mypyc/analysis/capsule_deps.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
+from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.deps import Dependency
 from mypyc.ir.func_ir import FuncIR
 from mypyc.ir.ops import Assign, CallC, PrimitiveOp
-from mypyc.ir.rtypes import RStruct, RTuple, RType, RUnion
+from mypyc.ir.rtypes import RStruct, RTuple, RType, RUnion, RVec
 
 
 def find_implicit_op_dependencies(fn: FuncIR) -> set[Dependency] | None:
@@ -48,6 +49,15 @@ def find_type_dependencies(fn: FuncIR, deps: set[Dependency] | None) -> set[Depe
     return deps
 
 
+def find_class_dependencies(cl: ClassIR) -> set[Dependency] | None:
+    """Find dependencies from class attribute types."""
+    deps: set[Dependency] | None = None
+    for base in cl.mro:
+        for attr_type in base.attributes.values():
+            deps = collect_type_deps(attr_type, deps)
+    return deps
+
+
 def collect_type_deps(typ: RType, deps: set[Dependency] | None) -> set[Dependency] | None:
     """Collect dependencies from an RType, recursively checking compound types."""
     if typ.dependencies is not None:
@@ -64,4 +74,6 @@ def collect_type_deps(typ: RType, deps: set[Dependency] | None) -> set[Dependenc
     elif isinstance(typ, RStruct):
         for item in typ.types:
             deps = collect_type_deps(item, deps)
+    elif isinstance(typ, RVec):
+        deps = collect_type_deps(typ.item_type, deps)
     return deps

mypyc/analysis/dataflow.py

@@ -197,6 +197,12 @@ def visit_assign_multi(self, op: AssignMulti) -> GenAndKill[T]:
     def visit_set_mem(self, op: SetMem) -> GenAndKill[T]:
         raise NotImplementedError
 
+    def visit_inc_ref(self, op: IncRef) -> GenAndKill[T]:
+        return self.visit_register_op(op)
+
+    def visit_dec_ref(self, op: DecRef) -> GenAndKill[T]:
+        return self.visit_register_op(op)
+
     def visit_call(self, op: Call) -> GenAndKill[T]:
         return self.visit_register_op(op)
 

mypyc/analysis/ircheck.py

@@ -63,12 +63,17 @@
     RPrimitive,
     RType,
     RUnion,
+    RVec,
     bytes_rprimitive,
     dict_rprimitive,
     int_rprimitive,
+    is_c_py_ssize_t_rprimitive,
+    is_fixed_width_rtype,
     is_float_rprimitive,
     is_object_rprimitive,
+    is_pointer_rprimitive,
     list_rprimitive,
+    pointer_rprimitive,
     range_rprimitive,
     set_rprimitive,
     str_rprimitive,
@@ -200,7 +205,7 @@ def can_coerce_to(src: RType, dest: RType) -> bool:
             if src.name in disjoint_types and dest.name in disjoint_types:
                 return src.name == dest.name
             return src.size == dest.size
-        if isinstance(src, RInstance):
+        if isinstance(src, (RInstance, RVec)):
             return is_object_rprimitive(dest)
         if isinstance(src, RUnion):
             # IR doesn't have the ability to narrow unions based on
@@ -211,6 +216,29 @@ def can_coerce_to(src: RType, dest: RType) -> bool:
     return True
 
 
+def is_valid_ptr_displacement_type(rtype: RType) -> bool:
+    """Check if rtype is a valid displacement type for pointer arithmetic."""
+    if not (is_fixed_width_rtype(rtype) or is_c_py_ssize_t_rprimitive(rtype)):
+        return False
+    assert isinstance(rtype, RPrimitive)
+    return rtype.size == pointer_rprimitive.size
+
+
+def is_pointer_arithmetic(op: IntOp) -> bool:
+    """Check if op is add/subtract targeting pointer_rprimitive and integer of the same size."""
+    if op.op not in (IntOp.ADD, IntOp.SUB):
+        return False
+    if not is_pointer_rprimitive(op.type):
+        return False
+    left = op.lhs.type
+    right = op.rhs.type
+    if is_pointer_rprimitive(left):
+        return is_valid_ptr_displacement_type(right)
+    if is_pointer_rprimitive(right):
+        return is_valid_ptr_displacement_type(left)
+    return False
+
+
 class OpChecker(OpVisitor[None]):
     def __init__(self, parent_fn: FuncIR) -> None:
         self.parent_fn = parent_fn
@@ -417,6 +445,7 @@ def visit_int_op(self, op: IntOp) -> None:
                 op_str in ("+", "-", "*", "/", "%")
                 or (op_str not in ("<<", ">>") and left.size != right.size)
             )
+            and not is_pointer_arithmetic(op)
         ):
             self.fail(op, f"Operand types have incompatible signs: {left}, {right}")
 

mypyc/analysis/selfleaks.py

@@ -11,6 +11,7 @@
     CallC,
     Cast,
     ComparisonOp,
+    DecRef,
     Extend,
     FloatComparisonOp,
     FloatNeg,
@@ -19,6 +20,7 @@
     GetElement,
     GetElementPtr,
     Goto,
+    IncRef,
     InitStatic,
     IntOp,
     KeepAlive,
@@ -90,6 +92,12 @@ def visit_assign_multi(self, op: AssignMulti) -> GenAndKill:
     def visit_set_mem(self, op: SetMem) -> GenAndKill:
         return CLEAN
 
+    def visit_inc_ref(self, op: IncRef) -> GenAndKill:
+        return CLEAN
+
+    def visit_dec_ref(self, op: DecRef) -> GenAndKill:
+        return CLEAN
+
     def visit_call(self, op: Call) -> GenAndKill:
         fn = op.fn
         if fn.class_name and fn.name == "__init__":

mypyc/ir/deps.py

@@ -56,3 +56,4 @@ def get_header(self) -> str:
 STRING_WRITER_EXTRA_OPS: Final = SourceDep("stringwriter_extra_ops.c")
 BYTEARRAY_EXTRA_OPS: Final = SourceDep("bytearray_extra_ops.c")
 STR_EXTRA_OPS: Final = SourceDep("str_extra_ops.c")
+VECS_EXTRA_OPS: Final = SourceDep("vecs_extra_ops.c")

mypyc/ir/ops.py

@@ -39,6 +39,7 @@ class to enable the new behavior. Sometimes adding a new abstract
     RTuple,
     RType,
     RUnion,
+    RVec,
     RVoid,
     bit_rprimitive,
     bool_rprimitive,
@@ -1690,7 +1691,7 @@ class GetElement(RegisterOp):
 
     def __init__(self, src: Value, field: str, line: int = -1) -> None:
         super().__init__(line)
-        assert isinstance(src.type, RStruct)
+        assert isinstance(src.type, (RStruct, RVec))
         self.type = src.type.field_type(field)
         self.src = src
         self.src_type = src.type
@@ -1720,7 +1721,7 @@ class GetElementPtr(RegisterOp):
 
     def __init__(self, src: Value, src_type: RType, field: str, line: int = -1) -> None:
         super().__init__(line)
-        assert not isinstance(src.type, RStruct)
+        assert not isinstance(src.type, (RStruct, RVec))
         self.type = pointer_rprimitive
         self.src = src
         self.src_type = src_type
@@ -1750,7 +1751,7 @@ class SetElement(RegisterOp):
 
     def __init__(self, src: Value, field: str, item: Value, line: int = -1) -> None:
         super().__init__(line)
-        assert isinstance(src.type, RStruct), src.type
+        assert isinstance(src.type, (RStruct, RVec)), src.type
         self.type = src.type
         self.src = src
         self.item = item