Use static encoding strings in define_class! methods

For optimization, and as a first step towards
#604

Author: madsmtm

crates/objc2/CHANGELOG.md

@@ -12,6 +12,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 * Added support for `&CStr` in method argument and return types.
 * Implement `Encode` on `Result<(), T>` when `T: OptionEncode`.
 * Allow returning `Result<T, Retained<NSError>>` from `define_class!`.
+* Optimized adding `define_class!` methods to avoid a few string allocations.
 
 ## Changed
 * **BREAKING**: Changed syntax for `define_class!` ivars (instance variables).

crates/objc2/src/__macros/define_class/checks.rs

@@ -6,6 +6,7 @@ use core::panic::{RefUnwindSafe, UnwindSafe};
 #[cfg(debug_assertions)]
 use std::collections::HashSet;
 
+use crate::encode::{EncodeArguments, EncodeReturn};
 use crate::runtime::{AnyClass, ClassBuilder, MethodImplementation, Sel};
 #[cfg(debug_assertions)]
 use crate::runtime::{AnyProtocol, MethodDescription};
@@ -215,21 +216,43 @@ impl<T: DefinedClass> ClassBuilderHelper<T> {
 
     // Addition: This restricts to callee `T`
     #[inline]
-    pub unsafe fn add_method<F>(&mut self, sel: Sel, func: F)
+    pub unsafe fn add_method<F>(&mut self, sel: Sel, func: F, encoding: &CStr)
     where
         F: MethodImplementation<Callee = T>,
     {
+        if cfg!(all(
+            debug_assertions,
+            not(feature = "disable-encoding-assertions")
+        )) {
+            self.builder
+                .verify_method(sel, F::Arguments::ENCODINGS, &F::Return::ENCODING_RETURN);
+        }
+
         // SAFETY: Checked by caller
-        unsafe { self.builder.add_method(sel, func) }
+        unsafe { self.builder.add_method_inner(sel, func.__imp(), encoding) }
     }
 
     #[inline]
-    pub unsafe fn add_class_method<F>(&mut self, sel: Sel, func: F)
+    pub unsafe fn add_class_method<F>(&mut self, sel: Sel, func: F, encoding: &CStr)
     where
         F: MethodImplementation<Callee = AnyClass>,
     {
+        if cfg!(all(
+            debug_assertions,
+            not(feature = "disable-encoding-assertions")
+        )) {
+            self.builder.verify_class_method(
+                sel,
+                F::Arguments::ENCODINGS,
+                &F::Return::ENCODING_RETURN,
+            );
+        }
+
         // SAFETY: Checked by caller
-        unsafe { self.builder.add_class_method(sel, func) }
+        unsafe {
+            self.builder
+                .add_class_method_inner(sel, func.__imp(), encoding)
+        }
     }
 }
 
@@ -259,7 +282,7 @@ pub struct ClassProtocolMethodsBuilder<'a, T: ?Sized> {
 impl<T: DefinedClass> ClassProtocolMethodsBuilder<'_, T> {
     // Addition: This restricts to callee `T`
     #[inline]
-    pub unsafe fn add_method<F>(&mut self, sel: Sel, func: F)
+    pub unsafe fn add_method<F>(&mut self, sel: Sel, func: F, encoding: &CStr)
     where
         F: MethodImplementation<Callee = T>,
     {
@@ -279,7 +302,7 @@ impl<T: DefinedClass> ClassProtocolMethodsBuilder<'_, T> {
         }
 
         // SAFETY: Checked by caller
-        unsafe { self.builder.add_method(sel, func) };
+        unsafe { self.builder.add_method(sel, func, encoding) };
 
         #[cfg(debug_assertions)]
         if !self.registered_instance_methods.insert(sel) {
@@ -288,7 +311,7 @@ impl<T: DefinedClass> ClassProtocolMethodsBuilder<'_, T> {
     }
 
     #[inline]
