[mypyc] Update vec value struct to have direct pointer to items

mypyc/codegen/emit.py

@@ -83,6 +83,18 @@
     NAMESPACE_TYPE_VAR: TYPE_VAR_PREFIX,
 }
 
+# Map from RVec._ctype to C macro prefix for VEC_*_INCREF/DECREF/BUF macros
+VEC_MACRO_PREFIX: Final = {
+    "VecI64": "VEC_I64",
+    "VecI32": "VEC_I32",
+    "VecI16": "VEC_I16",
+    "VecU8": "VEC_U8",
+    "VecFloat": "VEC_FLOAT",
+    "VecBool": "VEC_BOOL",
+    "VecT": "VEC_T",
+    "VecNested": "VEC_NESTED",
+}
+
 
 class HeaderDeclaration:
     """A representation of a declaration in C.
@@ -348,7 +360,7 @@ def ctype_spaced(self, rtype: RType) -> str:
     def set_undefined_value(self, target: str, rtype: RType) -> None:
         if isinstance(rtype, RVec):
             self.emit_line(f"{target}.len = -1;")
-            self.emit_line(f"{target}.buf = NULL;")
+            self.emit_line(f"{target}.items = NULL;")
         else:
             self.emit_line(f"{target} = {self.c_undefined_value(rtype)};")
 
@@ -574,8 +586,8 @@ def emit_inc_ref(self, dest: str, rtype: RType, *, rare: bool = False) -> None:
             for i, item_type in enumerate(rtype.types):
                 self.emit_inc_ref(f"{dest}.f{i}", item_type)
         elif isinstance(rtype, RVec):
-            # TODO: Only use the X variant if buf can be NULL
-            self.emit_line(f"Py_XINCREF({dest}.buf);")
+            prefix = VEC_MACRO_PREFIX[rtype._ctype]
+            self.emit_line(f"{prefix}_INCREF({dest});")
         elif not rtype.is_unboxed:
             # Always inline, since this is a simple but very hot op
             if rtype.may_be_immortal or not HAVE_IMMORTAL:
@@ -605,11 +617,8 @@ def emit_dec_ref(
             for i, item_type in enumerate(rtype.types):
                 self.emit_dec_ref(f"{dest}.f{i}", item_type, is_xdec=is_xdec, rare=rare)
         elif isinstance(rtype, RVec):
-            # TODO: Only use the X variant if buf can be NULL
-            if rare:
-                self.emit_line(f"CPy_XDecRef({dest}.buf);")
-            else:
-                self.emit_line(f"CPy_XDECREF({dest}.buf);")
+            prefix = VEC_MACRO_PREFIX[rtype._ctype]
+            self.emit_line(f"{prefix}_DECREF({dest});")
         elif not rtype.is_unboxed:
             if rare:
                 self.emit_line(f"CPy_{x}DecRef({dest});")
@@ -831,7 +840,7 @@ def emit_cast(
                 item_type_c = self.vec_item_type_c(typ)
                 check = (
                     f"(Py_TYPE({src}) == VecTApi.boxed_type && "
-                    f"((VecTObject *){src})->vec.buf->item_type == {item_type_c})"
+                    f"VEC_T_BUF(((VecTObject *){src})->vec)->item_type == {item_type_c})"
                 )
             else:
                 # Nested vec types (vec[vec[...]]). Check boxed type, item type, and depth.
@@ -842,8 +851,8 @@ def emit_cast(
                     type_value = self.vec_item_type_c(typ)
                 check = (
                     f"(Py_TYPE({src}) == VecNestedApi.boxed_type && "
-                    f"((VecNestedObject *){src})->vec.buf->item_type == {type_value} && "
-                    f"((VecNestedObject *){src})->vec.buf->depth == {depth})"
+                    f"VEC_NESTED_BUF(((VecNestedObject *){src})->vec)->item_type == {type_value} && "
+                    f"VEC_NESTED_BUF(((VecNestedObject *){src})->vec)->depth == {depth})"
                 )
             if likely:
                 check = f"(likely{check})"
@@ -1284,7 +1293,8 @@ def emit_gc_visit(self, target: str, rtype: RType) -> None:
             for i, item_type in enumerate(rtype.types):
                 self.emit_gc_visit(f"{target}.f{i}", item_type)
         elif isinstance(rtype, RVec):
-            self.emit_line(f"Py_VISIT({target}.buf);")
+            prefix = VEC_MACRO_PREFIX[rtype._ctype]
+            self.emit_line(f"if ({target}.items) {{ Py_VISIT({prefix}_BUF({target})); }}")
         elif self.ctype(rtype) == "PyObject *":
             # The simplest case.
             self.emit_line(f"Py_VISIT({target});")
@@ -1310,7 +1320,11 @@ def emit_gc_clear(self, target: str, rtype: RType) -> None:
             for i, item_type in enumerate(rtype.types):
                 self.emit_gc_clear(f"{target}.f{i}", item_type)
         elif isinstance(rtype, RVec):
-            self.emit_line(f"Py_CLEAR({target}.buf);")
+            prefix = VEC_MACRO_PREFIX[rtype._ctype]
+            self.emit_line(f"if ({target}.items) {{")
+            self.emit_line(f"    Py_DECREF({prefix}_BUF({target}));")
+            self.emit_line(f"    {target}.items = NULL;")
+            self.emit_line("}")
         elif self.ctype(rtype) == "PyObject *" and self.c_undefined_value(rtype) == "NULL":
             # The simplest case.
             self.emit_line(f"Py_CLEAR({target});")

mypyc/codegen/emitfunc.py

@@ -86,6 +86,18 @@
     is_tagged,
 )
 
+VEC_ITEMS_C_TYPE: Final = {
+    "VecI64": "int64_t *",
+    "VecI32": "int32_t *",
+    "VecI16": "int16_t *",
+    "VecU8": "uint8_t *",
+    "VecFloat": "double *",
+    "VecBool": "char *",
+    "VecT": "PyObject **",
+    "VecNested": "VecNestedBufItem *",
+    "VecNestedBufItem": "void *",
+}
+
 
 def native_function_type(fn: FuncIR, emitter: Emitter) -> str:
     return native_function_type_from_decl(fn.decl, emitter)
@@ -557,7 +569,7 @@ def visit_tuple_get(self, op: TupleGet) -> None:
         dest = self.reg(op)
         src = self.reg(op.src)
         self.emit_line(f"{dest} = {src}.f{op.index};")
-        if not op.is_borrowed:
+        if not op.is_borrowed and op.type.is_refcounted:
             self.emit_inc_ref(dest, op.type)
 
     def get_dest_assign(self, dest: Value) -> str:
@@ -794,7 +806,7 @@ def visit_load_mem(self, op: LoadMem) -> None:
         # TODO: we shouldn't dereference to type that are pointer type so far
         type = self.ctype(op.type)
         self.emit_line(f"{dest} = *({type} *){src};")
-        if not op.is_borrowed:
+        if not op.is_borrowed and op.type.is_refcounted:
             self.emit_inc_ref(dest, op.type)
 
     def visit_set_mem(self, op: SetMem) -> None:
@@ -829,8 +841,8 @@ def visit_get_element_ptr(self, op: GetElementPtr) -> None:
 
     def visit_set_element(self, op: SetElement) -> None:
         dest = self.reg(op)
-        item = self.reg(op.item)
         field = op.field
+        item = self.set_element_item(op.src.type, field, self.reg(op.item))
         if isinstance(op.src, Undef):
             # First assignment to an undefined struct is trivial.
             self.emit_line(f"{dest}.{field} = {item};")
@@ -843,7 +855,7 @@ def visit_set_element(self, op: SetElement) -> None:
             # TODO: Support tuples (or use RStruct for tuples)?
             src = self.reg(op.src)
             src_type = op.src.type
-            assert isinstance(src_type, RStruct), src_type
+            assert isinstance(src_type, (RStruct, RVec)), src_type
             init_items = []
             for n in src_type.names:
                 if n != field:
@@ -852,6 +864,11 @@ def visit_set_element(self, op: SetElement) -> None:
                     init_items.append(item)
             self.emit_line(f"{dest} = ({self.ctype(src_type)}) {{ {', '.join(init_items)} }};")
 
+    def set_element_item(self, src_type: RType, field: str, item: str) -> str:
+        if field == "items" and src_type._ctype in VEC_ITEMS_C_TYPE:
+            return f"({VEC_ITEMS_C_TYPE[src_type._ctype]}){item}"
+        return item
+
     def visit_load_address(self, op: LoadAddress) -> None:
         typ = op.type
         dest = self.reg(op)

mypyc/ir/rtypes.py

@@ -1017,7 +1017,7 @@ class RVec(RType):
     def __init__(self, item_type: RType) -> None:
         self.name = "vec[%s]" % item_type
         self.item_type = item_type
-        self.names = ["len", "buf"]
+        self.names = ["len", "items"]
         self.dependencies = (LIBRT_VECS,)
         if isinstance(item_type, RUnion):
             non_opt = optional_value_type(item_type)
@@ -1026,14 +1026,14 @@ def __init__(self, item_type: RType) -> None:
         if item_type in vec_buf_types:
             self._ctype = vec_c_types[item_type]
             self.buf_type = vec_buf_types[item_type]
-            self.types = [c_pyssize_t_rprimitive, self.buf_type]
+            self.types = [c_pyssize_t_rprimitive, pointer_rprimitive]
         elif isinstance(non_opt, RVec):
             self._ctype = "VecNested"
-            self.types = [c_pyssize_t_rprimitive, VecTBufObject]
+            self.types = [c_pyssize_t_rprimitive, pointer_rprimitive]
             self.buf_type = VecNestedBufObject
         else:
             self._ctype = "VecT"
-            self.types = [c_pyssize_t_rprimitive, VecTBufObject]
+            self.types = [c_pyssize_t_rprimitive, pointer_rprimitive]
             self.buf_type = VecTBufObject
 
     @property
@@ -1076,8 +1076,8 @@ def depth(self) -> int:
     def field_type(self, name: str) -> RType:
         if name == "len":
             return c_pyssize_t_rprimitive
-        elif name == "buf":
-            return object_rprimitive
+        elif name == "items":
+            return pointer_rprimitive
         assert False, f"RVec has no field '{name}'"
 
     def accept(self, visitor: RTypeVisitor[T]) -> T:
@@ -1346,7 +1346,7 @@ def check_native_int_range(rtype: RPrimitive, n: int) -> bool:
 
 # Struct type for vec[i64] (in most cases use RVec instead).
 VecI64 = RStruct(
-    name="VecI64", names=["len", "buf"], types=[c_pyssize_t_rprimitive, object_rprimitive]
+    name="VecI64", names=["len", "items"], types=[c_pyssize_t_rprimitive, pointer_rprimitive]
 )
 
 
@@ -1359,13 +1359,13 @@ def check_native_int_range(rtype: RPrimitive, n: int) -> bool:
 
 # Struct type for vec[t] (in most cases use RVec instead).
 VecT = RStruct(
-    name="VecT", names=["len", "buf"], types=[c_pyssize_t_rprimitive, object_rprimitive]
+    name="VecT", names=["len", "items"], types=[c_pyssize_t_rprimitive, pointer_rprimitive]
 )
 
 VecNestedBufItem = RStruct(
     name="VecNestedBufItem",
-    names=["len", "buf"],
-    types=[c_pyssize_t_rprimitive, object_non_refcounted_rprimitive],
+    names=["len", "items"],
+    types=[c_pyssize_t_rprimitive, pointer_rprimitive],
 )
 
 # Buffer for vec[vec[t]]
@@ -1383,7 +1383,7 @@ def check_native_int_range(rtype: RPrimitive, n: int) -> bool:
 
 # Struct type for vec[vec[...]] (in most cases use RVec instead).
 VecNested = RStruct(
-    name="VecNested", names=["len", "buf"], types=[c_pyssize_t_rprimitive, object_rprimitive]
+    name="VecNested", names=["len", "items"], types=[c_pyssize_t_rprimitive, pointer_rprimitive]
 )
 
 VecNestedBufObject_rprimitive = RPrimitive(

mypyc/irbuild/vec.py

@@ -2,20 +2,18 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Final, cast
+from typing import TYPE_CHECKING, Final
 
 from mypyc.common import IS_32_BIT_PLATFORM, PLATFORM_SIZE
 from mypyc.ir.ops import (
     ERR_MAGIC,
-    ERR_NEVER,
     Assign,
     BasicBlock,
     Branch,
     CallC,
     ComparisonOp,
     DecRef,
     GetElement,
-    GetElementPtr,
     Integer,
     IntOp,
     RaiseStandardError,
@@ -170,7 +168,7 @@ def vec_create_initialized(
     for_loop = builder.begin_for(
         items_start, items_end, Integer(step, c_pyssize_t_rprimitive), signed=False
     )
-    builder.set_mem(for_loop.index, item_type, init)
+    vec_set_mem_item(builder, for_loop.index, item_type, init)
     for_loop.finish()
 
     builder.keep_alive([vec], line)
@@ -190,7 +188,7 @@ def vec_create_from_values(
     item_type = vtype.item_type
     step = step_size(item_type)
     for value in values:
-        builder.set_mem(ptr, item_type, value)
+        vec_set_mem_item(builder, ptr, item_type, value)
         ptr = builder.int_add(ptr, step)
     builder.keep_alive([vec], line)
     return vec
@@ -244,14 +242,11 @@ def vec_len_native(builder: LowLevelIRBuilder, val: Value) -> Value:
 
 
 def vec_items(builder: LowLevelIRBuilder, vecobj: Value) -> Value:
-    """Return pointer to first item in vec's buf.
+    """Return pointer to first item in vec.
 
-    Safe to call even when buf is NULL (empty vec), since GetElementPtr
-    uses offsetof-based arithmetic instead of &((T*)p)->field.
+    The items field points directly to the first element in the buffer.
     """
