Skip to main content

ozton_motion/basic/
drive.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    task::{Context, Poll},
5    time::{Duration, Instant},
6};
7
8use ozton_control::{
9    Tolerances,
10    loops::{AngularPid, Feedback, Pid},
11};
12use ozton_drivetrain::{Drivetrain, model::Arcade};
13use ozton_tracking::{TracksForwardTravel, TracksHeading, TracksVelocity};
14use vexide::{
15    math::Angle,
16    time::{Sleep, sleep},
17};
18
19pub(crate) struct DriveState {
20    pub sleep: Sleep,
21    pub initial_forward_travel: f64,
22    pub start_time: Instant,
23    pub prev_time: Instant,
24    pub linear_settled: bool,
25    pub angular_settled: bool,
26}
27
28/// Drives the robot forward or backwards for a distance at a given heading.
29#[must_use = "futures do nothing unless you `.await` or poll them"]
30pub struct DriveFuture<'a, M, L, A, T>
31where
32    M: Arcade,
33    L: Feedback<State = f64, Signal = f64> + Unpin,
34    A: Feedback<State = Angle, Signal = f64> + Unpin,
35    T: TracksForwardTravel + TracksHeading + TracksVelocity,
36{
37    pub(crate) target_distance: f64,
38    pub(crate) target_heading: Angle,
39    pub(crate) timeout: Option<Duration>,
40    pub(crate) linear_tolerances: Tolerances,
41    pub(crate) angular_tolerances: Tolerances,
42    pub(crate) linear_controller: L,
43    pub(crate) angular_controller: A,
44    pub(crate) drivetrain: &'a mut Drivetrain<M, T>,
45
46    /// Internal future state ("local variables").
47    pub(crate) state: Option<DriveState>,
48}
49
50// MARK: Future Poll
51
52impl<M, L, A, T> Future for DriveFuture<'_, M, L, A, T>
53where
54    M: Arcade,
55    L: Feedback<State = f64, Signal = f64> + Unpin,
56    A: Feedback<State = Angle, Signal = f64> + Unpin,
57    T: TracksForwardTravel + TracksHeading + TracksVelocity,
58{
59    type Output = ();
60
61    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
62        let this = self.get_mut();
63        let state = this.state.get_or_insert_with(|| {
64            let now = Instant::now();
65            DriveState {
66                sleep: sleep(Duration::from_millis(5)),
67                initial_forward_travel: this.drivetrain.tracking.forward_travel(),
68                start_time: now,
69                prev_time: now,
70                linear_settled: false,
71                angular_settled: false,
72            }
73        });
74
75        if Pin::new(&mut state.sleep).poll(cx).is_pending() {
76            return Poll::Pending;
77        }
78
79        let dt = state.prev_time.elapsed();
80
81        let forward_travel = this.drivetrain.tracking.forward_travel();
82        let heading = this.drivetrain.tracking.heading();
83
84        let linear_error = (this.target_distance + state.initial_forward_travel) - forward_travel;
85        let angular_error = (this.target_heading - heading).wrapped_half();
86
87        // println!("{}", linear_error);
88
89        if this
90            .linear_tolerances
91            .check(linear_error, this.drivetrain.tracking.linear_velocity())
92        {
93            state.linear_settled = true;
94        }
95        if this.angular_tolerances.check(
96            angular_error.as_radians(),
97            this.drivetrain.tracking.angular_velocity(),
98        ) {
99            state.angular_settled = true;
100        }
101
102        if (state.linear_settled && state.angular_settled)
103            || this
104                .timeout
105                .is_some_and(|timeout| state.start_time.elapsed() > timeout)
106        {
107            drop(this.drivetrain.model.drive_arcade(0.0, 0.0));
108            return Poll::Ready(());
109        }
110
111        let linear_output = this.linear_controller.update(
112            forward_travel,
113            this.target_distance + state.initial_forward_travel,
114            dt,
115        );
116        let angular_output = this
117            .angular_controller
118            .update(heading, this.target_heading, dt);
119
120        drop(
121            this.drivetrain
122                .model
123                .drive_arcade(linear_output, angular_output),
124        );
125
126        state.sleep = sleep(Duration::from_millis(5));
127        state.prev_time = Instant::now();
128
129        cx.waker().wake_by_ref();
130        Poll::Pending
131    }
132}
133
134// MARK: Generic Modifiers
135
136impl<M, L, A, T> DriveFuture<'_, M, L, A, T>
137where
138    M: Arcade,
139    L: Feedback<State = f64, Signal = f64> + Unpin,
140    A: Feedback<State = Angle, Signal = f64> + Unpin,
141    T: TracksForwardTravel + TracksHeading + TracksVelocity,
142{
143    /// Modifies this motion's linear feedback controller.
144    pub fn with_linear_controller(&mut self, controller: L) -> &mut Self {
145        self.linear_controller = controller;
146        self
147    }
148
149    /// Modifies this motion's angular feedback controller.
150    pub fn with_angular_controller(&mut self, controller: A) -> &mut Self {
151        self.angular_controller = controller;
152        self
153    }
154
155    /// Modifies this motion's timeout duration.
156    pub const fn with_timeout(&mut self, timeout: Duration) -> &mut Self {
157        self.timeout = Some(timeout);
158        self
159    }
160
161    /// Removes this motion's timeout duration.
162    pub const fn without_timeout(&mut self) -> &mut Self {
163        self.timeout = None;
164        self
165    }
166
167    /// Modifies this motion's linear tolerances.
168    pub const fn with_linear_tolerances(&mut self, tolerances: Tolerances) -> &mut Self {
169        self.linear_tolerances = tolerances;
170        self
171    }
172
173    /// Modifies this motion's linear error tolerance.
174    pub const fn with_linear_error_tolerance(&mut self, tolerance: f64) -> &mut Self {
175        self.linear_tolerances.error_tolerance = Some(tolerance);
176        self
177    }
178
179    /// Removes this motion's linear error tolerance.
180    pub const fn without_linear_error_tolerance(&mut self) -> &mut Self {
181        self.linear_tolerances.error_tolerance = None;
182        self
183    }
184
185    /// Modifies this motion's linear velocity tolerance.
186    pub const fn with_linear_velocity_tolerance(&mut self, tolerance: f64) -> &mut Self {
187        self.linear_tolerances.velocity_tolerance = Some(tolerance);
188        self
189    }
190
191    /// Removes this motion's linear velocity tolerance.
192    pub const fn without_linear_velocity_tolerance(&mut self) -> &mut Self {
193        self.linear_tolerances.velocity_tolerance = None;
194        self
195    }
196
197    /// Modifies this motion's linear tolerance duration.
198    pub const fn with_linear_tolerance_duration(&mut self, duration: Duration) -> &mut Self {
199        self.linear_tolerances.duration = Some(duration);
200        self
201    }
202
203    /// Removes this motion's linear tolerance duration.
204    pub const fn without_linear_tolerance_duration(&mut self) -> &mut Self {
205        self.linear_tolerances.duration = None;
206        self
207    }
208
209    /// Removes this motion's linear and angular tolerance durations.
210    pub const fn without_tolerance_duration(&mut self) -> &mut Self {
211        self.linear_tolerances.duration = None;
212        self.angular_tolerances.duration = None;
213        self
214    }
215
216    /// Modifies this motion's angular tolerances.
217    pub const fn with_angular_tolerances(&mut self, tolerances: Tolerances) -> &mut Self {
218        self.angular_tolerances = tolerances;
219        self
220    }
221
222    /// Modifies this motion's angular error tolerance.
223    pub const fn with_angular_error_tolerance(&mut self, tolerance: f64) -> &mut Self {
224        self.angular_tolerances.error_tolerance = Some(tolerance);
225        self
226    }
227
228    /// Removes this motion's angular error tolerance.
229    pub const fn without_angular_error_tolerance(&mut self) -> &mut Self {
230        self.angular_tolerances.error_tolerance = None;
231        self
232    }
233
234    /// Modifies this motion's angular velocity tolerance.
235    pub const fn with_angular_velocity_tolerance(&mut self, tolerance: f64) -> &mut Self {
236        self.angular_tolerances.velocity_tolerance = Some(tolerance);
237        self
238    }
239
240    /// Removes this motion's angular velocity tolerance.
241    pub const fn without_angular_velocity_tolerance(&mut self) -> &mut Self {
242        self.angular_tolerances.velocity_tolerance = None;
243        self
244    }
245
246    /// Modifies this motion's angular tolerance duration.
247    pub const fn with_angular_tolerance_duration(&mut self, duration: Duration) -> &mut Self {
248        self.angular_tolerances.duration = Some(duration);
249        self
250    }
251
252    /// Removes this motion's angular tolerance duration.
253    pub const fn without_angular_tolerance_duration(&mut self) -> &mut Self {
254        self.angular_tolerances.duration = None;
255        self
256    }
257}
258
259// MARK: Linear PID Modifiers
260
261impl<M, A, T> DriveFuture<'_, M, Pid, A, T>
262where
263    M: Arcade,
264    A: Feedback<State = Angle, Signal = f64> + Unpin,
265    T: TracksForwardTravel + TracksHeading + TracksVelocity,
266{
267    /// Modifies this motion's linear PID gains.
268    pub const fn with_linear_gains(&mut self, kp: f64, ki: f64, kd: f64) -> &mut Self {
269        self.linear_controller.set_gains(kp, ki, kd);
270        self
271    }
272
273    /// Modifies this motion's linear proportional gain (`kp`).
274    pub const fn with_linear_kp(&mut self, kp: f64) -> &mut Self {
275        self.linear_controller.set_kp(kp);
276        self
277    }
278
279    /// Modifies this motion's linear integral gain (`ki`).
280    pub const fn with_linear_ki(&mut self, ki: f64) -> &mut Self {
281        self.linear_controller.set_ki(ki);
282        self
283    }
284
285    /// Modifies this motion's linear derivative gain (`kd`).
286    pub const fn with_linear_kd(&mut self, kd: f64) -> &mut Self {
287        self.linear_controller.set_kd(kd);
288        self
289    }
290
291    /// Modifies this motion's linear integration range.
292    pub const fn with_linear_integration_range(&mut self, integration_range: f64) -> &mut Self {
293        self.linear_controller
294            .set_integration_range(Some(integration_range));
295        self
296    }
297
298    /// Removes this motion's linear integration range.
299    pub const fn without_linear_integration_range(&mut self) -> &mut Self {
300        self.linear_controller.set_integration_range(None);
301        self
302    }
303
304    /// Modifies this motion's linear output limit.
305    pub const fn with_linear_output_limit(&mut self, limit: f64) -> &mut Self {
306        self.linear_controller.set_output_limit(Some(limit));
307        self
308    }
309
310    /// Removes this motion's linear output limit.
311    pub const fn without_linear_output_limit(&mut self) -> &mut Self {
312        self.linear_controller.set_output_limit(None);
313        self
314    }
315}
316
317// MARK: Angular PID Modifiers
318
319impl<M, L, T> DriveFuture<'_, M, L, AngularPid, T>
320where
321    M: Arcade,
322    L: Feedback<State = f64, Signal = f64> + Unpin,
323    T: TracksForwardTravel + TracksHeading + TracksVelocity,
324{
325    /// Modifies this motion's angular PID gains.
326    pub const fn with_angular_gains(&mut self, kp: f64, ki: f64, kd: f64) -> &mut Self {
327        self.angular_controller.set_gains(kp, ki, kd);
328        self
329    }
330
331    /// Modifies this motion's angular proportional gain (`kp`).
332    pub const fn with_angular_kp(&mut self, kp: f64) -> &mut Self {
333        self.angular_controller.set_kp(kp);
334        self
335    }
336
337    /// Modifies this motion's angular integral gain (`ki`).
338    pub const fn with_angular_ki(&mut self, ki: f64) -> &mut Self {
339        self.angular_controller.set_ki(ki);
340        self
341    }
342
343    /// Modifies this motion's angular derivative gain (`kd`).
344    pub const fn with_angular_kd(&mut self, kd: f64) -> &mut Self {
345        self.angular_controller.set_kd(kd);
346        self
347    }
348
349    /// Modifies this motion's angular integration range.
350    pub const fn with_angular_integration_range(&mut self, integration_range: Angle) -> &mut Self {
351        self.angular_controller
352            .set_integration_range(Some(integration_range));
353        self
354    }
355
356    /// Modifies this motion's angular output limit.
357    pub const fn with_angular_output_limit(&mut self, limit: f64) -> &mut Self {
358        self.angular_controller.set_output_limit(Some(limit));
359        self
360    }
361
362    /// Removes this motion's angular integration range.
363    pub const fn without_angular_integration_range(&mut self) -> &mut Self {
364        self.angular_controller.set_integration_range(None);
365        self
366    }
367
368    /// Removes this motion's angular output limit.
369    pub const fn without_angular_output_limit(&mut self) -> &mut Self {
370        self.angular_controller.set_output_limit(None);
371        self
372    }
373}