Don't allocate for PyNativeFunc when it's just a fn pointer

coolreader18 · coolreader18 · commit e2d7ab30aaf4 · 2020-04-02T13:57:08.000-05:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/derive/src/pyclass.rs b/derive/src/pyclass.rs
@@ -458,7 +458,7 @@ fn extract_impl_items(mut items: Vec<ItemSig>) -> Result<TokenStream2, Diagnosti
             let transform = if vec!["new", "call"].contains(&slot_ident.to_string().as_str()) {
                 quote! { ::rustpython_vm::function::IntoPyNativeFunc::into_func }
             } else {
-                quote! { Box::new }
+                quote! { ::rustpython_vm::__exports::smallbox! }
             };
             let into_func = quote_spanned! {item_ident.span()=>
                 #transform(Self::#item_ident)
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
@@ -72,6 +72,7 @@ is-macro = "0.1"
 result-like = "^0.2.1"
 foreign-types = "0.3"
 num_enum = "0.4"
+smallbox = "0.8"
 
 flame = { version = "0.2", optional = true }
 flamer = { version = "0.3", optional = true }
diff --git a/vm/src/function.rs b/vm/src/function.rs
@@ -4,6 +4,7 @@ use std::ops::RangeInclusive;
 
 use indexmap::IndexMap;
 use result_like::impl_option_like;
+use smallbox::{smallbox, space::S1, SmallBox};
 
 use crate::exceptions::PyBaseExceptionRef;
 use crate::obj::objtuple::PyTuple;
@@ -453,8 +454,11 @@ tuple_from_py_func_args!(A, B, C, D);
 tuple_from_py_func_args!(A, B, C, D, E);
 tuple_from_py_func_args!(A, B, C, D, E, F);
 
+/// A container that can hold a `dyn Fn*` trait object, but doesn't allocate if it's only a fn() pointer
+pub type FunctionBox<T> = SmallBox<T, S1>;
+
 /// A built-in Python function.
-pub type PyNativeFunc = Box<dyn Fn(&VirtualMachine, PyFuncArgs) -> PyResult + 'static>;
+pub type PyNativeFunc = FunctionBox<dyn Fn(&VirtualMachine, PyFuncArgs) -> PyResult + 'static>;
 
 /// Implemented by types that are or can generate built-in functions.
 ///
@@ -479,7 +483,7 @@ where
     F: Fn(&VirtualMachine, PyFuncArgs) -> PyResult + 'static,
 {
     fn into_func(self) -> PyNativeFunc {
-        Box::new(self)
+        smallbox!(self)
     }
 }
 
@@ -499,7 +503,7 @@ macro_rules! into_py_native_func_tuple {
             R: IntoPyObject,
         {
             fn into_func(self) -> PyNativeFunc {
-                Box::new(move |vm, args| {
+                smallbox!(move |vm: &VirtualMachine, args: PyFuncArgs| {
                     let ($($n,)*) = args.bind::<($($T,)*)>(vm)?;
 
                     (self)($($n,)* vm).into_pyobject(vm)
@@ -515,7 +519,7 @@ macro_rules! into_py_native_func_tuple {
             R: IntoPyObject,
         {
             fn into_func(self) -> PyNativeFunc {
-                Box::new(move |vm, args| {
+                smallbox!(move |vm: &VirtualMachine, args: PyFuncArgs| {
                     let (zelf, $($n,)*) = args.bind::<(PyRef<S>, $($T,)*)>(vm)?;
 
                     (self)(&zelf, $($n,)* vm).into_pyobject(vm)
@@ -597,3 +601,15 @@ pub fn single_or_tuple_any<T: PyValue, F: Fn(PyRef<T>) -> PyResult<bool>>(
     };
     checker.check(obj)
 }
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_functionbox_noalloc() {
+        fn py_func(_b: bool, _vm: &crate::VirtualMachine) -> i32 {
+            1
+        }
+        let f = super::IntoPyNativeFunc::into_func(py_func);
+        assert!(!f.is_heap());
+    }
+}
diff --git a/vm/src/lib.rs b/vm/src/lib.rs
@@ -88,4 +88,5 @@ pub use rustpython_bytecode::*;
 #[doc(hidden)]
 pub mod __exports {
     pub use maplit::hashmap;
+    pub use smallbox::smallbox;
 }
diff --git a/vm/src/obj/objgetset.rs b/vm/src/obj/objgetset.rs
@@ -2,16 +2,17 @@
 
 */
 use super::objtype::PyClassRef;
-use crate::function::{OptionalArg, OwnedParam, RefParam};
+use crate::function::{FunctionBox, OptionalArg, OwnedParam, RefParam};
 use crate::pyobject::{
     IntoPyObject, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
     TypeProtocol,
 };
 use crate::slots::SlotDescriptor;
 use crate::vm::VirtualMachine;
 
-pub type PyGetterFunc = Box<dyn Fn(&VirtualMachine, PyObjectRef) -> PyResult>;
-pub type PySetterFunc = Box<dyn Fn(&VirtualMachine, PyObjectRef, PyObjectRef) -> PyResult<()>>;
+pub type PyGetterFunc = FunctionBox<dyn Fn(&VirtualMachine, PyObjectRef) -> PyResult>;
+pub type PySetterFunc =
+    FunctionBox<dyn Fn(&VirtualMachine, PyObjectRef, PyObjectRef) -> PyResult<()>>;
 
 pub trait IntoPyGetterFunc<T> {
     fn into_getter(self) -> PyGetterFunc;
@@ -24,7 +25,7 @@ where
     R: IntoPyObject,
 {
     fn into_getter(self) -> PyGetterFunc {
-        Box::new(move |vm, obj| {
+        smallbox::smallbox!(move |vm: &VirtualMachine, obj| {
             let obj = T::try_from_object(vm, obj)?;
             (self)(obj, vm).into_pyobject(vm)
         })
@@ -38,7 +39,7 @@ where
     R: IntoPyObject,
 {
     fn into_getter(self) -> PyGetterFunc {
-        Box::new(move |vm, obj| {
+        smallbox::smallbox!(move |vm: &VirtualMachine, obj| {
             let zelf = PyRef::<S>::try_from_object(vm, obj)?;
             (self)(&zelf, vm).into_pyobject(vm)
         })
@@ -95,7 +96,7 @@ where
     R: IntoPyNoResult,
 {
     fn into_setter(self) -> PySetterFunc {
-        Box::new(move |vm, obj, value| {
+        smallbox::smallbox!(move |vm: &VirtualMachine, obj, value| {
             let obj = T::try_from_object(vm, obj)?;
             let value = V::try_from_object(vm, value)?;
             (self)(obj, value, vm).into_noresult()
@@ -111,7 +112,7 @@ where
     R: IntoPyNoResult,
 {
     fn into_setter(self) -> PySetterFunc {
-        Box::new(move |vm, obj, value| {
+        smallbox::smallbox!(move |vm: &VirtualMachine, obj, value| {
             let zelf = PyRef::<S>::try_from_object(vm, obj)?;
             let value = V::try_from_object(vm, value)?;
             (self)(&zelf, value, vm).into_noresult()
diff --git a/vm/src/slots.rs b/vm/src/slots.rs
@@ -1,4 +1,4 @@
-use crate::function::{OptionalArg, PyFuncArgs, PyNativeFunc};
+use crate::function::{FunctionBox, OptionalArg, PyFuncArgs, PyNativeFunc};
 use crate::pyobject::{IdProtocol, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject};
 use crate::VirtualMachine;
 
@@ -44,7 +44,7 @@ pub trait SlotCall: PyValue {
     fn call(&self, args: PyFuncArgs, vm: &VirtualMachine) -> PyResult;
 }
 
-pub type PyDescrGetFunc = Box<
+pub type PyDescrGetFunc = FunctionBox<
     dyn Fn(&VirtualMachine, PyObjectRef, Option<PyObjectRef>, OptionalArg<PyObjectRef>) -> PyResult,
 >;
 

Original file line number	Diff line number	Diff line change
`@@ -88,4 +88,5 @@ pub use rustpython_bytecode::*;`
`88`	`88`	`#[doc(hidden)]`
`89`	`89`	`pub mod __exports {`
`90`	`90`	`pub use maplit::hashmap;`
	`91`	`+ pub use smallbox::smallbox;`
`91`	`92`	`}`