-    pub unsafe fn add_class_method<F>(&mut self, sel: Sel, func: F)
+    pub unsafe fn add_class_method<F>(&mut self, sel: Sel, func: F, encoding: &CStr)
     where
         F: MethodImplementation<Callee = AnyClass>,
     {
@@ -308,7 +331,7 @@ impl<T: DefinedClass> ClassProtocolMethodsBuilder<'_, T> {
         }
 
         // SAFETY: Checked by caller
-        unsafe { self.builder.add_class_method(sel, func) };
+        unsafe { self.builder.add_class_method(sel, func, encoding) };
 
         #[cfg(debug_assertions)]
         if !self.registered_class_methods.insert(sel) {

crates/objc2/src/__macros/define_class/encoding.rs

@@ -0,0 +1,91 @@
+//! Helpers to construct encoding strings for methods at `const`-time.
+use objc2_encode::Encoding;
+
+pub const fn method_encoding_str_len(ret: &Encoding, args: &[Encoding]) -> usize {
+    let mut len = 0;
+    len += ret.str_len();
+    len += 1; // Encoding::Object
+    len += 1; // Encoding::Sel
+    let mut i = 0;
+    while i < args.len() {
+        len += args[i].str_len();
+        i += 1;
+    }
+    len
+}
+
+pub const fn method_encoding_str_array<const LEN: usize>(
+    ret: &Encoding,
+    args: &[Encoding],
+) -> [u8; LEN] {
+    let mut res: [u8; LEN] = [0; LEN];
+    let mut res_i = 0;
+
+    let mut i = 0;
+    // We use LEN even though it creates an oversized array.
+    let arr = ret.str_array::<LEN>();
+    while i < ret.str_len() {
+        res[res_i] = arr[i];
+        i += 1;
+        res_i += 1;
+    }
+
+    // Encoding::Object
+    res[res_i] = b'@';
+    res_i += 1;
+
+    // Encoding::Sel
+    res[res_i] = b':';
+    res_i += 1;
+
+    let mut i = 0;
+    while i < args.len() {
+        let mut j = 0;
+        // We use LEN even though it creates an oversized array.
+        let arr = args[i].str_array::<LEN>();
+        while j < args[i].str_len() {
+            res[res_i] = arr[j];
+            j += 1;
+            res_i += 1;
+        }
+        i += 1;
+    }
+    res
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::encode::{Encode, EncodeReturn};
+    use crate::runtime::{method_type_encoding, Bool, NSObject};
+
+    #[track_caller]
+    fn check(ret: &Encoding, args: &[Encoding]) {
+        // Compare against reference impl
+        let expected = method_type_encoding(ret, args).into_string().unwrap();
+
+        let len = method_encoding_str_len(ret, args);
+        assert_eq!(expected.len(), len, "incorrect length");
+        assert!(len < 100);
+
+        let arr = method_encoding_str_array::<100>(ret, args);
+        let (s, rest) = arr.split_at(len);
+        let s = str::from_utf8(s).unwrap();
+        assert_eq!(expected, s, "incorrect output: {arr:?}");
+
+        assert!(rest.iter().all(|x| *x == 0), "rest must be zero: {arr:?}");
+    }
+
+    #[test]
+    fn various() {
+        check(&<()>::ENCODING_RETURN, &[]);
+        check(
+            &<*const NSObject>::ENCODING_RETURN,
+            &[<&NSObject>::ENCODING],
+        );
+        check(
+            &<&&&&&&i32>::ENCODING_RETURN,
+            &[<&&&&&&NSObject>::ENCODING, <&&&&&&Bool>::ENCODING],
+        );
+    }
+}

crates/objc2/src/__macros/define_class/ivars.rs

@@ -52,8 +52,11 @@ use core::ffi::CStr;
 use core::mem;
 use core::ptr::{self, NonNull};
 
+use crate::__macros::{method_encoding_str_array, method_encoding_str_len};
 use crate::encode::{Encode, Encoding};
-use crate::runtime::{AnyClass, AnyObject, ClassBuilder, MessageReceiver, Sel};
+use crate::runtime::{
+    AnyClass, AnyObject, ClassBuilder, MessageReceiver, MethodImplementation, Sel,
+};
 use crate::{sel, ClassType, DefinedClass};
 
 /// A type representing the drop flags that may be set for a type.
@@ -144,9 +147,15 @@ where
     // Add dealloc if the class or the ivars need dropping.
     if mem::needs_drop::<T>() || mem::needs_drop::<T::Ivars>() {
         let func: unsafe extern "C-unwind" fn(_, _) = dealloc::<T>;
-        // SAFETY: The function signature is correct, and method contract is
-        // upheld inside `dealloc`.
-        unsafe { builder.add_method(sel!(dealloc), func) };
+
+        const ENCODING_LEN: usize = method_encoding_str_len(&Encoding::Void, &[]);
+        const ENCODING_ARRAY: [u8; ENCODING_LEN + 1] =
+            method_encoding_str_array::<{ ENCODING_LEN + 1 }>(&Encoding::Void, &[]);
+        const ENCODING: &CStr = unsafe { CStr::from_bytes_with_nul_unchecked(&ENCODING_ARRAY) };
+
+        // SAFETY: The function signature is correct, the encoding is correct,
+        // and method contract is upheld inside `dealloc`.
+        unsafe { builder.add_method_inner(sel!(dealloc), func.__imp(), ENCODING) };
     } else {
         // Users should not rely on this omission, it is only an optimization.
     }

crates/objc2/src/__macros/define_class/mod.rs

@@ -1,10 +1,12 @@
 mod checks;
+mod encoding;
 mod ivars;
 mod output_impls;
 mod register_impls;
 mod thunk;
 
 pub use self::checks::*;
+pub use self::encoding::*;
 pub use self::ivars::*;
 pub use self::thunk::*;
 

crates/objc2/src/__macros/define_class/register_impls.rs

@@ -348,12 +348,28 @@ macro_rules! __declare_class_register_thunk {
             // TODO: `<Self $(as $protocol)?>` once we emit on protocols.
         }
 
+        const __ENCODING_LEN: $crate::__macros::usize = $crate::__macros::method_encoding_str_len(
+            &<$for as $crate::__macros::ConvertDefinedFn<'_, __FnMarker, __RetainSemantics, __Kind<'_>>>::RETURN_ENCODING,
+            <$for as $crate::__macros::ConvertDefinedFn<'_, __FnMarker, __RetainSemantics, __Kind<'_>>>::ARGUMENT_ENCODINGS,
+        );
+
+        // SAFETY: The array is guaranteed to be contain no interior NUL
+        // bytes (`verify_name` in `objc2-encode`), and a NUL byte at the end
+        // because we specify `__ENCODING_LEN + 1`.
+        const __ENCODING: &$crate::__macros::CStr = unsafe {
+            $crate::__macros::CStr::from_bytes_with_nul_unchecked(&$crate::__macros::method_encoding_str_array::<{__ENCODING_LEN + 1}>(
+                &<$for as $crate::__macros::ConvertDefinedFn<'_, __FnMarker, __RetainSemantics, __Kind<'_>>>::RETURN_ENCODING,
+                <$for as $crate::__macros::ConvertDefinedFn<'_, __FnMarker, __RetainSemantics, __Kind<'_>>>::ARGUMENT_ENCODINGS,
+            ))
+        };
+
         unsafe {
             $crate::__declare_class_call_builder_method! {
                 ($builder)
                 ($crate::sel!($($sel)*))
                 (<Self as $crate::__macros::ConvertDefinedFn<'_, __FnMarker, __RetainSemantics, __Kind<'_>>>::THUNK)
                 ($($receiver_ty)?)
+                (__ENCODING)
             }
         };
     );
