From 39be1861642ccb9856e1e1191f5bf5769cd4f169 Mon Sep 17 00:00:00 2001
From: Uli Schlachter <psychon@znc.in>
Date: Sat, 3 Feb 2024 13:42:23 +0100
Subject: [PATCH] Optimise atom_manager! implementation

winit uses atom_manager! with 59 atom names. Due to some sub-optimal
code generation, this causes the generated new() function to end up with
a size of 111.6 KiB, or almost 5% of the .text size of a winit example.

The issue is that each "intern_atom()?" can cause an early exit and
needs to call drop() of the already generated Cookies. This leads to
some quadratic growth.

To work around this, in this commit I change the implementation of
new(). It now iterates over the atom names, using .map() to send the
InternAtom requests. Finally, collect() is then used to change this
Iterator<Item=Result<Cookie, Error>> into Result<Vec<Cookie>, Error>.
This even preserves the early-exit behaviour in case of errors.

The reply() function is then changed to turn the Vec into an iterator,
using next() to pull out the items. This relies on the evaluation of
this code to be well-defined (which e.g. wouldn't be guaranteed in C).
Luckily, Rust provides the needed guarantee:
https://github.com/rust-lang/reference/pull/888

Fixes: https://github.com/psychon/x11rb/issues/912
Signed-off-by: Uli Schlachter <psychon@znc.in>
---
 x11rb/src/x11_utils.rs | 35 ++++++++++++++++++-----------------
 1 file changed, 18 insertions(+), 17 deletions(-)

diff --git a/x11rb/src/x11_utils.rs b/x11rb/src/x11_utils.rs
index 55b4980f..e339c84a 100644
--- a/x11rb/src/x11_utils.rs
+++ b/x11rb/src/x11_utils.rs
@@ -44,17 +44,19 @@ pub use x11rb_protocol::x11_utils::{
 /// #[derive(Debug)]
 /// /// A handle to a response from the X11 server.
 /// struct AtomCollectionCookie<'c, C: ConnectionExt> {
-///     phantom: std::marker::PhantomData<&'c C>,
-///     _NET_WM_NAME: Cookie<'c, C, InternAtomReply>,
-///     _NET_WM_ICON: Cookie<'c, C, InternAtomReply>,
-///     ATOM_WITH_SPACES: Cookie<'c, C, InternAtomReply>,
-///     WHATEVER: Cookie<'c, C, InternAtomReply>,
+///     // please treat the actual members as private
+///     # phantom: std::marker::PhantomData<&'c C>,
+///     # _NET_WM_NAME: Cookie<'c, C, InternAtomReply>,
+///     # _NET_WM_ICON: Cookie<'c, C, InternAtomReply>,
+///     # ATOM_WITH_SPACES: Cookie<'c, C, InternAtomReply>,
+///     # WHATEVER: Cookie<'c, C, InternAtomReply>,
 /// }
 ///
 /// impl AtomCollection {
 ///     pub fn new<C: ConnectionExt>(
 ///         conn: &C,
 ///     ) -> Result<AtomCollectionCookie<'_, C>, ConnectionError> {
+///         // This is just an example for readability; the actual code is more efficient.
 ///         Ok(AtomCollectionCookie {
 ///             phantom: std::marker::PhantomData,
 ///             _NET_WM_NAME: conn.intern_atom(false, b"_NET_WM_NAME")?,
@@ -70,6 +72,7 @@ pub use x11rb_protocol::x11_utils::{
 ///     C: ConnectionExt,
 /// {
 ///     pub fn reply(self) -> Result<AtomCollection, ReplyError> {
+///         // This is just an example for readability; the actual code is different.
 ///         Ok(AtomCollection {
 ///             _NET_WM_NAME: self._NET_WM_NAME.reply()?.atom,
 ///             _NET_WM_ICON: self._NET_WM_ICON.reply()?.atom,
@@ -94,10 +97,8 @@ macro_rules! atom_manager {
         #[derive(Debug)]
         $(#[$cookie_meta])*
         $vis struct $cookie_name<'a, C: $crate::protocol::xproto::ConnectionExt> {
-            phantom: std::marker::PhantomData<&'a C>,
-            $(
-                $field_name: $crate::cookie::Cookie<'a, C, $crate::protocol::xproto::InternAtomReply>,
-            )*
+            __private_phantom: std::marker::PhantomData<&'a C>,
+            __private_cookies: Vec<$crate::cookie::Cookie<'a, C, $crate::protocol::xproto::InternAtomReply>>,
         }
 
         // Replies
@@ -114,23 +115,23 @@ macro_rules! atom_manager {
             $vis fn new<C: $crate::protocol::xproto::ConnectionExt>(
                 _conn: &C,
             ) -> ::std::result::Result<$cookie_name<'_, C>, $crate::errors::ConnectionError> {
+                let names = [
+                    $($crate::__atom_manager_atom_value!($field_name$(: $atom_value)?),)*
+                ];
+                let cookies: Result<Vec<_>, _> = names.into_iter().map(|name| _conn.intern_atom(false, name)).collect();
                 Ok($cookie_name {
-                    phantom: std::marker::PhantomData,
-                    $(
-                        $field_name: _conn.intern_atom(
-                            false,
-                            $crate::__atom_manager_atom_value!($field_name$(: $atom_value)?),
-                        )?,
-                    )*
+                    __private_phantom: std::marker::PhantomData,
+                    __private_cookies: cookies?,
                 })
             }
         }
 
         impl<'a, C: $crate::protocol::xproto::ConnectionExt> $cookie_name<'a, C> {
             $vis fn reply(self) -> ::std::result::Result<$struct_name, $crate::errors::ReplyError> {
+                let mut replies = self.__private_cookies.into_iter();
                 Ok($struct_name {
                     $(
-                        $field_name: self.$field_name.reply()?.atom,
+                        $field_name: replies.next().expect("new() should have constructed a Vec of the correct size").reply()?.atom,
                     )*
                 })
             }