From 89d45007bd05679cae11200ac7782f250807fbd8 Mon Sep 17 00:00:00 2001
From: NicEastvillage <nicoesterby@gmail.com>
Date: Sun, 22 Jun 2025 20:23:07 +0200
Subject: [PATCH 1/9] Add DesiredStateBuilder

---
 rlbot/examples/atba_agent/main.rs |   2 +
 rlbot/src/lib.rs                  |   1 +
 rlbot/src/state_builder.rs        | 321 ++++++++++++++++++++++++++++++
 3 files changed, 324 insertions(+)
 create mode 100644 rlbot/src/state_builder.rs

diff --git a/rlbot/examples/atba_agent/main.rs b/rlbot/examples/atba_agent/main.rs
index 7ec4ec5..df4f7e6 100644
--- a/rlbot/examples/atba_agent/main.rs
+++ b/rlbot/examples/atba_agent/main.rs
@@ -9,6 +9,8 @@ use rlbot::{
     },
     util::{AgentEnvironment, PacketQueue},
 };
+use rlbot::state_builder::DesiredStateBuilder;
+use rlbot_flat::flat::{DesiredGameStateBuilder, Vector3};
 
 #[allow(dead_code)]
 struct AtbaAgent {
diff --git a/rlbot/src/lib.rs b/rlbot/src/lib.rs
index af435c0..59aee0f 100644
--- a/rlbot/src/lib.rs
+++ b/rlbot/src/lib.rs
@@ -9,6 +9,7 @@ use thiserror::Error;
 
 pub mod agents;
 pub mod util;
+pub mod state_builder;
 
 #[cfg(feature = "glam")]
 pub use rlbot_flat::glam;
diff --git a/rlbot/src/state_builder.rs b/rlbot/src/state_builder.rs
new file mode 100644
index 0000000..26b8728
--- /dev/null
+++ b/rlbot/src/state_builder.rs
@@ -0,0 +1,321 @@
+use rlbot_flat::flat::{DesiredBallState, DesiredCarState, DesiredGameState, DesiredMatchInfo, Rotator, RotatorPartial, Vector3, Vector3Partial};
+
+/// Utility for easy construction of [DesiredGameState]s using builder patterns.
+/// 
+/// Example:
+/// ```example
+/// let dgs: DesiredGameState = DesiredStateBuilder::new()
+///             .car(0, |c| c
+///                 .location(Vector3::default())
+///                 .boost(100.))
+///             .ball(0, |b| b
+///                 .location_z(0.)
+///                 .velocity_z(0.)
+///             )
+///             .build();
+/// ```
+#[derive(Default, Debug, Clone, PartialEq, PartialOrd)]
+pub struct DesiredStateBuilder {
+    state: DesiredGameState
+}
+
+#[allow(dead_code)]
+impl DesiredStateBuilder {
+    pub fn new() -> Self {
+        Self {
+            state: DesiredGameState::default()
+        }
+    }
+
+    /// Modify the desired match info.
+    pub fn match_info(mut self, build: impl FnOnce(DesiredMatchInfoBuilder) -> DesiredMatchInfoBuilder) -> Self {
+        build(DesiredMatchInfoBuilder::new(&mut self.state.match_info.get_or_insert_default()));
+        self
+    }
+
+    /// Modify the desired ball at the given index.
+    pub fn ball(mut self, index: usize, build: impl FnOnce(DesiredBallBuilder) -> DesiredBallBuilder) -> Self {
+        while self.state.ball_states.len() <= index {
+            self.state.ball_states.push(Default::default());
+        }
+        build(DesiredBallBuilder::new(&mut self.state.ball_states[index]));
+        self
+    }
+
+    /// Modify the desired car at the given index.
+    pub fn car(mut self, index: usize, build: impl FnOnce(DesiredCarBuilder) -> DesiredCarBuilder) -> Self {
+        while self.state.car_states.len() <= index {
+            self.state.car_states.push(Default::default());
+        }
+        build(DesiredCarBuilder::new(&mut self.state.car_states[index]));
+        self
+    }
+
+    /// Extract the resulting [DesiredGameState].
+    pub fn build(self) -> DesiredGameState {
+        self.state
+    }
+}
+
+/// Allows for easy modification of a [DesiredMatchInfo]. See [DesiredStateBuilder].
+#[derive(Debug, PartialEq, PartialOrd)]
+pub struct DesiredMatchInfoBuilder<'a> {
+    info: &'a mut DesiredMatchInfo,
+}
+
+#[allow(dead_code)]
+impl<'a> DesiredMatchInfoBuilder<'a> {
+    pub fn new(info: &'a mut DesiredMatchInfo) -> Self {
+        Self {
+            info
+        }
+    }
+
+    /// Set the desired world gravity z.
+    pub fn gravity_z(self, gravity_z: f32) -> Self {
+        self.info.world_gravity_z.get_or_insert_default().val = gravity_z;
+        self
+    }
+
+    /// Set the desired game speed.
+    pub fn game_speed(self, game_speed: f32) -> Self {
+        self.info.game_speed.get_or_insert_default().val = game_speed;
+        self
+    }
+}
+
+/// Allows for easy modification of a [DesiredCarState]. See [DesiredStateBuilder].
+#[derive(Debug, PartialEq, PartialOrd)]
+pub struct DesiredCarBuilder<'a> {
+    car: &'a mut DesiredCarState,
+}
+
+#[allow(dead_code)]
+impl<'a> DesiredCarBuilder<'a> {
+    pub fn new(car: &'a mut DesiredCarState) -> Self {
+        Self {
+            car
+        }
+    }
+
+    /// Set the boost amount of the desired car.
+    pub fn boost(self, amount: f32) -> Self {
+        self.car.boost_amount = Some(amount.into());
+        self
+    }
+
+    /// Set the location of the desired car.
+    pub fn location(self, loc: impl Into<Vector3>) -> Self {
+        let loc: Vector3Partial = loc.into().into();
+        self.car.physics.get_or_insert_default().location.replace(loc.into());
+        self
+    }
+
+    /// Set the location x value of the desired car.
+    pub fn location_x(self, x: f32) -> Self {
+        self.car.physics.get_or_insert_default().location.get_or_insert_default().x.get_or_insert_default().val = x;
+        self
+    }
+
+    /// Set the location y value of the desired car.
+    pub fn location_y(self, y: f32) -> Self {
+        self.car.physics.get_or_insert_default().location.get_or_insert_default().y.get_or_insert_default().val = y;
+        self
+    }
+
+    /// Set the location z value of the desired car.
+    pub fn location_z(self, z: f32) -> Self {
+        self.car.physics.get_or_insert_default().location.get_or_insert_default().z.get_or_insert_default().val = z;
+        self
+    }
+
+    /// Set the velocity of the desired car.
+    pub fn velocity(self, vel: impl Into<Vector3>) -> Self {
+        let vel: Vector3Partial = vel.into().into();
+        self.car.physics.get_or_insert_default().velocity.replace(vel.into());
+        self
+    }
+
+    /// Set the velocity x value of the desired car.
+    pub fn velocity_x(self, x: f32) -> Self {
+        self.car.physics.get_or_insert_default().velocity.get_or_insert_default().x.get_or_insert_default().val = x;
+        self
+    }
+
+    /// Set the velocity y value of the desired car.
+    pub fn velocity_y(self, y: f32) -> Self {
+        self.car.physics.get_or_insert_default().velocity.get_or_insert_default().y.get_or_insert_default().val = y;
+        self
+    }
+
+    /// Set the velocity z value of the desired car.
+    pub fn velocity_z(self, z: f32) -> Self {
+        self.car.physics.get_or_insert_default().velocity.get_or_insert_default().z.get_or_insert_default().val = z;
+        self
+    }
+
+    /// Set the rotation of the desired car.
+    pub fn rotation(self, rot: impl Into<Rotator>) -> Self {
+        let rot: RotatorPartial = rot.into().into();
+        self.car.physics.get_or_insert_default().rotation.replace(rot.into());
+        self
+    }
+
+    /// Set the rotation pitch of the desired car.
+    pub fn rotation_pitch(self, pitch: f32) -> Self {
+        self.car.physics.get_or_insert_default().rotation.get_or_insert_default().pitch.get_or_insert_default().val = pitch;
+        self
+    }
+
+    /// Set the rotation yaw of the desired car.
+    pub fn rotation_yaw(self, yaw: f32) -> Self {
+        self.car.physics.get_or_insert_default().rotation.get_or_insert_default().yaw.get_or_insert_default().val = yaw;
+        self
+    }
+
+    /// Set the rotation roll of the desired car.
+    pub fn rotation_roll(self, roll: f32) -> Self {
+        self.car.physics.get_or_insert_default().rotation.get_or_insert_default().roll.get_or_insert_default().val = roll;
+        self
+    }
+
+    /// Set the angular velocity of the desired car.
+    pub fn angular_velocity(self, ang_vel: impl Into<Vector3>) -> Self {
+        let ang_vel: Vector3Partial = ang_vel.into().into();
+        self.car.physics.get_or_insert_default().angular_velocity.replace(ang_vel.into());
+        self
+    }
+
+    /// Set the angular velocity x value of the desired car.
+    pub fn angular_velocity_x(self, x: f32) -> Self {
+        self.car.physics.get_or_insert_default().angular_velocity.get_or_insert_default().x.get_or_insert_default().val = x;
+        self
+    }
+
+    /// Set the angular velocity y value of the desired car.
+    pub fn angular_velocity_y(self, y: f32) -> Self {
+        self.car.physics.get_or_insert_default().angular_velocity.get_or_insert_default().y.get_or_insert_default().val = y;
+        self
+    }
+
+    /// Set the angular velocity z value of the desired car.
+    pub fn angular_velocity_z(self, z: f32) -> Self {
+        self.car.physics.get_or_insert_default().angular_velocity.get_or_insert_default().z.get_or_insert_default().val = z;
+        self
+    }
+}
+
+/// Allows for easy modification of a [DesiredBallState]. See [DesiredStateBuilder].
+#[derive(Debug, PartialEq, PartialOrd)]
+pub struct DesiredBallBuilder<'a> {
+    ball: &'a mut DesiredBallState,
+}
+
+#[allow(dead_code)]
+impl<'a> DesiredBallBuilder<'a> {
+    pub fn new(ball: &'a mut DesiredBallState) -> Self {
+        Self {
+            ball
+        }
+    }
+
+    /// Set the location of the desired ball.
+    pub fn location(self, loc: impl Into<Vector3>) -> Self {
+        let loc: Vector3Partial = loc.into().into();
+        self.ball.physics.location.replace(loc.into());
+        self
+    }
+
+    /// Set the location x value of the desired ball.
+    pub fn location_x(self, x: f32) -> Self {
+        self.ball.physics.location.get_or_insert_default().x.get_or_insert_default().val = x;
+        self
+    }
+
+    /// Set the location y value of the desired ball.
+    pub fn location_y(self, y: f32) -> Self {
+        self.ball.physics.location.get_or_insert_default().y.get_or_insert_default().val = y;
+        self
+    }
+
+    /// Set the location z value of the desired ball.
+    pub fn location_z(self, z: f32) -> Self {
+        self.ball.physics.location.get_or_insert_default().z.get_or_insert_default().val = z;
+        self
+    }
+
+    /// Set the velocity of the desired ball.
+    pub fn velocity(self, vel: impl Into<Vector3>) -> Self {
+        let vel: Vector3Partial = vel.into().into();
+        self.ball.physics.velocity.replace(vel.into());
+        self
+    }
+
+    /// Set the velocity x value of the desired ball.
+    pub fn velocity_x(self, x: f32) -> Self {
+        self.ball.physics.velocity.get_or_insert_default().x.get_or_insert_default().val = x;
+        self
+    }
+
+    /// Set the velocity y value of the desired ball.
+    pub fn velocity_y(self, y: f32) -> Self {
+        self.ball.physics.velocity.get_or_insert_default().y.get_or_insert_default().val = y;
+        self
+    }
+
+    /// Set the velocity z value of the desired ball.
+    pub fn velocity_z(self, z: f32) -> Self {
+        self.ball.physics.velocity.get_or_insert_default().z.get_or_insert_default().val = z;
+        self
+    }
+
+    /// Set the rotation of the desired ball.
+    pub fn rotation(self, rot: impl Into<Rotator>) -> Self {
+        let rot: RotatorPartial = rot.into().into();
+        self.ball.physics.rotation.replace(rot.into());
+        self
+    }
+
+    /// Set the rotation pitch of the desired ball.
+    pub fn rotation_pitch(self, pitch: f32) -> Self {
+        self.ball.physics.rotation.get_or_insert_default().pitch.get_or_insert_default().val = pitch;
+        self
+    }
+
+    /// Set the rotation yaw of the desired ball.
+    pub fn rotation_yaw(self, yaw: f32) -> Self {
+        self.ball.physics.rotation.get_or_insert_default().yaw.get_or_insert_default().val = yaw;
+        self
+    }
+
+    /// Set the rotation roll of the desired ball.
+    pub fn rotation_roll(self, roll: f32) -> Self {
+        self.ball.physics.rotation.get_or_insert_default().roll.get_or_insert_default().val = roll;
+        self
+    }
+
+    /// Set the angular velocity of the desired ball.
+    pub fn angular_velocity(self, ang_vel: impl Into<Vector3>) -> Self {
+        let ang_vel: Vector3Partial = ang_vel.into().into();
+        self.ball.physics.angular_velocity.replace(ang_vel.into());
+        self
+    }
+
+    /// Set the angular velocity x value of the desired ball.
+    pub fn angular_velocity_x(self, x: f32) -> Self {
+        self.ball.physics.angular_velocity.get_or_insert_default().x.get_or_insert_default().val = x;
+        self
+    }
+
+    /// Set the angular velocity y value of the desired ball.
+    pub fn angular_velocity_y(self, y: f32) -> Self {
+        self.ball.physics.angular_velocity.get_or_insert_default().y.get_or_insert_default().val = y;
+        self
+    }
+
+    /// Set the angular velocity z value of the desired ball.
+    pub fn angular_velocity_z(self, z: f32) -> Self {
+        self.ball.physics.angular_velocity.get_or_insert_default().z.get_or_insert_default().val = z;
+        self
+    }
+}