@@ -385,8 +401,9 @@ macro_rules! __declare_class_call_builder_method {
         ($sel:expr)
         ($fnptr:expr)
         ()
+        ($encoding:expr)
     } => {
-        $builder.add_class_method($sel, $fnptr)
+        $builder.add_class_method($sel, $fnptr, $encoding)
     };
 
     // Instance method
@@ -395,8 +412,9 @@ macro_rules! __declare_class_call_builder_method {
         ($sel:expr)
         ($fnptr:expr)
         ($_receiver_ty:ty)
+        ($encoding:expr)
     } => {
-        $builder.add_method($sel, $fnptr)
+        $builder.add_method($sel, $fnptr, $encoding)
     };
 }
 

crates/objc2/src/__macros/define_class/thunk.rs

@@ -20,7 +20,10 @@
 //! few arguments compared to the selector.
 use core::marker::PhantomData;
 
+use objc2_encode::Encoding;
+
 use super::super::{ConvertArgument, ConvertError, KindDefined, RetainSemantics};
+use crate::encode::{EncodeArguments, EncodeReturn};
 use crate::rc::Retained;
 use crate::runtime::{MethodImplementation, Sel};
 use crate::ClassType;
@@ -49,6 +52,22 @@ pub trait ConvertDefinedFn<'f, FnMarker, MethodFamily, Kind>: ClassType {
     /// The generated thunk.
     const THUNK: <Self::Func as FnToExternFn<FnMarker, MethodFamily, Kind, Self>>::ExternFn =
         { Self::Func::THUNK };
