[mypyc] Add experimental C extension librt.vecs (part 1/2) (#20653)

The extension defines the `vec[t]` type and associated functions. `vec`
is a sequence type optimized for use in compiled code. `vec` values are
immutable and represented internally as a struct with a length and
buffer fields. Only the buffer object is mutable. This allows very fast
access to the length field in compiled code, but any operations that
change the length of a `vec` value must return a modified `vec` value
(e.g. `append`).

Related issue: mypyc/mypyc#840.

This PR doesn't define any mypyc primitives -- only use via generic
Python operations is supported at this point. I have a local branch that
implements efficient mypyc primitives for vecs, and there are
significant performance gains over `list` objects in benchmarks.

This PR only adds support for item types `i64`, `i32`, `i16`, `u8`,
`float` and `bool` (e.g. `vec[i32]`). A follow-up PR will add support
for reference item types (e.g. `vec[str]`) and nested vecs.

See the comment at the top of `mypyc/lib-rt/vecs/librt_vecs.c` for a
more detailed description. It also documents some implementation details
of reference item types and nested vec types (which are not included yet
in this PR).

Currently a fairly minimal API is supported. This example shows some of
the supported operations:
```
from librt.vecs import vec, append, remove, pop

v = vec[i32]([2, 3])
v[1] = v[0] + 2
v = append(v, 4)
v = remove(v, 2)
v, x = pop(v)
```
Additionally, read-only slicing and equality is supported. The `in`
operator and iteration only work through the sequence API, but I will
add proper support for these later on.

Author: JukkaL

mypy/typeshed/stubs/librt/librt/vecs.pyi

@@ -0,0 +1,21 @@
+from typing import TypeVar, Generic, Iterable, Iterator, overload
+from mypy_extensions import i64
+
+T = TypeVar("T")
+
+class vec(Generic[T]):
+    @overload
+    def __init__(self) -> None: ...
+    @overload
+    def __init__(self, items: Iterable[T], /) -> None: ...
+    def __len__(self) -> i64: ...
+    @overload
+    def __getitem__(self, i: i64, /) -> T: ...
+    @overload
+    def __getitem__(self, i: slice, /) -> vec[T]: ...
+    def __setitem__(self, i: i64, o: T, /) -> None: ...
+    def __iter__(self) -> Iterator[T]: ...
+
+def append(v: vec[T], o: T, /) -> vec[T]: ...
+def remove(v: vec[T], o: T, /) -> vec[T]: ...
+def pop(v: vec[T], i: i64 = -1, /) -> tuple[vec[T], T]: ...

mypyc/build.py

@@ -104,6 +104,20 @@ class ModDesc(NamedTuple):
         ],
         ["base64"],
     ),
+    ModDesc(
+        "librt.vecs",
+        [
+            "vecs/librt_vecs.c",
+            "vecs/vec_i64.c",
+            "vecs/vec_i32.c",
+            "vecs/vec_i16.c",
+            "vecs/vec_u8.c",
+            "vecs/vec_float.c",
+            "vecs/vec_bool.c",
+        ],
+        ["vecs/librt_vecs.h", "vecs/vec_template.c"],
+        ["vecs"],
+    ),
 ]
 
 try:

mypyc/ir/deps.py

@@ -48,6 +48,7 @@ def get_header(self) -> str:
 
 LIBRT_STRINGS: Final = Capsule("librt.strings")
 LIBRT_BASE64: Final = Capsule("librt.base64")
+LIBRT_VECS: Final = Capsule("librt.vecs")
 
 BYTES_EXTRA_OPS: Final = SourceDep("bytes_extra_ops.c")
 BYTES_WRITER_EXTRA_OPS: Final = SourceDep("byteswriter_extra_ops.c")

mypyc/lib-rt/setup.py

@@ -132,5 +132,19 @@ def run(self) -> None:
                 include_dirs=[".", "base64"],
                 extra_compile_args=cflags,
             ),
+            Extension(
+                "librt.vecs",
+                [
+                    "vecs/librt_vecs.c",
+                    "vecs/vec_i64.c",
+                    "vecs/vec_i32.c",
+                    "vecs/vec_i16.c",
+                    "vecs/vec_u8.c",
+                    "vecs/vec_float.c",
+                    "vecs/vec_bool.c",
+                ],
+                include_dirs=[".", "vecs"],
+                extra_compile_args=cflags,
+            ),
         ]
     )