Imu traits take 2

firmware/src/imu/drivers/bmi160/mod.rs

@@ -2,8 +2,8 @@ mod math;
 
 use self::math::{discrete_to_radians, GyroFsr};
 use crate::aliases::I2c;
-use crate::imu::{FusedImu, Quat};
-use crate::utils;
+use crate::imu::{FusedData, Imu, Quat};
+use crate::utils::{self, nb2a};
 
 use ::bmi160::{AccelerometerPowerMode, GyroscopePowerMode, SensorSelector};
 use defmt::{debug, trace};
@@ -74,14 +74,8 @@ impl<I: I2c> Bmi160<I> {
 		// Map converts from tuple -> struct
 		.map_err(|(i2c, error)| InitError { i2c, error })
 	}
-}
-
-impl<I: I2c> FusedImu for Bmi160<I> {
-	type Error = BmiError<I>;
 
-	const IMU_TYPE: ImuType = ImuType::Bmi160;
-
-	fn quat(&mut self) -> nb::Result<Quat, Self::Error> {
+	fn quat(&mut self) -> nb::Result<Quat, <Self as Imu>::Error> {
 		let data = self.driver.data(SensorSelector::new().gyro())?;
 		let gyro_vel_euler = data.gyro.unwrap();
 
@@ -99,10 +93,21 @@ impl<I: I2c> FusedImu for Bmi160<I> {
 	}
 }
 
+impl<I: I2c> Imu for Bmi160<I> {
+	type Error = BmiError<I>;
+	type Data = FusedData; //TODO: Stop pretending we are fused.
+
+	const IMU_TYPE: ImuType = ImuType::Bmi160;
+
+	async fn next_data(&mut self) -> Result<Self::Data, Self::Error> {
+		nb2a(|| self.quat()).await.map(|quat| FusedData { q: quat })
+	}
+}
+
 #[allow(dead_code)]
 pub fn new_imu(
 	i2c: impl crate::aliases::I2c,
 	delay: &mut impl DelayMs<u32>,
-) -> impl crate::imu::FusedImu {
+) -> impl Imu<Data = FusedData> {
 	Bmi160::new(i2c, delay).expect("Failed to initialize BMI160")
 }

firmware/src/imu/drivers/mpu6050.rs

@@ -1,6 +1,6 @@
 use crate::aliases::I2c;
-use crate::imu::{FusedImu, Quat};
-use crate::utils;
+use crate::imu::{FusedData, Imu, Quat};
+use crate::utils::{self, nb2a};
 
 use defmt::{debug, trace, warn};
 use embedded_hal::blocking::delay::DelayMs;
@@ -47,14 +47,8 @@ impl<I: I2c> Mpu6050<I> {
 		// Map converts from tuple -> struct
 		.map_err(|(i2c, error)| InitError { i2c, error })
 	}
-}
-
-impl<I: I2c> FusedImu for Mpu6050<I> {
-	type Error = mpu6050_dmp::error::Error<I>;
 
-	const IMU_TYPE: ImuType = ImuType::Mpu6050;
-
-	fn quat(&mut self) -> nb::Result<Quat, Self::Error> {
+	fn quat(&mut self) -> nb::Result<Quat, <Self as Imu>::Error> {
 		if self.mpu.get_fifo_count()? >= 28 {
 			let data = self.mpu.read_fifo(&mut self.fifo_buf)?;
 			let opt = data.get(..16);
@@ -73,10 +67,21 @@ impl<I: I2c> FusedImu for Mpu6050<I> {
 	}
 }
 
+impl<I: I2c> Imu for Mpu6050<I> {
+	type Error = mpu6050_dmp::error::Error<I>;
+	type Data = FusedData;
+
+	const IMU_TYPE: ImuType = ImuType::Mpu6050;
+
+	async fn next_data(&mut self) -> Result<Self::Data, Self::Error> {
+		nb2a(|| self.quat()).await.map(|quat| FusedData { q: quat })
+	}
+}
+
 #[allow(dead_code)]
 pub fn new_imu(
 	i2c: impl crate::aliases::I2c,
 	delay: &mut impl DelayMs<u32>,
-) -> impl crate::imu::FusedImu {
+) -> impl Imu<Data = FusedData> {
 	Mpu6050::new(i2c, delay).expect("Failed to create mpu6050")
 }

firmware/src/imu/drivers/stubbed.rs

@@ -1,4 +1,4 @@
-use crate::imu::{FusedImu, Quat};
+use crate::imu::{FusedData, Imu, Quat};
 
 use defmt::debug;
 use embedded_hal::blocking::delay::DelayMs;
@@ -7,21 +7,24 @@ use firmware_protocol::ImuType;
 /// Fakes an IMU for easier testing.
 struct FakeImu;
 
-impl FusedImu for FakeImu {
+impl Imu for FakeImu {
 	type Error = ();
+	type Data = FusedData;
 
 	const IMU_TYPE: ImuType = ImuType::Unknown(0xFF);
 
-	fn quat(&mut self) -> nb::Result<Quat, Self::Error> {
-		Ok(Quat::identity())
+	async fn next_data(&mut self) -> Result<Self::Data, Self::Error> {
+		Ok(FusedData {
+			q: Quat::identity(),
+		})
 	}
 }
 
 #[allow(dead_code)]
 pub fn new_imu(
 	_i2c: impl crate::aliases::I2c,
 	_delay: &mut impl DelayMs<u32>,
-) -> impl crate::imu::FusedImu {
+) -> impl Imu<Data = FusedData> {
 	debug!("Created FakeImu");
 	FakeImu
 }

firmware/src/imu/mod.rs

@@ -3,48 +3,77 @@ mod fusion;
 
 use defmt::{debug, info, trace, warn};
 use embassy_executor::task;
-use embassy_futures::yield_now;
 use firmware_protocol::ImuType;
 
 use crate::{
 	aliases::ඞ::{DelayConcrete, I2cConcrete},
-	utils::{nb2a, Unreliable},
+	utils::Unreliable,
 };
 
 pub type Quat = nalgebra::UnitQuaternion<f32>;
+pub type Accel = nalgebra::Vector3<f32>;
+pub type Gyro = nalgebra::Vector3<f32>;
 
-pub trait FusedImu {
-	type Error: core::fmt::Debug;
+pub struct UnfusedData {
+	pub accel: Accel,
+	pub gyro: Gyro,
+}
+
+pub struct FusedData {
+	pub q: Quat,
+}
+
+/// Represents a sensor fusion algorithm that will take an imu's `UnfusedData` and do math to turn
+/// it into `FusedData`, suitable for use as orientation.
+pub trait Fuser {
+	// TODO: Does a one-in, one-out api here work? Should we reuse a standard trait like a
+	// sink/stream/iterator?
+	// Note: Intentionally not async rn, this should only be doing math, not io or any internal
+	// awaiting.
+	fn process(&mut self, unfused: &UnfusedData) -> FusedData;
+}
+
+pub trait Imu {
+	type Error: core::fmt::Debug; // TODO: Maybe use defmt instead?
+	/// The data that the imu outputs.
+	type Data;
 
 	const IMU_TYPE: ImuType;
-	// TODO: This should be async
-	fn quat(&mut self) -> nb::Result<Quat, Self::Error>;
+	/// Performs IO to get the next data from the imu.
+	async fn next_data(&mut self) -> Result<Self::Data, Self::Error>;
+}
+
+pub struct FusedImu<I: Imu, F: Fuser> {
+	pub imu: I,
+	pub fuser: F,
+}
+impl<I: Imu<Data = UnfusedData>, F: Fuser> Imu for FusedImu<I, F> {
+	type Error = I::Error;
+	type Data = FusedData;
+
+	const IMU_TYPE: ImuType = I::IMU_TYPE;
+
+	async fn next_data(&mut self) -> Result<Self::Data, Self::Error> {
+		let unfused = self.imu.next_data().await?;
+		Ok(self.fuser.process(&unfused))
+	}
 }
 
 /// Gets data from the IMU
 #[task]
 pub async fn imu_task(
 	quat_signal: &'static Unreliable<Quat>,
 	i2c: I2cConcrete<'static>,
-	delay: DelayConcrete,
-) -> ! {
-	imu_task_inner(quat_signal, i2c, delay).await
-}
-
-/// Same as [`imu_task()`] but this version's arguments are type erased behind impl
-/// Trait to avoid accidentally accessing concrete behavior.
-async fn imu_task_inner(
-	quat_signal: &Unreliable<Quat>,
-	i2c: impl crate::aliases::I2c,
-	mut delay: impl crate::aliases::Delay,
+	mut delay: DelayConcrete,
 ) -> ! {
 	debug!("Imu task");
+
 	let mut imu = new_imu(i2c, &mut delay);
 	info!("Initialized IMU!");
 
 	loop {
-		let q = match nb2a(|| imu.quat()).await {
-			Ok(q) => q,
+		let q = match imu.next_data().await {
+			Ok(q) => q.q,
 			Err(err) => {
 				warn!("Error in IMU: {}", defmt::Debug2Format(&err));
 				continue;
@@ -58,14 +87,13 @@ async fn imu_task_inner(
 			q.coords.w
 		);
 		quat_signal.signal(q);
-		yield_now().await // Yield to ensure fairness
 	}
 }
 
 fn new_imu(
 	i2c: impl crate::aliases::I2c,
 	delay: &mut impl crate::aliases::Delay,
-) -> impl FusedImu {
+) -> impl Imu<Data = FusedData> {
 	use crate::imu::drivers as d;
 
 	#[cfg(feature = "imu-bmi160")]

firmware/src/main.rs

@@ -8,6 +8,8 @@
 #![feature(alloc_error_handler)]
 // We want to do some floating point math at compile time
 #![feature(const_fn_floating_point_arithmetic)]
+// We need async traits for efficient yet generic IMUs
+#![feature(async_fn_in_trait)]
 #![deny(unsafe_op_in_unsafe_fn)]
 
 load_dotenv::try_load_dotenv!();