+
+    /// Helper for getting the return encoding of the extern function.
+    const RETURN_ENCODING: Encoding = <<<Self::Func as FnToExternFn<
+        FnMarker,
+        MethodFamily,
+        Kind,
+        Self,
+    >>::ExternFn as MethodImplementation>::Return as EncodeReturn>::ENCODING_RETURN;
+
+    /// Helper for getting the argument encodings of the extern function.
+    const ARGUMENT_ENCODINGS: &'static [Encoding] = <<<Self::Func as FnToExternFn<
+        FnMarker,
+        MethodFamily,
+        Kind,
+        Self,
+    >>::ExternFn as MethodImplementation>::Arguments as EncodeArguments>::ENCODINGS;
 }
 
 /// Function pointers that can be converted to `extern "C-unwind"` function

crates/objc2/src/__macros/mod.rs

@@ -47,7 +47,7 @@ pub use core::marker::{PhantomData, Sized};
 pub use core::mem::{size_of, transmute, ManuallyDrop, MaybeUninit};
 pub use core::ops::Deref;
 pub use core::option::Option::{self, None, Some};
-pub use core::primitive::{bool, isize, str, u8};
+pub use core::primitive::{bool, isize, str, u8, usize};
 pub use core::{compile_error, concat, env, module_path, panic, stringify};
 pub use std::sync::Once;
 
@@ -60,8 +60,9 @@ pub use self::convert::{
     ConvertArgument, ConvertArguments, ConvertError, ConvertReturn, TupleExtender,
 };
 pub use self::define_class::{
-    class_c_name, define_class, ClassBuilderHelper, ClassFnKind, ClassFnResultKind,
-    ClassProtocolMethodsBuilder, ConvertDefinedFn, DefinedIvarsHelper, LifetimeAssign, MethodKind,
+    class_c_name, define_class, method_encoding_str_array, method_encoding_str_len,
+    ClassBuilderHelper, ClassFnKind, ClassFnResultKind, ClassProtocolMethodsBuilder,
+    ConvertDefinedFn, DefinedIvarsHelper, FnToExternFn, LifetimeAssign, MethodKind,
     MethodResultKind, ThreadKindAutoTraits,
 };
 pub use self::extern_class::{DoesNotImplDrop, MainThreadOnlyDoesNotImplSendSync, ValidThreadKind};