Use PyList instead of Vec<T>

VirxEC · Dec 25, 2024 · cdea642 · cdea642
1 parent e01d3a2
commit cdea642
Show file tree

Hide file tree

Showing 6 changed files with 87 additions and 25 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "rlbot-flatbuffers-py"
-version = "0.11.4"
+version = "0.11.5"
 edition = "2021"
 description = "A Python module implemented in Rust for serializing and deserializing RLBot's flatbuffers"
 repository = "https://github.com/VirxEC/rlbot_flatbuffers_py"

diff --git a/codegen/structs.rs b/codegen/structs.rs
@@ -105,7 +105,7 @@ impl StructBindGenerator {
         file_contents.push(Cow::Borrowed(if is_frozen {
             "use pyo3::{pyclass, pymethods, types::PyBytes, Bound, PyResult, Python};"
         } else {
-            "use pyo3::{pyclass, pymethods, types::PyBytes, Bound, Py, PyResult, Python};"
+            "use pyo3::{pyclass, pymethods, types::{PyAnyMethods, PyBytes, PyList, PyListMethods}, Bound, Py, PyResult, Python};"
         }));
 
         file_contents.push(Cow::Borrowed(""));
@@ -306,6 +306,10 @@ impl StructBindGenerator {
                     needs_python = true;
                     format!("{variable_name}=None")
                 }
+                RustType::Vec(InnerVecType::Custom(_)) if !self.is_frozen => {
+                    needs_python = true;
+                    format!("{variable_name}=Vec::new()")
+                }
                 _ => {
                     format!("{variable_name}=Default::default()")
                 }
@@ -409,15 +413,18 @@ impl StructBindGenerator {
                     }
                 }
                 RustType::Box(inner_type) | RustType::Custom(inner_type) if !self.is_frozen => {
+                    write_fmt!(self, "            {variable_name}: {variable_name}.unwrap_or_else(|| super::{inner_type}::py_default(py)),",);
+                }
+                RustType::Vec(InnerVecType::U8) => {
                     write_fmt!(
                         self,
-                        "            {variable_name}: {variable_name}.unwrap_or_else(|| super::{inner_type}::py_default(py)),",
+                        "            {variable_name}: {variable_name}.unwrap_or_else(|| PyBytes::new(py, &[]).unbind()),"
                     );
                 }
