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#[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 pub(crate) state: Option<DriveState>,
48}
49
50impl<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 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
134impl<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 pub fn with_linear_controller(&mut self, controller: L) -> &mut Self {
145 self.linear_controller = controller;
146 self
147 }
148
149 pub fn with_angular_controller(&mut self, controller: A) -> &mut Self {
151 self.angular_controller = controller;
152 self
153 }
154
155 pub const fn with_timeout(&mut self, timeout: Duration) -> &mut Self {
157 self.timeout = Some(timeout);
158 self
159 }
160
161 pub const fn without_timeout(&mut self) -> &mut Self {
163 self.timeout = None;
164 self
165 }
166
167 pub const fn with_linear_tolerances(&mut self, tolerances: Tolerances) -> &mut Self {
169 self.linear_tolerances = tolerances;
170 self
171 }
172
173 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 pub const fn without_linear_error_tolerance(&mut self) -> &mut Self {
181 self.linear_tolerances.error_tolerance = None;
182 self
183 }
184
185 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 pub const fn without_linear_velocity_tolerance(&mut self) -> &mut Self {
193 self.linear_tolerances.velocity_tolerance = None;
194 self
195 }
196
197 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 pub const fn without_linear_tolerance_duration(&mut self) -> &mut Self {
205 self.linear_tolerances.duration = None;
206 self
207 }
208
209 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 pub const fn with_angular_tolerances(&mut self, tolerances: Tolerances) -> &mut Self {
218 self.angular_tolerances = tolerances;
219 self
220 }
221
222 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 pub const fn without_angular_error_tolerance(&mut self) -> &mut Self {
230 self.angular_tolerances.error_tolerance = None;
231 self
232 }
233
234 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 pub const fn without_angular_velocity_tolerance(&mut self) -> &mut Self {
242 self.angular_tolerances.velocity_tolerance = None;
243 self
244 }
245
246 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 pub const fn without_angular_tolerance_duration(&mut self) -> &mut Self {
254 self.angular_tolerances.duration = None;
255 self
256 }
257}
258
259impl<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 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 pub const fn with_linear_kp(&mut self, kp: f64) -> &mut Self {
275 self.linear_controller.set_kp(kp);
276 self
277 }
278
279 pub const fn with_linear_ki(&mut self, ki: f64) -> &mut Self {
281 self.linear_controller.set_ki(ki);
282 self
283 }
284
285 pub const fn with_linear_kd(&mut self, kd: f64) -> &mut Self {
287 self.linear_controller.set_kd(kd);
288 self
289 }
290
291 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 pub const fn without_linear_integration_range(&mut self) -> &mut Self {
300 self.linear_controller.set_integration_range(None);
301 self
302 }
303
304 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 pub const fn without_linear_output_limit(&mut self) -> &mut Self {
312 self.linear_controller.set_output_limit(None);
313 self
314 }
315}
316
317impl<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 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 pub const fn with_angular_kp(&mut self, kp: f64) -> &mut Self {
333 self.angular_controller.set_kp(kp);
334 self
335 }
336
337 pub const fn with_angular_ki(&mut self, ki: f64) -> &mut Self {
339 self.angular_controller.set_ki(ki);
340 self
341 }
342
343 pub const fn with_angular_kd(&mut self, kd: f64) -> &mut Self {
345 self.angular_controller.set_kd(kd);
346 self
347 }
348
349 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 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 pub const fn without_angular_integration_range(&mut self) -> &mut Self {
364 self.angular_controller.set_integration_range(None);
365 self
366 }
367
368 pub const fn without_angular_output_limit(&mut self) -> &mut Self {
370 self.angular_controller.set_output_limit(None);
371 self
372 }
373}