-    vtype = cast(RVec, vecobj.type)
-    buf = builder.get_element(vecobj, "buf")
-    return builder.add(GetElementPtr(buf, vtype.buf_type, "items"))
+    return builder.get_element(vecobj, "items")
 
 
 def vec_item_ptr(builder: LowLevelIRBuilder, vecobj: Value, index: Value) -> Value:
@@ -269,6 +264,20 @@ def vec_item_ptr(builder: LowLevelIRBuilder, vecobj: Value, index: Value) -> Val
     return builder.int_add(items_addr, delta)
 
 
+def vec_load_mem_item(
+    builder: LowLevelIRBuilder, ptr: Value, item_type: RType, *, can_borrow: bool = False
+) -> Value:
+    """Load a vec item from storage, converting nested vec slots to RVec values."""
+    return builder.load_mem(ptr, item_type, borrow=can_borrow)
+
+
+def vec_set_mem_item(
+    builder: LowLevelIRBuilder, ptr: Value, item_type: RType, item: Value
+) -> None:
+    """Store a vec item, converting RVec values to nested storage items."""
+    builder.set_mem(ptr, item_type, item)
+
+
 def vec_check_and_adjust_index(
     builder: LowLevelIRBuilder, lenv: Value, index: Value, line: int
 ) -> Value:
@@ -324,7 +333,7 @@ def vec_get_item_unsafe(
     index = as_platform_int(builder, index, line)
     vtype = base.type
     item_addr = vec_item_ptr(builder, base, index)
-    result = builder.load_mem(item_addr, vtype.item_type, borrow=can_borrow)
+    result = vec_load_mem_item(builder, item_addr, vtype.item_type, can_borrow=can_borrow)
     builder.keep_alives.append(base)
     return result
 
@@ -343,9 +352,9 @@ def vec_set_item(
     if item_type.is_refcounted:
         # Read an unborrowed reference to cause a decref to be
         # generated for the old item.
-        old_item = builder.load_mem(item_addr, item_type, borrow=True)
+        old_item = vec_load_mem_item(builder, item_addr, item_type, can_borrow=True)
         builder.add(DecRef(old_item))
-    builder.set_mem(item_addr, item_type, item)
+    vec_set_mem_item(builder, item_addr, item_type, item)
     builder.keep_alive([base], line)
 
 
@@ -358,32 +367,23 @@ def vec_init_item_unsafe(
     item_addr = vec_item_ptr(builder, base, index)
     item_type = vtype.item_type
     item = builder.coerce(item, item_type, line)
-    builder.set_mem(item_addr, item_type, item)
+    vec_set_mem_item(builder, item_addr, item_type, item)
     builder.keep_alive([base], line)
 
 
 def convert_to_t_ext_item(builder: LowLevelIRBuilder, item: Value) -> Value:
     vec_len = builder.add(GetElement(item, "len"))
-    vec_buf = builder.add(GetElement(item, "buf"))
+    vec_items = builder.add(GetElement(item, "items"))
     temp = builder.add(SetElement(Undef(VecNestedBufItem), "len", vec_len))
-    return builder.add(SetElement(temp, "buf", vec_buf))
+    return builder.add(SetElement(temp, "items", vec_items))
 
 
 def convert_from_t_ext_item(builder: LowLevelIRBuilder, item: Value, vec_type: RVec) -> Value:
-    """Convert a value of type VecNestedBufItem to the corresponding RVec value."""
-    api_name = vec_api_by_item_type.get(vec_type.item_type)
-    if api_name is not None:
-        name = f"{api_name}.convert_from_nested"
-    elif isinstance(vec_type.item_type, RVec):
-        name = "VecNestedApi.convert_from_nested"
-    else:
-        name = "VecTApi.convert_from_nested"
-
-    return builder.add(
-        CallC(
-            name, [item], vec_type, steals=[True], is_borrowed=False, error_kind=ERR_NEVER, line=-1
-        )
-    )
+    """Convert an owned VecNestedBufItem to the corresponding RVec value."""
+    vec_len = builder.add(GetElement(item, "len"))
+    vec_items = builder.add(GetElement(item, "items"))
+    temp = builder.add(SetElement(Undef(vec_type), "len", vec_len))
+    return builder.add(SetElement(temp, "items", vec_items))
 
 
 def vec_item_type(builder: LowLevelIRBuilder, item_type: RType, line: int) -> Value:
@@ -551,7 +551,7 @@ def vec_contains(builder: LowLevelIRBuilder, vec: Value, target: Value, line: in
     for_loop = builder.begin_for(
         items_start, items_end, Integer(step, c_pyssize_t_rprimitive), signed=False
     )
-    item = builder.load_mem(for_loop.index, item_type, borrow=True)
+    item = vec_load_mem_item(builder, for_loop.index, item_type, can_borrow=True)
     comp = builder.binary_op(item, target, "==", line)
     false = BasicBlock()
     builder.add(Branch(comp, true, false, Branch.BOOL))