-                RustType::Vec(InnerVecType::U8) => {
+                RustType::Vec(InnerVecType::Custom(_)) if !self.is_frozen => {
                     write_fmt!(
                         self,
-                        "            {variable_name}: {variable_name}.unwrap_or_else(|| PyBytes::new(py, &[]).unbind()),"
+                        "            {variable_name}: PyList::new(py, {variable_name}).unwrap().unbind(),"
                     );
                 }
                 _ => write_fmt!(self, "            {variable_name},"),
@@ -490,16 +497,17 @@ impl StructBindGenerator {
                             write_str!(self, "                .iter()");
                             write_str!(self, "                .map(ToString::to_string)");
                         }
-                        InnerVecType::Custom(_) => {
+                        InnerVecType::Custom(type_name) => {
+                            if !self.is_frozen {
+                                write_str!(self, "                .bind_borrowed(py)");
+                            }
+
                             write_str!(self, "                .iter()");
-                            write_str!(
-                                self,
-                                if self.is_frozen {
-                                    "                .map(|x| x.__repr__(py))"
-                                } else {
-                                    "                .map(|x| x.borrow(py).__repr__(py))"
-                                }
-                            );
+                            if self.is_frozen {
+                                write_str!(self, "                .map(|x| x.__repr__(py))");
+                            } else {
+                                write_fmt!(self, "                .map(|x| x.downcast_into::<super::{type_name}>().unwrap().borrow().__repr__(py))");
+                            }
                         }
                     }
                     write_str!(self, "                .collect::<Vec<String>>()");
@@ -759,7 +767,7 @@ impl Generator for StructBindGenerator {
                     if self.is_frozen {
                         format!("Vec<super::{inner_type}>")
                     } else {
-                        format!("Vec<Py<super::{inner_type}>>")
+                        String::from("Py<PyList>")
                     }
                 }
                 RustType::Box(inner_type) => {
@@ -815,6 +823,9 @@ impl Generator for StructBindGenerator {
 
             let end = match &variable_info.rust_type {
                 RustType::Vec(InnerVecType::U8) => Cow::Borrowed("PyBytes::new(py, &[]).unbind()"),
+                RustType::Vec(InnerVecType::Custom(_)) if !self.is_frozen => {
+                    Cow::Borrowed("PyList::empty(py).unbind()")
+                }
                 RustType::Vec(_) => Cow::Borrowed("Vec::new()"),
                 RustType::Option(_, _) => Cow::Borrowed("None"),
                 RustType::Union(inner_type)
@@ -890,7 +901,7 @@ impl Generator for StructBindGenerator {
                 RustType::Vec(InnerVecType::String | InnerVecType::Base(_)) => {
                     write_fmt!(self, "            {variable_name}: flat_t.{variable_name},")
                 }
-                RustType::Vec(InnerVecType::Custom(_)) => {
+                RustType::Vec(InnerVecType::Custom(type_name)) => {
                     if self.is_frozen {
                         let map_out = if variable_info.frozen_needs_py {
                             "|x| x.into_gil(py)"
@@ -905,7 +916,7 @@ impl Generator for StructBindGenerator {
                     } else {
                         write_fmt!(
                             self,
-                            "            {variable_name}: flat_t.{variable_name}.into_iter().map(|x| crate::into_py_from(py, x)).collect(),",
+                            "            {variable_name}: PyList::new(py, flat_t.{variable_name}.into_iter().map(|x| crate::into_py_from::<_, super::{type_name}>(py, x))).unwrap().unbind(),",
                         )
                     }
                 }
@@ -1041,7 +1052,7 @@ impl Generator for StructBindGenerator {
                     } else {
                         write_fmt!(
                             self,
-                            "            {variable_name}: py_type.{variable_name}.iter().map(|x| crate::from_py_into(py, x)).collect(),",
+                            "            {variable_name}: py_type.{variable_name}.bind_borrowed(py).iter().map(|x| crate::from_pyany_into(py, x)).collect(),",
                         )
                     }
                 }
@@ -1052,7 +1063,10 @@ impl Generator for StructBindGenerator {
                         "crate::from_py_into(py, x)"
                     };
 
-                    write_fmt!(self, "            {variable_name}: py_type.{variable_name}.as_ref().map(|x| Box::new({inner})),");
+                    write_fmt!(
+                        self,
+                        "            {variable_name}: py_type.{variable_name}.as_ref().map(|x| Box::new({inner})),"
+                    );
                 }
                 RustType::Option(InnerOptionType::String, _) => {
                     write_fmt!(

diff --git a/codegen/unions.rs b/codegen/unions.rs
@@ -139,9 +139,11 @@ impl UnionBindGenerator {
             let variable_name = variable_info.name.as_str();
 
             if variable_info.value.is_some() {
-                write_fmt!(self,
+                write_fmt!(
+                    self,
                     "            Some({}Union::{variable_name}(item)) => format!(\"{}({{}})\", item{borrow_str}.__repr__(py)),",
-                    self.struct_name, self.struct_name
+                    self.struct_name,
+                    self.struct_name
                 );
             } else {
                 write_fmt!(

diff --git a/pybench.py b/pybench.py
@@ -39,7 +39,7 @@ def test_ballpred():
 
     times = []
 
-    ballPred = flat.BallPrediction([flat.PredictionSlice(1) for _ in range(120 * 8)])
+    ballPred = flat.BallPrediction([flat.PredictionSlice(1) for _ in range(120 * 10)])
 
     print(len(ballPred.pack()))
 
@@ -72,11 +72,35 @@ def find_slice_at_time(ball_prediction: flat.BallPrediction, game_time: float):
 
 
 def test_loop():
-    times = []
-
     ballPred = flat.BallPrediction([flat.PredictionSlice(1) for _ in range(120 * 6)])
 
+    start = time_ns()
     for _ in range(100):
+
+        li = []
+        for t in range(1, 301):
+            ball_in_future = find_slice_at_time(ballPred, t / 60)
+            li.append(ball_in_future)
+    print(f"Total time: {(time_ns() - start) / 1_000_000:.3f}ms")
+
+    times = []
+    for _ in range(50_000):
+        start = time_ns()
+
+        li = []
+        for t in range(1, 301):
+            ball_in_future = find_slice_at_time(ballPred, t / 60)
+            li.append(ball_in_future)
+
+        times.append(time_ns() - start)
+
+    print(f"Total time: {sum(times) / 1_000_000_000:.3f}s")
+    avg_time_ns = sum(times) / len(times)
+    print(f"Average time per: {avg_time_ns / 1000:.1f}us")
+    print(f"Minimum time per: {min(times) / 1000:.1f}us")
+
+    times = []
+    for _ in range(50_000):
         start = time_ns()
 
         li = [find_slice_at_time(ballPred, t / 60) for t in range(1, 301)]
@@ -88,6 +112,19 @@ def test_loop():
     print(f"Average time per: {avg_time_ns / 1000:.1f}us")
     print(f"Minimum time per: {min(times) / 1000:.1f}us")
 
+    times = []
+    for _ in range(1_000_000):
+        start = time_ns()
+
+        li = list(ballPred.slices[1:602:2])
+
+        times.append(time_ns() - start)
+
+    print(f"Total time: {sum(times) / 1_000_000_000:.3f}s")
+    avg_time_ns = sum(times) / len(times)
+    print(f"Average time per: {avg_time_ns / 1000:.1f}us")
+    print(f"Minimum time per: {min(times) / 1000:.1f}us")
+
 
 if __name__ == "__main__":
     test_gtp()

diff --git a/src/lib.rs b/src/lib.rs
@@ -80,6 +80,15 @@ where
     (&*obj.borrow(py)).into_gil(py)
 }
 
+#[inline(never)]
+fn from_pyany_into<T, U>(py: Python, obj: Bound<PyAny>) -> U
+where
+    T: PyClass,
+    U: for<'a> FromGil<&'a T>,
+{
+    (&*obj.downcast_into::<T>().unwrap().borrow()).into_gil(py)
+}
+
 pub trait PyDefault: Sized + PyClass {
     fn py_default(py: Python) -> Py<Self>;
 }