From 73a71d8b05f175cc3eb292f548d65bda1988240e Mon Sep 17 00:00:00 2001
From: NicEastvillage <nicoesterby@gmail.com>
Date: Sun, 22 Jun 2025 20:30:16 +0200
Subject: [PATCH 2/9] Remove unnecesary imports and fmt

---
 rlbot/examples/atba_agent/main.rs |   2 -
 rlbot/src/lib.rs                  |   2 +-
 rlbot/src/state_builder.rs        | 272 +++++++++++++++++++++++++-----
 3 files changed, 228 insertions(+), 48 deletions(-)

diff --git a/rlbot/examples/atba_agent/main.rs b/rlbot/examples/atba_agent/main.rs
index df4f7e6..7ec4ec5 100644
--- a/rlbot/examples/atba_agent/main.rs
+++ b/rlbot/examples/atba_agent/main.rs
@@ -9,8 +9,6 @@ use rlbot::{
     },
     util::{AgentEnvironment, PacketQueue},
 };
-use rlbot::state_builder::DesiredStateBuilder;
-use rlbot_flat::flat::{DesiredGameStateBuilder, Vector3};
 
 #[allow(dead_code)]
 struct AtbaAgent {
diff --git a/rlbot/src/lib.rs b/rlbot/src/lib.rs
index 59aee0f..6a33913 100644
--- a/rlbot/src/lib.rs
+++ b/rlbot/src/lib.rs
@@ -8,8 +8,8 @@ use rlbot_flat::planus::{self, ReadAsRoot};
 use thiserror::Error;
 
 pub mod agents;
-pub mod util;
 pub mod state_builder;
+pub mod util;
 
 #[cfg(feature = "glam")]
 pub use rlbot_flat::glam;
diff --git a/rlbot/src/state_builder.rs b/rlbot/src/state_builder.rs
index 26b8728..201fb9e 100644
--- a/rlbot/src/state_builder.rs
+++ b/rlbot/src/state_builder.rs
@@ -1,7 +1,10 @@
-use rlbot_flat::flat::{DesiredBallState, DesiredCarState, DesiredGameState, DesiredMatchInfo, Rotator, RotatorPartial, Vector3, Vector3Partial};
+use rlbot_flat::flat::{
+    DesiredBallState, DesiredCarState, DesiredGameState, DesiredMatchInfo, Rotator, RotatorPartial,
+    Vector3, Vector3Partial,
+};
 
 /// Utility for easy construction of [DesiredGameState]s using builder patterns.
-/// 
+///
 /// Example:
 /// ```example
 /// let dgs: DesiredGameState = DesiredStateBuilder::new()
@@ -16,25 +19,34 @@ use rlbot_flat::flat::{DesiredBallState, DesiredCarState, DesiredGameState, Desi
 /// ```
 #[derive(Default, Debug, Clone, PartialEq, PartialOrd)]
 pub struct DesiredStateBuilder {
-    state: DesiredGameState
+    state: DesiredGameState,
 }
 
 #[allow(dead_code)]
 impl DesiredStateBuilder {
     pub fn new() -> Self {
         Self {
-            state: DesiredGameState::default()
+            state: DesiredGameState::default(),
         }
     }
 
     /// Modify the desired match info.
-    pub fn match_info(mut self, build: impl FnOnce(DesiredMatchInfoBuilder) -> DesiredMatchInfoBuilder) -> Self {
-        build(DesiredMatchInfoBuilder::new(&mut self.state.match_info.get_or_insert_default()));
+    pub fn match_info(
+        mut self,
+        build: impl FnOnce(DesiredMatchInfoBuilder) -> DesiredMatchInfoBuilder,
+    ) -> Self {
+        build(DesiredMatchInfoBuilder::new(
+            &mut self.state.match_info.get_or_insert_default(),
+        ));
         self
     }
 
     /// Modify the desired ball at the given index.
-    pub fn ball(mut self, index: usize, build: impl FnOnce(DesiredBallBuilder) -> DesiredBallBuilder) -> Self {
+    pub fn ball(
+        mut self,
+        index: usize,
+        build: impl FnOnce(DesiredBallBuilder) -> DesiredBallBuilder,
+    ) -> Self {
         while self.state.ball_states.len() <= index {
             self.state.ball_states.push(Default::default());
         }
@@ -43,7 +55,11 @@ impl DesiredStateBuilder {
     }
 
     /// Modify the desired car at the given index.
-    pub fn car(mut self, index: usize, build: impl FnOnce(DesiredCarBuilder) -> DesiredCarBuilder) -> Self {
+    pub fn car(
+        mut self,
+        index: usize,
+        build: impl FnOnce(DesiredCarBuilder) -> DesiredCarBuilder,
+    ) -> Self {
         while self.state.car_states.len() <= index {
             self.state.car_states.push(Default::default());
         }
@@ -66,9 +82,7 @@ pub struct DesiredMatchInfoBuilder<'a> {
 #[allow(dead_code)]
 impl<'a> DesiredMatchInfoBuilder<'a> {
     pub fn new(info: &'a mut DesiredMatchInfo) -> Self {
-        Self {
-            info
-        }
+        Self { info }
     }
 
     /// Set the desired world gravity z.
@@ -93,9 +107,7 @@ pub struct DesiredCarBuilder<'a> {
 #[allow(dead_code)]
 impl<'a> DesiredCarBuilder<'a> {
     pub fn new(car: &'a mut DesiredCarState) -> Self {
-        Self {
-            car
-        }
+        Self { car }
     }
 
     /// Set the boost amount of the desired car.
@@ -107,100 +119,200 @@ impl<'a> DesiredCarBuilder<'a> {
     /// Set the location of the desired car.
     pub fn location(self, loc: impl Into<Vector3>) -> Self {
         let loc: Vector3Partial = loc.into().into();
-        self.car.physics.get_or_insert_default().location.replace(loc.into());
+        self.car
+            .physics
+            .get_or_insert_default()
+            .location
+            .replace(loc.into());
         self
     }
 
     /// Set the location x value of the desired car.
     pub fn location_x(self, x: f32) -> Self {
-        self.car.physics.get_or_insert_default().location.get_or_insert_default().x.get_or_insert_default().val = x;
+        self.car
+            .physics
+            .get_or_insert_default()
+            .location
+            .get_or_insert_default()
+            .x
+            .get_or_insert_default()
+            .val = x;
         self
     }
 
     /// Set the location y value of the desired car.
     pub fn location_y(self, y: f32) -> Self {
-        self.car.physics.get_or_insert_default().location.get_or_insert_default().y.get_or_insert_default().val = y;
+        self.car
+            .physics
+            .get_or_insert_default()
+            .location
+            .get_or_insert_default()
+            .y
+            .get_or_insert_default()
+            .val = y;
         self
     }
 
     /// Set the location z value of the desired car.
     pub fn location_z(self, z: f32) -> Self {
-        self.car.physics.get_or_insert_default().location.get_or_insert_default().z.get_or_insert_default().val = z;
+        self.car
+            .physics
+            .get_or_insert_default()
+            .location
+            .get_or_insert_default()
+            .z
+            .get_or_insert_default()
+            .val = z;
         self
     }
 
     /// Set the velocity of the desired car.
     pub fn velocity(self, vel: impl Into<Vector3>) -> Self {
         let vel: Vector3Partial = vel.into().into();
-        self.car.physics.get_or_insert_default().velocity.replace(vel.into());
+        self.car
+            .physics
+            .get_or_insert_default()
+            .velocity
+            .replace(vel.into());
         self
     }
 
     /// Set the velocity x value of the desired car.
     pub fn velocity_x(self, x: f32) -> Self {
-        self.car.physics.get_or_insert_default().velocity.get_or_insert_default().x.get_or_insert_default().val = x;
+        self.car
+            .physics
+            .get_or_insert_default()
+            .velocity
+            .get_or_insert_default()
+            .x
+            .get_or_insert_default()
+            .val = x;
         self
     }
 
     /// Set the velocity y value of the desired car.
     pub fn velocity_y(self, y: f32) -> Self {
-        self.car.physics.get_or_insert_default().velocity.get_or_insert_default().y.get_or_insert_default().val = y;
+        self.car
+            .physics
+            .get_or_insert_default()
+            .velocity
+            .get_or_insert_default()
+            .y
+            .get_or_insert_default()
+            .val = y;
         self
     }
 
     /// Set the velocity z value of the desired car.
     pub fn velocity_z(self, z: f32) -> Self {
-        self.car.physics.get_or_insert_default().velocity.get_or_insert_default().z.get_or_insert_default().val = z;
+        self.car
+            .physics
+            .get_or_insert_default()
+            .velocity
+            .get_or_insert_default()
+            .z
+            .get_or_insert_default()
+            .val = z;
         self
     }
 
     /// Set the rotation of the desired car.
     pub fn rotation(self, rot: impl Into<Rotator>) -> Self {
         let rot: RotatorPartial = rot.into().into();
-        self.car.physics.get_or_insert_default().rotation.replace(rot.into());
+        self.car
+            .physics
+            .get_or_insert_default()
+            .rotation
+            .replace(rot.into());
         self
     }
 
     /// Set the rotation pitch of the desired car.
     pub fn rotation_pitch(self, pitch: f32) -> Self {
-        self.car.physics.get_or_insert_default().rotation.get_or_insert_default().pitch.get_or_insert_default().val = pitch;
+        self.car
+            .physics
+            .get_or_insert_default()
+            .rotation
+            .get_or_insert_default()
+            .pitch
+            .get_or_insert_default()
+            .val = pitch;
         self
     }
 
     /// Set the rotation yaw of the desired car.
     pub fn rotation_yaw(self, yaw: f32) -> Self {
-        self.car.physics.get_or_insert_default().rotation.get_or_insert_default().yaw.get_or_insert_default().val = yaw;
+        self.car
+            .physics
+            .get_or_insert_default()
+            .rotation
+            .get_or_insert_default()
+            .yaw
+            .get_or_insert_default()
+            .val = yaw;
         self
     }
 
     /// Set the rotation roll of the desired car.
     pub fn rotation_roll(self, roll: f32) -> Self {
-        self.car.physics.get_or_insert_default().rotation.get_or_insert_default().roll.get_or_insert_default().val = roll;
+        self.car
+            .physics
+            .get_or_insert_default()
+            .rotation
+            .get_or_insert_default()
+            .roll
+            .get_or_insert_default()
+            .val = roll;
         self
     }
 
     /// Set the angular velocity of the desired car.
     pub fn angular_velocity(self, ang_vel: impl Into<Vector3>) -> Self {
         let ang_vel: Vector3Partial = ang_vel.into().into();
-        self.car.physics.get_or_insert_default().angular_velocity.replace(ang_vel.into());
+        self.car
+            .physics
+            .get_or_insert_default()
+            .angular_velocity
+            .replace(ang_vel.into());
         self
     }
 
     /// Set the angular velocity x value of the desired car.
     pub fn angular_velocity_x(self, x: f32) -> Self {
-        self.car.physics.get_or_insert_default().angular_velocity.get_or_insert_default().x.get_or_insert_default().val = x;
+        self.car
+            .physics
+            .get_or_insert_default()
+            .angular_velocity
+            .get_or_insert_default()
+            .x
+            .get_or_insert_default()
+            .val = x;
         self
     }
 
     /// Set the angular velocity y value of the desired car.
     pub fn angular_velocity_y(self, y: f32) -> Self {
-        self.car.physics.get_or_insert_default().angular_velocity.get_or_insert_default().y.get_or_insert_default().val = y;
+        self.car
+            .physics
+            .get_or_insert_default()
+            .angular_velocity
+            .get_or_insert_default()
+            .y
+            .get_or_insert_default()
+            .val = y;
         self
     }
 
     /// Set the angular velocity z value of the desired car.
     pub fn angular_velocity_z(self, z: f32) -> Self {
-        self.car.physics.get_or_insert_default().angular_velocity.get_or_insert_default().z.get_or_insert_default().val = z;
+        self.car
+            .physics
+            .get_or_insert_default()
+            .angular_velocity
+            .get_or_insert_default()
+            .z
+            .get_or_insert_default()
+            .val = z;
         self
     }
 }
@@ -214,9 +326,7 @@ pub struct DesiredBallBuilder<'a> {
 #[allow(dead_code)]
 impl<'a> DesiredBallBuilder<'a> {
     pub fn new(ball: &'a mut DesiredBallState) -> Self {
-        Self {
-            ball
-        }
+        Self { ball }
     }
 
     /// Set the location of the desired ball.
@@ -228,19 +338,37 @@ impl<'a> DesiredBallBuilder<'a> {
 
     /// Set the location x value of the desired ball.
     pub fn location_x(self, x: f32) -> Self {
-        self.ball.physics.location.get_or_insert_default().x.get_or_insert_default().val = x;
+        self.ball
+            .physics
+            .location
+            .get_or_insert_default()
+            .x
+            .get_or_insert_default()
+            .val = x;
         self
     }
 
     /// Set the location y value of the desired ball.
     pub fn location_y(self, y: f32) -> Self {
-        self.ball.physics.location.get_or_insert_default().y.get_or_insert_default().val = y;
+        self.ball
+            .physics
+            .location
+            .get_or_insert_default()
+            .y
+            .get_or_insert_default()
+            .val = y;
         self
     }
 
     /// Set the location z value of the desired ball.
     pub fn location_z(self, z: f32) -> Self {
-        self.ball.physics.location.get_or_insert_default().z.get_or_insert_default().val = z;
+        self.ball
+            .physics
+            .location
+            .get_or_insert_default()
+            .z
+            .get_or_insert_default()
+            .val = z;
         self
     }
 
@@ -253,19 +381,37 @@ impl<'a> DesiredBallBuilder<'a> {
 
     /// Set the velocity x value of the desired ball.
     pub fn velocity_x(self, x: f32) -> Self {
-        self.ball.physics.velocity.get_or_insert_default().x.get_or_insert_default().val = x;
+        self.ball
+            .physics
+            .velocity
+            .get_or_insert_default()
+            .x
+            .get_or_insert_default()
+            .val = x;
         self
     }
 
     /// Set the velocity y value of the desired ball.
     pub fn velocity_y(self, y: f32) -> Self {
-        self.ball.physics.velocity.get_or_insert_default().y.get_or_insert_default().val = y;
+        self.ball
+            .physics
+            .velocity
+            .get_or_insert_default()
+            .y
+            .get_or_insert_default()
+            .val = y;
         self
     }
 
     /// Set the velocity z value of the desired ball.
     pub fn velocity_z(self, z: f32) -> Self {
-        self.ball.physics.velocity.get_or_insert_default().z.get_or_insert_default().val = z;
+        self.ball
+            .physics
+            .velocity
+            .get_or_insert_default()
+            .z
+            .get_or_insert_default()
+            .val = z;
         self
     }
 
@@ -278,19 +424,37 @@ impl<'a> DesiredBallBuilder<'a> {
 
     /// Set the rotation pitch of the desired ball.
     pub fn rotation_pitch(self, pitch: f32) -> Self {
-        self.ball.physics.rotation.get_or_insert_default().pitch.get_or_insert_default().val = pitch;
+        self.ball
+            .physics
+            .rotation
+            .get_or_insert_default()
+            .pitch
+            .get_or_insert_default()
+            .val = pitch;
         self
     }
 
     /// Set the rotation yaw of the desired ball.
     pub fn rotation_yaw(self, yaw: f32) -> Self {
-        self.ball.physics.rotation.get_or_insert_default().yaw.get_or_insert_default().val = yaw;
+        self.ball
+            .physics
+            .rotation
+            .get_or_insert_default()
+            .yaw
+            .get_or_insert_default()
+            .val = yaw;
         self
     }
 
     /// Set the rotation roll of the desired ball.
     pub fn rotation_roll(self, roll: f32) -> Self {
-        self.ball.physics.rotation.get_or_insert_default().roll.get_or_insert_default().val = roll;
+        self.ball
+            .physics
+            .rotation
+            .get_or_insert_default()
+            .roll
+            .get_or_insert_default()
+            .val = roll;
         self
     }
 
@@ -303,19 +467,37 @@ impl<'a> DesiredBallBuilder<'a> {
 
     /// Set the angular velocity x value of the desired ball.
     pub fn angular_velocity_x(self, x: f32) -> Self {
-        self.ball.physics.angular_velocity.get_or_insert_default().x.get_or_insert_default().val = x;
+        self.ball
+            .physics
+            .angular_velocity
+            .get_or_insert_default()
+            .x
+            .get_or_insert_default()
+            .val = x;
         self
     }
 
     /// Set the angular velocity y value of the desired ball.
     pub fn angular_velocity_y(self, y: f32) -> Self {
-        self.ball.physics.angular_velocity.get_or_insert_default().y.get_or_insert_default().val = y;
+        self.ball
+            .physics
+            .angular_velocity
+            .get_or_insert_default()
+            .y
+            .get_or_insert_default()
+            .val = y;
         self
     }
 
     /// Set the angular velocity z value of the desired ball.
     pub fn angular_velocity_z(self, z: f32) -> Self {
-        self.ball.physics.angular_velocity.get_or_insert_default().z.get_or_insert_default().val = z;
+        self.ball
+            .physics
+            .angular_velocity
+            .get_or_insert_default()
+            .z
+            .get_or_insert_default()
+            .val = z;
         self
     }
 }

From d44876b24fc3e2f41a673335d75904e1323454d8 Mon Sep 17 00:00:00 2001
From: NicEastvillage <nicoesterby@gmail.com>
Date: Sun, 22 Jun 2025 20:31:52 +0200
Subject: [PATCH 3/9] Bump version

---
 Cargo.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Cargo.toml b/Cargo.toml
index 4ca3c9d..c4eba53 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -2,7 +2,7 @@
 members = ["rlbot", "rlbot_flat"]
 resolver = "3"
 package.license-file = "LICENSE"
-package.version = "0.1.0"
+package.version = "0.1.1"
 
 [profile.release]
 lto = "fat"

From 875dde1b098f0b2bf5ebbd66b41f667de68ac486 Mon Sep 17 00:00:00 2001
From: NicEastvillage <nicoesterby@gmail.com>
Date: Mon, 23 Jun 2025 20:44:11 +0200
Subject: [PATCH 4/9] Add easy way to state set all cars/balls

---
 rlbot/src/state_builder.rs | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/rlbot/src/state_builder.rs b/rlbot/src/state_builder.rs
index 201fb9e..a77eba2 100644
--- a/rlbot/src/state_builder.rs
+++ b/rlbot/src/state_builder.rs
@@ -53,6 +53,14 @@ impl DesiredStateBuilder {
         build(DesiredBallBuilder::new(&mut self.state.ball_states[index]));
         self
     }
+    
+    /// Modify all desired balls.
+    pub fn all_balls(mut self, build: impl Fn(usize, DesiredBallBuilder) -> DesiredBallBuilder) -> Self {
+        for (i, ball) in self.state.ball_states.iter_mut().enumerate() {
+            build(i, DesiredBallBuilder::new(ball));
+        }
+        self
+    }
 
     /// Modify the desired car at the given index.
     pub fn car(
@@ -67,6 +75,14 @@ impl DesiredStateBuilder {
         self
     }
 
+    /// Modify all desired cars.
+    pub fn all_cars(mut self, build: impl Fn(usize, DesiredCarBuilder) -> DesiredCarBuilder) -> Self {
+        for (i, car) in self.state.car_states.iter_mut().enumerate() {
+            build(i, DesiredCarBuilder::new(car));
+        }
+        self
+    }
+
     /// Extract the resulting [DesiredGameState].
     pub fn build(self) -> DesiredGameState {
         self.state

From c269461363254ddbe395dbe5cfbdd3366104c0b0 Mon Sep 17 00:00:00 2001
From: NicEastvillage <nicoesterby@gmail.com>
Date: Mon, 23 Jun 2025 20:48:30 +0200
Subject: [PATCH 5/9] Fix state builder all_cars/balls only affecting already
 set cars/balls

---
 rlbot/src/state_builder.rs | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/rlbot/src/state_builder.rs b/rlbot/src/state_builder.rs
index a77eba2..6178e19 100644
--- a/rlbot/src/state_builder.rs
+++ b/rlbot/src/state_builder.rs
@@ -1,3 +1,4 @@
+use std::ops::Range;
 use rlbot_flat::flat::{
     DesiredBallState, DesiredCarState, DesiredGameState, DesiredMatchInfo, Rotator, RotatorPartial,
     Vector3, Vector3Partial,
@@ -55,8 +56,11 @@ impl DesiredStateBuilder {
     }
     
     /// Modify all desired balls.
-    pub fn all_balls(mut self, build: impl Fn(usize, DesiredBallBuilder) -> DesiredBallBuilder) -> Self {
-        for (i, ball) in self.state.ball_states.iter_mut().enumerate() {
+    pub fn all_balls(mut self, range: Range<usize>, build: impl Fn(usize, DesiredBallBuilder) -> DesiredBallBuilder) -> Self {
+        while self.state.ball_states.len() < range.end {
+            self.state.ball_states.push(Default::default());
+        }
+        for (i, ball) in self.state.ball_states[range].iter_mut().enumerate() {
             build(i, DesiredBallBuilder::new(ball));
         }
         self
@@ -76,8 +80,11 @@ impl DesiredStateBuilder {
     }
 
     /// Modify all desired cars.
-    pub fn all_cars(mut self, build: impl Fn(usize, DesiredCarBuilder) -> DesiredCarBuilder) -> Self {
-        for (i, car) in self.state.car_states.iter_mut().enumerate() {
+    pub fn all_cars(mut self, range: Range<usize>, build: impl Fn(usize, DesiredCarBuilder) -> DesiredCarBuilder) -> Self {
+        while self.state.ball_states.len() < range.end {
+            self.state.ball_states.push(Default::default());
+        }
+        for (i, car) in self.state.car_states[range].iter_mut().enumerate() {
             build(i, DesiredCarBuilder::new(car));
         }
         self

From 9dd47f52d14d6eec20afdfea9e827a572225500a Mon Sep 17 00:00:00 2001
From: NicEastvillage <nicoesterby@gmail.com>
Date: Mon, 23 Jun 2025 21:13:04 +0200
Subject: [PATCH 6/9] fmt

---
 rlbot/src/state_builder.rs | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/rlbot/src/state_builder.rs b/rlbot/src/state_builder.rs
index 6178e19..82f79e1 100644
--- a/rlbot/src/state_builder.rs
+++ b/rlbot/src/state_builder.rs
@@ -1,8 +1,8 @@
-use std::ops::Range;
 use rlbot_flat::flat::{
     DesiredBallState, DesiredCarState, DesiredGameState, DesiredMatchInfo, Rotator, RotatorPartial,
     Vector3, Vector3Partial,
 };
+use std::ops::Range;
 
 /// Utility for easy construction of [DesiredGameState]s using builder patterns.
 ///
@@ -54,9 +54,13 @@ impl DesiredStateBuilder {
         build(DesiredBallBuilder::new(&mut self.state.ball_states[index]));
         self
     }
-    
+
     /// Modify all desired balls.
-    pub fn all_balls(mut self, range: Range<usize>, build: impl Fn(usize, DesiredBallBuilder) -> DesiredBallBuilder) -> Self {
+    pub fn all_balls(
+        mut self,
+        range: Range<usize>,
+        build: impl Fn(usize, DesiredBallBuilder) -> DesiredBallBuilder,
+    ) -> Self {
         while self.state.ball_states.len() < range.end {
             self.state.ball_states.push(Default::default());
         }
@@ -80,7 +84,11 @@ impl DesiredStateBuilder {
     }
 
     /// Modify all desired cars.
-    pub fn all_cars(mut self, range: Range<usize>, build: impl Fn(usize, DesiredCarBuilder) -> DesiredCarBuilder) -> Self {
+    pub fn all_cars(
+        mut self,
+        range: Range<usize>,
+        build: impl Fn(usize, DesiredCarBuilder) -> DesiredCarBuilder,
+    ) -> Self {
         while self.state.ball_states.len() < range.end {
             self.state.ball_states.push(Default::default());
         }

From 7f7835e58694bee0ca041338e8406ef145dc07fd Mon Sep 17 00:00:00 2001
From: NicEastvillage <nicoesterby@gmail.com>
Date: Wed, 25 Jun 2025 19:31:29 +0200
Subject: [PATCH 7/9] Use extension traits instead

---
 Cargo.toml                 |   2 +-
 rlbot/src/state_builder.rs | 697 ++++++++++++-------------------------
 2 files changed, 217 insertions(+), 482 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index c4eba53..4ca3c9d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -2,7 +2,7 @@
 members = ["rlbot", "rlbot_flat"]
 resolver = "3"
 package.license-file = "LICENSE"
-package.version = "0.1.1"
+package.version = "0.1.0"
 
 [profile.release]
 lto = "fat"
diff --git a/rlbot/src/state_builder.rs b/rlbot/src/state_builder.rs
index 82f79e1..b2650b7 100644
--- a/rlbot/src/state_builder.rs
+++ b/rlbot/src/state_builder.rs
@@ -1,534 +1,269 @@
-use rlbot_flat::flat::{
-    DesiredBallState, DesiredCarState, DesiredGameState, DesiredMatchInfo, Rotator, RotatorPartial,
-    Vector3, Vector3Partial,
-};
-use std::ops::Range;
+use rlbot_flat::flat::{DesiredBallState, DesiredCarState, DesiredGameState, DesiredMatchInfo, DesiredPhysics, Rotator, RotatorPartial, Vector3, Vector3Partial};
 
-/// Utility for easy construction of [DesiredGameState]s using builder patterns.
+/// Extension methods for easy construction of [DesiredGameState].
 ///
 /// Example:
-/// ```example
-/// let dgs: DesiredGameState = DesiredStateBuilder::new()
-///             .car(0, |c| c
-///                 .location(Vector3::default())
-///                 .boost(100.))
-///             .ball(0, |b| b
-///                 .location_z(0.)
-///                 .velocity_z(0.)
-///             )
-///             .build();
+/// ```rust
+/// use rlbot::state_builder::{DesiredCarStateExt, DesiredGameStateExt, DesiredPhysicsExt};
+/// use rlbot_flat::flat::{DesiredGameState, Vector3};
+/// let mut dgs = DesiredGameState::default();
+/// dgs.mod_car(0, |c| {
+///     c.set_location(Vector3::default());
+///     c.set_boost(100.);
+/// });
+/// dgs.mod_balls(0..5, |_, b| {
+///     b.set_location_z(0.);
+///     b.set_velocity_z(0.);
+/// });
 /// ```
-#[derive(Default, Debug, Clone, PartialEq, PartialOrd)]
-pub struct DesiredStateBuilder {
-    state: DesiredGameState,
+pub trait DesiredGameStateExt {
+
+    fn mod_match_info(
+        &mut self,
+        build: impl FnOnce(&mut DesiredMatchInfo),
+    );
+
+    fn mod_car(
+        &mut self,
+        index: usize,
+        build: impl FnOnce(&mut DesiredCarState),
+    );
+
+    fn mod_cars(
+        &mut self,
+        indices: impl IntoIterator<Item = usize>,
+        build: impl Fn(usize, &mut DesiredCarState),
+    );
+
+    fn mod_ball(
+        &mut self,
+        index: usize,
+        build: impl FnOnce(&mut DesiredBallState),
+    );
+
+    fn mod_balls(
+        &mut self,
+        indices: impl IntoIterator<Item = usize>,
+        build: impl Fn(usize, &mut DesiredBallState),
+    );
 }
 
 #[allow(dead_code)]
-impl DesiredStateBuilder {
-    pub fn new() -> Self {
-        Self {
-            state: DesiredGameState::default(),
-        }
-    }
+impl DesiredGameStateExt for DesiredGameState {
 
     /// Modify the desired match info.
-    pub fn match_info(
-        mut self,
-        build: impl FnOnce(DesiredMatchInfoBuilder) -> DesiredMatchInfoBuilder,
-    ) -> Self {
-        build(DesiredMatchInfoBuilder::new(
-            &mut self.state.match_info.get_or_insert_default(),
-        ));
-        self
+    fn mod_match_info(
+        &mut self,
+        build: impl FnOnce(&mut DesiredMatchInfo),
+    ) {
+        build(self.match_info.get_or_insert_default());
     }
 
-    /// Modify the desired ball at the given index.
-    pub fn ball(
-        mut self,
+    /// Modify the desired car at the given index.
+    fn mod_car(
+        &mut self,
         index: usize,
-        build: impl FnOnce(DesiredBallBuilder) -> DesiredBallBuilder,
-    ) -> Self {
-        while self.state.ball_states.len() <= index {
-            self.state.ball_states.push(Default::default());
+        build: impl FnOnce(&mut DesiredCarState),
+    ) {
+        if self.car_states.len() <= index {
+            self.car_states.resize(index + 1, Default::default());
         }
-        build(DesiredBallBuilder::new(&mut self.state.ball_states[index]));
-        self
+        build(&mut self.car_states[index]);
     }
 
-    /// Modify all desired balls.
-    pub fn all_balls(
-        mut self,
-        range: Range<usize>,
-        build: impl Fn(usize, DesiredBallBuilder) -> DesiredBallBuilder,
-    ) -> Self {
-        while self.state.ball_states.len() < range.end {
-            self.state.ball_states.push(Default::default());
-        }
-        for (i, ball) in self.state.ball_states[range].iter_mut().enumerate() {
-            build(i, DesiredBallBuilder::new(ball));
+    /// Modify all desired cars.
+    fn mod_cars(
+        &mut self,
+        indices: impl IntoIterator<Item = usize>,
+        build: impl Fn(usize, &mut DesiredCarState),
+    ) {
+        for i in indices {
+            if self.car_states.len() <= i {
+                self.car_states.resize(i + 1, Default::default());
+            }
+            build(i, &mut self.car_states[i]);
         }
-        self
     }
 
-    /// Modify the desired car at the given index.
-    pub fn car(
-        mut self,
+    /// Modify the desired ball at the given index.
+    fn mod_ball(
+        &mut self,
         index: usize,
-        build: impl FnOnce(DesiredCarBuilder) -> DesiredCarBuilder,
-    ) -> Self {
-        while self.state.car_states.len() <= index {
-            self.state.car_states.push(Default::default());
+        build: impl FnOnce(&mut DesiredBallState),
+    ) {
+        if self.ball_states.len() <= index {
+            self.ball_states.resize(index + 1, Default::default());
         }
-        build(DesiredCarBuilder::new(&mut self.state.car_states[index]));
-        self
+        build(&mut self.ball_states[index]);
     }
 
-    /// Modify all desired cars.
-    pub fn all_cars(
-        mut self,
-        range: Range<usize>,
-        build: impl Fn(usize, DesiredCarBuilder) -> DesiredCarBuilder,
-    ) -> Self {
-        while self.state.ball_states.len() < range.end {
-            self.state.ball_states.push(Default::default());
-        }
-        for (i, car) in self.state.car_states[range].iter_mut().enumerate() {
-            build(i, DesiredCarBuilder::new(car));
+    /// Modify all desired balls.
+    fn mod_balls(
+        &mut self,
+        indices: impl IntoIterator<Item = usize>,
+        build: impl Fn(usize, &mut DesiredBallState),
+    ) {
+        for i in indices {
+            if self.ball_states.len() <= i {
+                self.ball_states.resize(i + 1, Default::default());
+            }
+            build(i, &mut self.ball_states[i]);
         }
-        self
-    }
-
-    /// Extract the resulting [DesiredGameState].
-    pub fn build(self) -> DesiredGameState {
-        self.state
     }
 }
 
-/// Allows for easy modification of a [DesiredMatchInfo]. See [DesiredStateBuilder].
-#[derive(Debug, PartialEq, PartialOrd)]
-pub struct DesiredMatchInfoBuilder<'a> {
-    info: &'a mut DesiredMatchInfo,
+/// Extension methods for easy construction of a [DesiredMatchInfo].
+pub trait DesiredMatchInfoExt {
+    fn set_gravity_z(&mut self, gravity: f32);
+    fn set_game_speed(&mut self, speed: f32);
 }
 
 #[allow(dead_code)]
-impl<'a> DesiredMatchInfoBuilder<'a> {
-    pub fn new(info: &'a mut DesiredMatchInfo) -> Self {
-        Self { info }
-    }
+impl DesiredMatchInfoExt for DesiredMatchInfo {
 
     /// Set the desired world gravity z.
-    pub fn gravity_z(self, gravity_z: f32) -> Self {
-        self.info.world_gravity_z.get_or_insert_default().val = gravity_z;
-        self
+    fn set_gravity_z(&mut self, gravity_z: f32) {
+        self.world_gravity_z.get_or_insert_default().val = gravity_z;
     }
 
     /// Set the desired game speed.
-    pub fn game_speed(self, game_speed: f32) -> Self {
-        self.info.game_speed.get_or_insert_default().val = game_speed;
-        self
+    fn set_game_speed(&mut self, speed: f32) {
+        self.game_speed.get_or_insert_default().val = speed;
     }
 }
 
-/// Allows for easy modification of a [DesiredCarState]. See [DesiredStateBuilder].
-#[derive(Debug, PartialEq, PartialOrd)]
-pub struct DesiredCarBuilder<'a> {
-    car: &'a mut DesiredCarState,
+/// Extension methods for easy construction of a [DesiredCarState].
+pub trait DesiredCarStateExt {
+    fn set_boost(&mut self, amount: f32);
+    fn mod_physics(&mut self, build: impl FnOnce(&mut DesiredPhysics));
 }
 
-#[allow(dead_code)]
-impl<'a> DesiredCarBuilder<'a> {
-    pub fn new(car: &'a mut DesiredCarState) -> Self {
-        Self { car }
-    }
-
-    /// Set the boost amount of the desired car.
-    pub fn boost(self, amount: f32) -> Self {
-        self.car.boost_amount = Some(amount.into());
-        self
-    }
-
-    /// Set the location of the desired car.
-    pub fn location(self, loc: impl Into<Vector3>) -> Self {
-        let loc: Vector3Partial = loc.into().into();
-        self.car
-            .physics
-            .get_or_insert_default()
-            .location
-            .replace(loc.into());
-        self
-    }
-
-    /// Set the location x value of the desired car.
-    pub fn location_x(self, x: f32) -> Self {
-        self.car
-            .physics
-            .get_or_insert_default()
-            .location
-            .get_or_insert_default()
-            .x
-            .get_or_insert_default()
-            .val = x;
-        self
-    }
-
-    /// Set the location y value of the desired car.
-    pub fn location_y(self, y: f32) -> Self {
-        self.car
-            .physics
-            .get_or_insert_default()
-            .location
-            .get_or_insert_default()
-            .y
-            .get_or_insert_default()
-            .val = y;
-        self
-    }
-
-    /// Set the location z value of the desired car.
-    pub fn location_z(self, z: f32) -> Self {
-        self.car
-            .physics
-            .get_or_insert_default()
-            .location
-            .get_or_insert_default()
-            .z
-            .get_or_insert_default()
-            .val = z;
-        self
-    }
-
-    /// Set the velocity of the desired car.
-    pub fn velocity(self, vel: impl Into<Vector3>) -> Self {
-        let vel: Vector3Partial = vel.into().into();
-        self.car
-            .physics
-            .get_or_insert_default()
-            .velocity
-            .replace(vel.into());
-        self
-    }
-
-    /// Set the velocity x value of the desired car.
-    pub fn velocity_x(self, x: f32) -> Self {
-        self.car
-            .physics
-            .get_or_insert_default()
-            .velocity
-            .get_or_insert_default()
-            .x
-            .get_or_insert_default()
-            .val = x;
-        self
-    }
-
-    /// Set the velocity y value of the desired car.
-    pub fn velocity_y(self, y: f32) -> Self {
-        self.car
-            .physics
-            .get_or_insert_default()
-            .velocity
-            .get_or_insert_default()
-            .y
-            .get_or_insert_default()
-            .val = y;
-        self
-    }
+impl DesiredCarStateExt for DesiredCarState {
 
-    /// Set the velocity z value of the desired car.
-    pub fn velocity_z(self, z: f32) -> Self {
-        self.car
-            .physics
-            .get_or_insert_default()
-            .velocity
-            .get_or_insert_default()
-            .z
-            .get_or_insert_default()
-            .val = z;
-        self
+    /// Set the boost amount of this car.
+    fn set_boost(&mut self, amount: f32) {
+        self.boost_amount.get_or_insert_default().val = amount;
     }
 
-    /// Set the rotation of the desired car.
-    pub fn rotation(self, rot: impl Into<Rotator>) -> Self {
-        let rot: RotatorPartial = rot.into().into();
-        self.car
-            .physics
-            .get_or_insert_default()
-            .rotation
-            .replace(rot.into());
-        self
-    }
-
-    /// Set the rotation pitch of the desired car.
-    pub fn rotation_pitch(self, pitch: f32) -> Self {
-        self.car
-            .physics
-            .get_or_insert_default()
-            .rotation
-            .get_or_insert_default()
-            .pitch
-            .get_or_insert_default()
-            .val = pitch;
-        self
-    }
-
-    /// Set the rotation yaw of the desired car.
-    pub fn rotation_yaw(self, yaw: f32) -> Self {
-        self.car
-            .physics
-            .get_or_insert_default()
-            .rotation
-            .get_or_insert_default()
-            .yaw
-            .get_or_insert_default()
-            .val = yaw;
-        self
-    }
-
-    /// Set the rotation roll of the desired car.
-    pub fn rotation_roll(self, roll: f32) -> Self {
-        self.car
-            .physics
-            .get_or_insert_default()
-            .rotation
-            .get_or_insert_default()
-            .roll
-            .get_or_insert_default()
-            .val = roll;
-        self
-    }
-
-    /// Set the angular velocity of the desired car.
-    pub fn angular_velocity(self, ang_vel: impl Into<Vector3>) -> Self {
-        let ang_vel: Vector3Partial = ang_vel.into().into();
-        self.car
-            .physics
-            .get_or_insert_default()
-            .angular_velocity
-            .replace(ang_vel.into());
-        self
-    }
-
-    /// Set the angular velocity x value of the desired car.
-    pub fn angular_velocity_x(self, x: f32) -> Self {
-        self.car
-            .physics
-            .get_or_insert_default()
-            .angular_velocity
-            .get_or_insert_default()
-            .x
-            .get_or_insert_default()
-            .val = x;
-        self
-    }
-
-    /// Set the angular velocity y value of the desired car.
-    pub fn angular_velocity_y(self, y: f32) -> Self {
-        self.car
-            .physics
-            .get_or_insert_default()
-            .angular_velocity
-            .get_or_insert_default()
-            .y
-            .get_or_insert_default()
-            .val = y;
-        self
-    }
-
-    /// Set the angular velocity z value of the desired car.
-    pub fn angular_velocity_z(self, z: f32) -> Self {
-        self.car
-            .physics
-            .get_or_insert_default()
-            .angular_velocity
-            .get_or_insert_default()
-            .z
-            .get_or_insert_default()
-            .val = z;
-        self
+    /// Modify the physics of this car.
+    fn mod_physics(&mut self, build: impl FnOnce(&mut DesiredPhysics)) {
+        build(self.physics.get_or_insert_default());
     }
 }
 
-/// Allows for easy modification of a [DesiredBallState]. See [DesiredStateBuilder].
-#[derive(Debug, PartialEq, PartialOrd)]
-pub struct DesiredBallBuilder<'a> {
-    ball: &'a mut DesiredBallState,
+/// Extension methods for easy construction of a [DesiredBallState].
+pub trait DesiredBallStateExt {
+    fn mod_physics(&mut self, build: impl FnOnce(&mut DesiredPhysics));
 }
 
-#[allow(dead_code)]
-impl<'a> DesiredBallBuilder<'a> {
-    pub fn new(ball: &'a mut DesiredBallState) -> Self {
-        Self { ball }
-    }
-
-    /// Set the location of the desired ball.
-    pub fn location(self, loc: impl Into<Vector3>) -> Self {
-        let loc: Vector3Partial = loc.into().into();
-        self.ball.physics.location.replace(loc.into());
-        self
-    }
-
-    /// Set the location x value of the desired ball.
-    pub fn location_x(self, x: f32) -> Self {
-        self.ball
-            .physics
-            .location
-            .get_or_insert_default()
-            .x
-            .get_or_insert_default()
-            .val = x;
-        self
-    }
-
-    /// Set the location y value of the desired ball.
-    pub fn location_y(self, y: f32) -> Self {
-        self.ball
-            .physics
-            .location
-            .get_or_insert_default()
-            .y
-            .get_or_insert_default()
-            .val = y;
-        self
-    }
+impl DesiredBallStateExt for DesiredBallState {
 
-    /// Set the location z value of the desired ball.
-    pub fn location_z(self, z: f32) -> Self {
-        self.ball
-            .physics
-            .location
-            .get_or_insert_default()
-            .z
-            .get_or_insert_default()
-            .val = z;
-        self
-    }
-
-    /// Set the velocity of the desired ball.
-    pub fn velocity(self, vel: impl Into<Vector3>) -> Self {
-        let vel: Vector3Partial = vel.into().into();
-        self.ball.physics.velocity.replace(vel.into());
-        self
-    }
-
-    /// Set the velocity x value of the desired ball.
-    pub fn velocity_x(self, x: f32) -> Self {
-        self.ball
-            .physics
-            .velocity
-            .get_or_insert_default()
-            .x
-            .get_or_insert_default()
-            .val = x;
-        self
-    }
-
-    /// Set the velocity y value of the desired ball.
-    pub fn velocity_y(self, y: f32) -> Self {
-        self.ball
-            .physics
-            .velocity
-            .get_or_insert_default()
-            .y
-            .get_or_insert_default()
-            .val = y;
-        self
-    }
-
-    /// Set the velocity z value of the desired ball.
-    pub fn velocity_z(self, z: f32) -> Self {
-        self.ball
-            .physics
-            .velocity
-            .get_or_insert_default()
-            .z
-            .get_or_insert_default()
-            .val = z;
-        self
-    }
-
-    /// Set the rotation of the desired ball.
-    pub fn rotation(self, rot: impl Into<Rotator>) -> Self {
-        let rot: RotatorPartial = rot.into().into();
-        self.ball.physics.rotation.replace(rot.into());
-        self
-    }
-
-    /// Set the rotation pitch of the desired ball.
-    pub fn rotation_pitch(self, pitch: f32) -> Self {
-        self.ball
-            .physics
-            .rotation
-            .get_or_insert_default()
-            .pitch
-            .get_or_insert_default()
-            .val = pitch;
-        self
-    }
-
-    /// Set the rotation yaw of the desired ball.
-    pub fn rotation_yaw(self, yaw: f32) -> Self {
-        self.ball
-            .physics
-            .rotation
-            .get_or_insert_default()
-            .yaw
-            .get_or_insert_default()
-            .val = yaw;
-        self
-    }
-
-    /// Set the rotation roll of the desired ball.
-    pub fn rotation_roll(self, roll: f32) -> Self {
-        self.ball
-            .physics
-            .rotation
-            .get_or_insert_default()
-            .roll
-            .get_or_insert_default()
-            .val = roll;
-        self
-    }
-
-    /// Set the angular velocity of the desired ball.
-    pub fn angular_velocity(self, ang_vel: impl Into<Vector3>) -> Self {
-        let ang_vel: Vector3Partial = ang_vel.into().into();
-        self.ball.physics.angular_velocity.replace(ang_vel.into());
-        self
+    /// Modify the physics of this ball.
+    fn mod_physics(&mut self, build: impl FnOnce(&mut DesiredPhysics)) {
+        build(&mut self.physics)
     }
+}
 
-    /// Set the angular velocity x value of the desired ball.
-    pub fn angular_velocity_x(self, x: f32) -> Self {
-        self.ball
-            .physics
-            .angular_velocity
-            .get_or_insert_default()
-            .x
-            .get_or_insert_default()
-            .val = x;
-        self
-    }
+/// Extension methods for easy construction of [DesiredPhysics].
+pub trait DesiredPhysicsExt {
+    fn set_location(&mut self, loc: impl Into<Vector3>);
+    fn set_location_x(&mut self, x: f32);
+    fn set_location_y(&mut self, y: f32);
+    fn set_location_z(&mut self, z: f32);
+    fn set_velocity(&mut self, vel: impl Into<Vector3>);
+    fn set_velocity_x(&mut self, x: f32);
+    fn set_velocity_y(&mut self, y: f32);
+    fn set_velocity_z(&mut self, z: f32);
+    fn set_rotation(&mut self, rot: impl Into<Rotator>);
+    fn set_rotation_pitch(&mut self, pitch: f32);
+    fn set_rotation_yaw(&mut self, yaw: f32);
+    fn set_rotation_roll(&mut self, roll: f32);
+    fn set_angular_velocity(&mut self, ang_vel: impl Into<Vector3>);
+    fn set_angular_velocity_x(&mut self, x: f32);
+    fn set_angular_velocity_y(&mut self, y: f32);
+    fn set_angular_velocity_z(&mut self, z: f32);
+}
 
-    /// Set the angular velocity y value of the desired ball.
-    pub fn angular_velocity_y(self, y: f32) -> Self {
-        self.ball
-            .physics
-            .angular_velocity
-            .get_or_insert_default()
-            .y
-            .get_or_insert_default()
-            .val = y;
-        self
-    }
+macro_rules! physics_path {
+    ( $self:ident slf ) => { $self };
+    ( $self:ident physics) => { $self.physics };
+    ( $self:ident optional_physics) => { $self.physics.get_or_insert_default() };
+}
 
-    /// Set the angular velocity z value of the desired ball.
-    pub fn angular_velocity_z(self, z: f32) -> Self {
-        self.ball
-            .physics
-            .angular_velocity
-            .get_or_insert_default()
-            .z
-            .get_or_insert_default()
-            .val = z;
-        self
-    }
+macro_rules! desired_physics_ext {
+    ( $t:ty; $p:ident ) => {
+        impl DesiredPhysicsExt for $t {
+            fn set_location(&mut self, loc: impl Into<Vector3>) {
+                let loc: Vector3Partial = loc.into().into();
+                physics_path!(self $p).location.replace(loc.into());
+            }
+
+            fn set_location_x(&mut self, x: f32) {
+                physics_path!(self $p).location.get_or_insert_default().x.get_or_insert_default().val = x;
+            }
+
+            fn set_location_y(&mut self, y: f32) {
+                physics_path!(self $p).location.get_or_insert_default().y.get_or_insert_default().val = y;
+            }
+
+            fn set_location_z(&mut self, z: f32) {
+                physics_path!(self $p).location.get_or_insert_default().z.get_or_insert_default().val = z;
+            }
+
+            fn set_velocity(&mut self, vel: impl Into<Vector3>) {
+                let vel: Vector3Partial = vel.into().into();
+                physics_path!(self $p).velocity.replace(vel.into());
+            }
+
+            fn set_velocity_x(&mut self, x: f32) {
+                physics_path!(self $p).velocity.get_or_insert_default().x.get_or_insert_default().val = x;
+            }
+
+            fn set_velocity_y(&mut self, y: f32) {
+                physics_path!(self $p).velocity.get_or_insert_default().y.get_or_insert_default().val = y;
+            }
+
+            fn set_velocity_z(&mut self, z: f32) {
+                physics_path!(self $p).velocity.get_or_insert_default().z.get_or_insert_default().val = z;
+            }
+
+            fn set_rotation(&mut self, rot: impl Into<Rotator>) {
+                let rot: RotatorPartial = rot.into().into();
+                physics_path!(self $p).rotation.replace(rot.into());
+            }
+
+            fn set_rotation_pitch(&mut self, pitch: f32) {
+                physics_path!(self $p).rotation.get_or_insert_default().pitch.get_or_insert_default().val = pitch;
+            }
+
+            fn set_rotation_yaw(&mut self, yaw: f32) {
+                physics_path!(self $p).rotation.get_or_insert_default().yaw.get_or_insert_default().val = yaw;
+            }
+
+            fn set_rotation_roll(&mut self, roll: f32) {
+                physics_path!(self $p).rotation.get_or_insert_default().roll.get_or_insert_default().val = roll;
+            }
+
+            fn set_angular_velocity(&mut self, ang_vel: impl Into<Vector3>) {
+                let ang_vel: Vector3Partial = ang_vel.into().into();
+                physics_path!(self $p).angular_velocity.replace(ang_vel.into());
+            }
+
+            fn set_angular_velocity_x(&mut self, x: f32) {
+                physics_path!(self $p).angular_velocity.get_or_insert_default().x.get_or_insert_default().val = x;
+            }
+
+            fn set_angular_velocity_y(&mut self, y: f32) {
+                physics_path!(self $p).angular_velocity.get_or_insert_default().y.get_or_insert_default().val = y;
+            }
+
+            fn set_angular_velocity_z(&mut self, z: f32) {
+                physics_path!(self $p).angular_velocity.get_or_insert_default().z.get_or_insert_default().val = z;
+            }
+        }
+    };
 }
+
+desired_physics_ext!(DesiredPhysics; slf);
+desired_physics_ext!(DesiredBallState; physics);
+desired_physics_ext!(DesiredCarState; optional_physics);

From 8d851a1ec3be0b0f831811e94672159f98ddb9aa Mon Sep 17 00:00:00 2001
From: NicEastvillage <nicoesterby@gmail.com>
Date: Thu, 3 Jul 2025 22:29:05 +0200
Subject: [PATCH 8/9] Add PR suggestions

---
 rlbot/src/state_builder.rs | 26 +++++++++++---------------
 1 file changed, 11 insertions(+), 15 deletions(-)

diff --git a/rlbot/src/state_builder.rs b/rlbot/src/state_builder.rs
index b2650b7..2adb5cf 100644
--- a/rlbot/src/state_builder.rs
+++ b/rlbot/src/state_builder.rs
@@ -5,16 +5,16 @@ use rlbot_flat::flat::{DesiredBallState, DesiredCarState, DesiredGameState, Desi
 /// Example:
 /// ```rust
 /// use rlbot::state_builder::{DesiredCarStateExt, DesiredGameStateExt, DesiredPhysicsExt};
-/// use rlbot_flat::flat::{DesiredGameState, Vector3};
+/// use rlbot::flat::{DesiredGameState, Vector3};
 /// let mut dgs = DesiredGameState::default();
 /// dgs.mod_car(0, |c| {
 ///     c.set_location(Vector3::default());
 ///     c.set_boost(100.);
 /// });
-/// dgs.mod_balls(0..5, |_, b| {
+/// dgs.mod_balls((0..5).map(|i| (i, |b| {
 ///     b.set_location_z(0.);
 ///     b.set_velocity_z(0.);
-/// });
+/// })));
 /// ```
 pub trait DesiredGameStateExt {
 
@@ -31,8 +31,7 @@ pub trait DesiredGameStateExt {
 
     fn mod_cars(
         &mut self,
-        indices: impl IntoIterator<Item = usize>,
-        build: impl Fn(usize, &mut DesiredCarState),
+        build: impl IntoIterator<Item = (usize, impl Fn(&mut DesiredCarState))>,
     );
 
     fn mod_ball(
@@ -43,8 +42,7 @@ pub trait DesiredGameStateExt {
 
     fn mod_balls(
         &mut self,
-        indices: impl IntoIterator<Item = usize>,
-        build: impl Fn(usize, &mut DesiredBallState),
+        build: impl IntoIterator<Item = (usize, impl Fn(&mut DesiredBallState))>,
     );
 }
 
@@ -74,14 +72,13 @@ impl DesiredGameStateExt for DesiredGameState {
     /// Modify all desired cars.
     fn mod_cars(
         &mut self,
-        indices: impl IntoIterator<Item = usize>,
-        build: impl Fn(usize, &mut DesiredCarState),
+        build: impl IntoIterator<Item = (usize, impl Fn(&mut DesiredCarState))>,
     ) {
-        for i in indices {
+        for (i, func) in build {
             if self.car_states.len() <= i {
                 self.car_states.resize(i + 1, Default::default());
             }
-            build(i, &mut self.car_states[i]);
+            func(&mut self.car_states[i]);
         }
     }
 
@@ -100,14 +97,13 @@ impl DesiredGameStateExt for DesiredGameState {
     /// Modify all desired balls.
     fn mod_balls(
         &mut self,
-        indices: impl IntoIterator<Item = usize>,
-        build: impl Fn(usize, &mut DesiredBallState),
+        build: impl IntoIterator<Item = (usize, impl Fn(&mut DesiredBallState))>,
     ) {
-        for i in indices {
+        for (i, func) in build {
             if self.ball_states.len() <= i {
                 self.ball_states.resize(i + 1, Default::default());
             }
-            build(i, &mut self.ball_states[i]);
+            func(&mut self.ball_states[i]);
         }
     }
 }

From be2ec93e00ba40faf6c17913042afea91f31ff66 Mon Sep 17 00:00:00 2001
From: NicEastvillage <nicoesterby@gmail.com>
Date: Thu, 3 Jul 2025 22:30:44 +0200
Subject: [PATCH 9/9] fmt

---
 rlbot/src/state_builder.rs | 61 +++++++++++++-------------------------
 1 file changed, 20 insertions(+), 41 deletions(-)

diff --git a/rlbot/src/state_builder.rs b/rlbot/src/state_builder.rs
index 2adb5cf..6937ad8 100644
--- a/rlbot/src/state_builder.rs
+++ b/rlbot/src/state_builder.rs
@@ -1,4 +1,7 @@
-use rlbot_flat::flat::{DesiredBallState, DesiredCarState, DesiredGameState, DesiredMatchInfo, DesiredPhysics, Rotator, RotatorPartial, Vector3, Vector3Partial};
+use rlbot_flat::flat::{
+    DesiredBallState, DesiredCarState, DesiredGameState, DesiredMatchInfo, DesiredPhysics, Rotator,
+    RotatorPartial, Vector3, Vector3Partial,
+};
 
 /// Extension methods for easy construction of [DesiredGameState].
 ///
@@ -17,28 +20,13 @@ use rlbot_flat::flat::{DesiredBallState, DesiredCarState, DesiredGameState, Desi
 /// })));
 /// ```
 pub trait DesiredGameStateExt {
+    fn mod_match_info(&mut self, build: impl FnOnce(&mut DesiredMatchInfo));
 
-    fn mod_match_info(
-        &mut self,
-        build: impl FnOnce(&mut DesiredMatchInfo),
-    );
-
-    fn mod_car(
-        &mut self,
-        index: usize,
-        build: impl FnOnce(&mut DesiredCarState),
-    );
+    fn mod_car(&mut self, index: usize, build: impl FnOnce(&mut DesiredCarState));
 
-    fn mod_cars(
-        &mut self,
-        build: impl IntoIterator<Item = (usize, impl Fn(&mut DesiredCarState))>,
-    );
+    fn mod_cars(&mut self, build: impl IntoIterator<Item = (usize, impl Fn(&mut DesiredCarState))>);
 
-    fn mod_ball(
-        &mut self,
-        index: usize,
-        build: impl FnOnce(&mut DesiredBallState),
-    );
+    fn mod_ball(&mut self, index: usize, build: impl FnOnce(&mut DesiredBallState));
 
     fn mod_balls(
         &mut self,
@@ -48,21 +36,13 @@ pub trait DesiredGameStateExt {
 
 #[allow(dead_code)]
 impl DesiredGameStateExt for DesiredGameState {
-
     /// Modify the desired match info.
-    fn mod_match_info(
-        &mut self,
-        build: impl FnOnce(&mut DesiredMatchInfo),
-    ) {
+    fn mod_match_info(&mut self, build: impl FnOnce(&mut DesiredMatchInfo)) {
         build(self.match_info.get_or_insert_default());
     }
 
     /// Modify the desired car at the given index.
-    fn mod_car(
-        &mut self,
-        index: usize,
-        build: impl FnOnce(&mut DesiredCarState),
-    ) {
+    fn mod_car(&mut self, index: usize, build: impl FnOnce(&mut DesiredCarState)) {
         if self.car_states.len() <= index {
             self.car_states.resize(index + 1, Default::default());
         }
@@ -83,11 +63,7 @@ impl DesiredGameStateExt for DesiredGameState {
     }
 
     /// Modify the desired ball at the given index.
-    fn mod_ball(
-        &mut self,
-        index: usize,
-        build: impl FnOnce(&mut DesiredBallState),
-    ) {
+    fn mod_ball(&mut self, index: usize, build: impl FnOnce(&mut DesiredBallState)) {
         if self.ball_states.len() <= index {
             self.ball_states.resize(index + 1, Default::default());
         }
@@ -116,7 +92,6 @@ pub trait DesiredMatchInfoExt {
 
 #[allow(dead_code)]
 impl DesiredMatchInfoExt for DesiredMatchInfo {
-
     /// Set the desired world gravity z.
     fn set_gravity_z(&mut self, gravity_z: f32) {
         self.world_gravity_z.get_or_insert_default().val = gravity_z;
@@ -135,7 +110,6 @@ pub trait DesiredCarStateExt {
 }
 
 impl DesiredCarStateExt for DesiredCarState {
-
     /// Set the boost amount of this car.
     fn set_boost(&mut self, amount: f32) {
         self.boost_amount.get_or_insert_default().val = amount;
@@ -153,7 +127,6 @@ pub trait DesiredBallStateExt {
 }
 
 impl DesiredBallStateExt for DesiredBallState {
-
     /// Modify the physics of this ball.
     fn mod_physics(&mut self, build: impl FnOnce(&mut DesiredPhysics)) {
         build(&mut self.physics)
@@ -181,9 +154,15 @@ pub trait DesiredPhysicsExt {
 }
 
 macro_rules! physics_path {
-    ( $self:ident slf ) => { $self };
-    ( $self:ident physics) => { $self.physics };
-    ( $self:ident optional_physics) => { $self.physics.get_or_insert_default() };
+    ( $self:ident slf ) => {
+        $self
+    };
+    ( $self:ident physics) => {
+        $self.physics
+    };
+    ( $self:ident optional_physics) => {
+        $self.physics.get_or_insert_default()
+    };
 }
 
 macro_rules! desired_physics_ext {