oled: split out the ssd1306 driver into its own module, and reimplement the oled design with separate rendering and drawing tasks

src/animation.rs

@@ -84,7 +84,7 @@ impl Animation {
         }
     }
 
-    pub async fn apply<S: AnimationActor + Debug>(&self, sfc: &mut S) {
+    pub async fn apply<S: AnimationActor>(&self, sfc: &mut S) {
         let from = if let Some(val) = self.from {
             val
         } else {
@@ -100,7 +100,7 @@ impl Animation {
             return;
         }
         let step_time = self.duration / steps.into();
-        trace!("fade={self:?} steps={steps} time={step_time} sfc={sfc:?}");
+        trace!("fade={self:?} steps={steps} time={step_time}");
         if from > to {
             let range = (to..=from).rev();
             for opacity in range {

src/bin/main.rs

@@ -103,7 +103,7 @@ async fn main(spawner: Spawner) {
     {
         use esp_hal::i2c::master::{Config, I2c};
 
-        let mut rst = Output::new(peripherals.GPIO21, esp_hal::gpio::Level::Low, OutputConfig::default());
+        let rst = Output::new(peripherals.GPIO21, esp_hal::gpio::Level::Low, OutputConfig::default());
         let i2c = I2c::new(
             peripherals.I2C0, 
             Config::default().with_frequency(Rate::from_khz(400))

src/display.rs

@@ -17,13 +17,13 @@ pub struct BikeOutput<T: SmartLedsWriteAsync> {
     controls: DisplayControls
 }
 
-impl<T:SmartLedsWriteAsync> GammaCorrected for BikeOutput<T> where T::Color: PixelBlend<Rgb<u8>> + PixelFormat + Debug + WithGamma, T::Error: Debug {
+impl<T:SmartLedsWriteAsync> GammaCorrected for BikeOutput<T> where T::Color: PixelBlend<Rgb<u8>> + PixelFormat + Debug + WithGamma + Default + Clone, T::Error: Debug {
     fn set_gamma(&mut self, gamma: GammaCurve) {
         self.writer.controls().set_gamma(gamma);
     }
 }
 
-impl<T: SmartLedsWriteAsync> BikeOutput<T> where T::Color: PixelBlend<Rgb<u8>> + PixelFormat + WithGamma + 'static,  T::Error: core::fmt::Debug {
+impl<T: SmartLedsWriteAsync> BikeOutput<T> where T::Color: PixelBlend<Rgb<u8>> + PixelFormat + WithGamma + 'static + Default + Clone + Copy,  T::Error: core::fmt::Debug {
     pub fn new(target: T, max_mw: u32, controls: DisplayControls) -> Self {
         Self {
             pixbuf: [Default::default(); NUM_PIXELS],
@@ -37,7 +37,7 @@ impl<T: SmartLedsWriteAsync> BikeOutput<T> where T::Color: PixelBlend<Rgb<u8>> +
     }
 }
 
-impl<'a, T: SmartLedsWriteAsync + 'a> OutputAsync<'a, SegmentSpace> for BikeOutput<T> where T::Color: PixelBlend<Rgb<u8>> + Debug + 'static + AsMilliwatts + PixelFormat + WithGamma, [T::Color; NUM_PIXELS]: AsMilliwatts + WithGamma + Copy,  T::Error: core::fmt::Debug {
+impl<'a, T: SmartLedsWriteAsync + 'a> OutputAsync<'a, SegmentSpace> for BikeOutput<T> where T::Color: PixelSink<T::Color> + PixelBlend<Rgb<u8>> + Default + Clone + Debug + 'static + AsMilliwatts + PixelFormat + WithGamma, [T::Color; NUM_PIXELS]: AsMilliwatts + WithGamma + Copy,  T::Error: core::fmt::Debug {
     async fn commit_async(&mut self) -> Result<(), T::Error> {
         self.writer.controls().set_brightness(self.controls.brightness());
         self.writer.controls().set_on(self.controls.is_on());
@@ -45,7 +45,7 @@ impl<'a, T: SmartLedsWriteAsync + 'a> OutputAsync<'a, SegmentSpace> for BikeOutp
         self.writer.write(&self.pixbuf).await
     }
     
-    type HardwarePixel = T::Color;
+    //type HardwarePixel = T::Color;
     type Error = T::Error;
     type Controls = DisplayControls;
     
@@ -54,10 +54,10 @@ impl<'a, T: SmartLedsWriteAsync + 'a> OutputAsync<'a, SegmentSpace> for BikeOutp
     }
 }
 
-impl<'a, T: SmartLedsWriteAsync> Sample<'a, SegmentSpace> for BikeOutput<T> where T::Color: Debug + PixelFormat + 'a, [T::Color; NUM_PIXELS]: Sample<'a, SegmentSpace, Output = T::Color> {
+impl<'a, T: SmartLedsWriteAsync> Sample<'a, SegmentSpace> for BikeOutput<T> where T::Color: 'a + Debug + PixelFormat, [T::Color; NUM_PIXELS]: Sample<'a, SegmentSpace, Output = T::Color> {
     type Output = T::Color;
 
-    fn sample(&mut self, rect: &Rectangle<SegmentSpace>) -> impl Iterator<Item = (Coordinates<SegmentSpace>, &'a mut T::Color)> {
+    fn sample(&mut self, rect: &Rectangle<SegmentSpace>) -> impl Iterator<Item = (Coordinates<SegmentSpace>, &'a mut Self::Output)> {
         self.pixbuf.sample(rect)
     }
 }

src/graphics/mod.rs

@@ -0,0 +1 @@
+pub mod ssd1306;

src/graphics/ssd1306.rs

@@ -0,0 +1,251 @@
+#![cfg(feature="oled")]
+use core::{cmp::min, sync::atomic::{AtomicBool, AtomicU8}};
+
+use alloc::sync::Arc;
+use display_interface::DisplayError;
+use embedded_graphics::prelude::*;
+use esp_hal::rng;
+use figments::{liber8tion::{interpolate::Fract8, noise}, mappings::embedded_graphics::Matrix2DSpace, prelude::*};
+use embedded_graphics::pixelcolor::BinaryColor;
+use figments::pixels::PixelSink;
+use figments_render::{gamma::GammaCurve, output::{Brightness, GammaCorrected, NullControls, OutputAsync}};
+use log::*;
+
+pub struct SsdOutput([u8; 128 * 64 / 8], ssd1306::Ssd1306Async<ssd1306::prelude::I2CInterface<esp_hal::i2c::master::I2c<'static, esp_hal::Async>>, ssd1306::prelude::DisplaySize128x64, ssd1306::mode::BasicMode>, SsdControls);
+
+#[derive(Copy, Clone)]
+pub struct SsdPixel(*mut u8, u8, Coordinates<Matrix2DSpace>);
+
+impl SsdPixel {
+    fn set_pixel(&self, value: BinaryColor) {
+        if value.is_on() {
+            unsafe { *self.0 |= 1 << self.1 };
+        } else {
+            unsafe { *self.0 &= !(1 << self.1) };
+        }
+    }
+
+    fn get_pixel(&self) -> BinaryColor {
+        if unsafe { *self.0 >> self.1 & 0x01 == 0x01 } {
+            BinaryColor::On
+        } else {
+            BinaryColor::Off
+        }
+    }
+}
+
+impl PixelSink<BinaryColor> for SsdPixel {
+    fn set(&mut self, pixel: &BinaryColor) {
+        self.set_pixel(*pixel);
+    }
+}
+
+impl PixelBlend<BinaryColor> for SsdPixel {
+    fn blend_pixel(self, overlay: BinaryColor, opacity: Fract8) -> Self {
+        let scale = 32;
+        let x = self.2.x * scale;
+        let y = self.2.y * scale;
+        let stiple_idx = noise::inoise8(x as i16, y as i16);
+        if opacity >= stiple_idx {
+            self.set_pixel(overlay);
+        }
+        self
+    }
+
+    fn multiply(self, overlay: BinaryColor) -> Self {
+        if overlay != self.get_pixel() {
+            self.set_pixel(overlay);
+        }
+        self
+    }
+}
+
+static mut DUMMY: u8 = 0;
+
+impl Default for SsdPixel {
+    fn default() -> Self {
+        Self (
+            unsafe { &mut DUMMY },
+            0,
+            Coordinates::new(0, 0)
+        )
+    }
+}
+
+pub struct SsdSampler<'a> {
+    buf: &'a mut SsdOutput,
+    rect: figments::prelude::Rectangle<Matrix2DSpace>,
+    pos: figments::prelude::Coordinates<Matrix2DSpace>,
+    pixel: Option<SsdPixel>
+}
+
+impl<'a> Iterator for SsdSampler<'a> {
+    type Item = (Coordinates<Matrix2DSpace>, &'a mut SsdPixel);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.pos.x > min(127, self.rect.right()) {
+            self.pos.x = self.rect.left();
+            self.pos.y += 1;
+        }
+
+        if self.pos.y > min(63, self.rect.bottom()) {
+            return None;
+        }
+
+        self.pixel = self.buf.get_pixel(self.pos);
+        let old_pos = self.pos;
+        self.pos.x += 1;
+
+        let pixelref = unsafe { &mut *(self.pixel.as_mut().unwrap() as *mut SsdPixel) };
+
+        Some((old_pos, pixelref))
+    }
+}
+
+impl SsdOutput {
+    pub fn get_pixel(&mut self, coords: Coordinates<Matrix2DSpace>) -> Option<SsdPixel> {
+        if coords.x < 0 || coords.y < 0 || coords.x >= 128 || coords.y >= 64 {
+            return None;
+        }
+        let idx = (coords.y / 8 * 128 + coords.x) as usize;
+        let bit = (coords.y % 8) as u8;
+        //info!("sample {coords:?} {idx}");
+        let pixref = unsafe { &mut self.0[idx] as *mut u8};
+        Some(SsdPixel(pixref, bit, coords))
+    }
+
+    pub fn new(display: ssd1306::Ssd1306Async<ssd1306::prelude::I2CInterface<esp_hal::i2c::master::I2c<'static, esp_hal::Async>>, ssd1306::prelude::DisplaySize128x64, ssd1306::mode::BasicMode>) -> Self {
+        Self([0; 128 * 64 / 8], display, Default::default())
+    }
+}
+
+impl<'a> Sample<'a, Matrix2DSpace> for SsdOutput {
+    type Output = SsdPixel;
+    
+    fn sample(&mut self, rect: &figments::prelude::Rectangle<Matrix2DSpace>) -> impl Iterator<Item = (Coordinates<Matrix2DSpace>, &'a mut Self::Output)> {
+        let bufref = unsafe {
+            &mut *(self as *mut Self)
+        };
+        let rect = figments::geometry::Rectangle::new(
+            figments::geometry::Coordinates::new(rect.left().clamp(0, 128), rect.top().clamp(0, 64)),
+            figments::geometry::Coordinates::new(rect.right().clamp(0, 128), rect.bottom().clamp(0, 64)),
+        );
+        SsdSampler {
+            buf: bufref,
+            pos: rect.top_left,
+            rect,
+            pixel: None
+        }
+    }
+}
+
+impl DrawTarget for SsdOutput {
+    type Color = BinaryColor;
+
+    type Error = ();
+
+    fn draw_iter<I>(&mut self, pixels: I) -> Result<(), Self::Error>
+    where
+        I: IntoIterator<Item = Pixel<Self::Color>> {
+        for epix in pixels {
+            /*let point = figments::geometry::Rectangle::new(
+                epix.0.into(), epix.0.into()
+            );
+            for (coords, pix) in self.sample(&point) {
+                pix.set_pixel(epix.1);
+            }*/
+            self.get_pixel(epix.0.into()).unwrap().set_pixel(epix.1);
+        }
+        Ok(())
+    }
+
+    fn fill_contiguous<I>(&mut self, area: &embedded_graphics::primitives::Rectangle, colors: I) -> Result<(), Self::Error>
+        where
+            I: IntoIterator<Item = Self::Color>, {
+        let rect = Rectangle::new(
+            area.top_left.into(),
+            area.bottom_right().unwrap().into()
+        );
+        for (color, (coords, pix)) in colors.into_iter().zip(self.sample(&rect)) {
+            pix.set_pixel(color);
+        }
+        Ok(())
+    }
+
+    fn fill_solid(&mut self, area: &embedded_graphics::primitives::Rectangle, color: Self::Color) -> Result<(), Self::Error> {
+        let rect = Rectangle::new(
+            area.top_left.into(),
+            area.bottom_right().unwrap().into()
+        );
+        for (coords, pix) in self.sample(&rect) {
+            pix.set_pixel(color);
+        }
+        Ok(())
+    }
+
+    fn clear(&mut self, color: Self::Color) -> Result<(), Self::Error> {
+        if color.is_on() {
+            self.0.fill(255);
+        } else {
+            self.0.fill(0);
+        }
+        Ok(())
+    }
+}
+
+impl OriginDimensions for SsdOutput {    
+    fn size(&self) -> Size {
+        Size::new(128, 64)
+    }
+}
+
+impl<'a> OutputAsync<'a, Matrix2DSpace> for SsdOutput {
+    type Error = DisplayError;
+
+    type Controls = SsdControls;
+
+    async fn commit_async(&mut self)  -> Result<(), Self::Error> {
+        self.1.set_brightness(ssd1306::prelude::Brightness::custom(1, self.2.data.brightness.load(core::sync::atomic::Ordering::Relaxed))).await.unwrap();
+        self.1.set_display_on(self.2.data.is_on.load(core::sync::atomic::Ordering::Relaxed)).await.unwrap();
+        self.1.draw(&self.0).await
+    }
+
+    fn controls(&self) -> Option<&Self::Controls> {
+        Some(&self.2)
+    }
+}
+
+// TODO: We can probably just replace this entirely with a DisplayControls instance
+#[derive(Debug)]
+struct SsdControlData {
+    is_on: AtomicBool,
+    brightness: AtomicU8
+}
+
+impl Default for SsdControlData {
+    fn default() -> Self {
+        Self {
+            is_on: AtomicBool::new(true),
+            brightness: AtomicU8::new(255)
+        }
+    }
+}
+
+#[derive(Default, Debug, Clone)]
+pub struct SsdControls {
+    data: Arc<SsdControlData>
+}
+
+impl Brightness for SsdControls {
+    fn set_brightness(&mut self, brightness: u8) {
+        self.data.brightness.store(brightness, core::sync::atomic::Ordering::Relaxed);
+    }
+
+    fn set_on(&mut self, is_on: bool) {
+        self.data.is_on.store(is_on, core::sync::atomic::Ordering::Relaxed);
+    }
+}
+
+impl GammaCorrected for SsdControls {
+    fn set_gamma(&mut self, gamma: GammaCurve) {} // Not handled
+}

src/lib.rs

@@ -9,6 +9,7 @@ pub mod ego;
 pub mod animation;
 pub mod idle;
 pub mod logging;
+pub mod graphics;
 
 extern crate alloc;
 

src/tasks/oled.rs

@@ -1,19 +1,23 @@
-use core::{cell::RefCell, f32::consts::PI};
+use core::{cell::{Cell, RefCell}, cmp::min, f32::consts::PI, fmt::Binary, ops::DerefMut};
 
-use alloc::format;
-use embassy_sync::pubsub::DynSubscriber;
-use embassy_time::Timer;
+use alloc::{format, sync::Arc};
+use display_interface::DisplayError;
+use embassy_executor::{raw, Spawner};
+use embassy_sync::{blocking_mutex::raw::{CriticalSectionRawMutex, NoopRawMutex}, mutex::Mutex, pubsub::DynSubscriber};
+use embassy_time::{Delay, Duration, Instant, Timer};
 use embedded_graphics::{image::Image, mono_font::{ascii::{FONT_6X10, FONT_10X20, FONT_9X18_BOLD}, MonoTextStyleBuilder}, pixelcolor::BinaryColor, prelude::*, primitives::{Line, PrimitiveStyleBuilder, Rectangle, StyledDrawable}, text::{Alignment, Baseline, Text}};
 use esp_hal::{i2c::master::I2c, Async, gpio::Output};
-use figments::{liber8tion::{interpolate::{ease_in_out_quad, scale8}, trig::{cos8, sin8}}, prelude::Fract8Ops};
-use nalgebra::Vector2;
+use figments::{liber8tion::{interpolate::{ease_in_out_quad, scale8, Fract8}, trig::{cos8, sin8}}, mappings::{embedded_graphics::{EmbeddedGraphicsSampler, Matrix2DSpace}, linear::LinearSpace}, pixels::{PixelBlend, PixelSink}, prelude::{CoordinateSpace, Coordinates, Fract8Ops}, render::{RenderSource, Sample}, surface::{BufferedSurface, BufferedSurfacePool, Surface, SurfaceBuilder, Surfaces}};
+use figments_render::output::{Brightness, NullControls, OutputAsync};
+use futures::{join, FutureExt};
+use nalgebra::{Matrix, Vector2};
 use ssd1306::{mode::{BufferedGraphicsModeAsync, DisplayConfigAsync}, prelude::{DisplayRotation, I2CInterface}, size::DisplaySize128x64, I2CDisplayInterface, Ssd1306Async};
 use log::*;
 use embedded_graphics::mono_font::MonoTextStyle;
 use embedded_graphics::primitives::PrimitiveStyle;
 use nalgebra::ComplexField;
 
-use crate::{backoff::Backoff, ego::engine::MotionState, events::{Notification, Prediction, Scene, SensorSource, Telemetry}, tasks::ui};
+use crate::{animation::{AnimatedSurface, Animation}, backoff::Backoff, ego::engine::MotionState, events::{Notification, Prediction, Scene, SensorSource, Telemetry}, graphics::ssd1306::{SsdControls, SsdOutput, SsdPixel}, tasks::ui};
 
 mod images {
     use embedded_graphics::{
@@ -23,7 +27,7 @@ mod images {
     include!("../../target/images.rs");
 }
 
-#[derive(Default)]
+#[derive(Default, Debug)]
 struct OledUI {
     scene: Scene,
     motion: MotionState,
@@ -36,7 +40,13 @@ struct OledUI {
     sleep: bool,
 }
 
-struct UiPainter<'a>(&'a mut Ssd1306Async<I2CInterface<I2c<'static, Async>>, DisplaySize128x64, BufferedGraphicsModeAsync<DisplaySize128x64>>);
+#[derive(Default, Debug)]
+struct OledUniforms {
+    ui: OledUI,
+    frame: usize,
+    current_screen: Screen,
+    next_screen: Screen
+}
 
 #[derive(Debug, Default, Clone, Copy)]
 enum Screen {
@@ -48,132 +58,108 @@ enum Screen {
     Waking
 }
 
-impl UiPainter<'_> {
-    pub async fn screen_transition(&mut self, from: Screen, to: Screen, state: &OledUI) {
-        warn!("Transitioning screens from {from:?} to {to:?}");
-        for pos in 0..16 {
-            self.blank();
-            // Draw the screen, then paint the swipe across it from left to right
-            self.draw_screen(from, 0, state);
-            Rectangle::new(Point::zero(), Size { width: figments::liber8tion::interpolate::ease_in_out_quad(pos*8) as u32, height: 64}).draw_styled(&DRAW_STYLE, self.0).unwrap();
-            self.commit().await;
-        }
-
-        for pos in 0..16 {
-            self.blank();
-            // Paint the next screen, then cover the right half with the swipe
-            self.draw_screen(to, 0, state);
-            let pos = figments::liber8tion::interpolate::ease_in_out_quad(pos*8) as u32;
-            Rectangle::new(Point::new(pos as i32, 0), Size { width: 128 - pos, height: 64}).draw_styled(&DRAW_STYLE, self.0).unwrap();
-            self.commit().await;
-        }
-
-        // Finally draw the untouched screen
-        self.blank();
-        self.draw_screen(to, 0, state);
-        self.commit().await;
-    }
-
-    pub fn draw_screen(&mut self, screen: Screen, frame: usize, ui_state: &OledUI) {
-        match screen {
+impl Screen {
+    pub fn draw_screen<'a, T>(&self, sampler: &mut EmbeddedGraphicsSampler<'a, T>, state: &OledUniforms) where T: Sample<'a, Matrix2DSpace>, T::Output: PixelSink<BinaryColor> {
+        match self {
             Screen::Blank => (),
             Screen::Bootsplash => {
-                Image::new(&images::BOOT_LOGO, Point::zero()).draw(self.0).unwrap();
+                Image::new(&images::BOOT_LOGO, Point::zero()).draw(sampler).unwrap();
                 const SPARKLE_COUNT: i32 = 8;
                 for n in 0..SPARKLE_COUNT {
-                    let sparkle_center = Point::new(128u8.scale8(cos8(frame.wrapping_mul(n as usize))) as i32, 64u8.scale8(sin8(frame.wrapping_mul(n as usize) as u8)) as i32);
-                    let offset = (frame / 2 % 32) as i32 - 16;
-                    let rotation = PI * 2.0 * (sin8(frame) as f32 / 255.0);
+                    let sparkle_center = Point::new(128u8.scale8(cos8(state.frame.wrapping_mul(n as usize))) as i32, 64u8.scale8(sin8(state.frame.wrapping_mul(n as usize) as u8)) as i32);
+                    let offset = (state.frame / 2 % 32) as i32 - 16;
+                    let rotation = PI * 2.0 * (sin8(state.frame) as f32 / 255.0);
                     let normal = Point::new((rotation.cos() * offset as f32) as i32, (rotation.sin() * offset as f32) as i32);
                     let cross_normal = Point::new((rotation.sin() * offset as f32) as i32, (rotation.cos() * offset as f32) as i32);
                     // Draw horizontal
                     Line::new(sparkle_center - normal, sparkle_center + normal)
-                        .draw_styled(&SPARKLE_STYLE, self.0).unwrap();
+                        .draw_styled(&SPARKLE_STYLE, sampler).unwrap();
                     // Draw vertical
                     Line::new(sparkle_center - cross_normal, sparkle_center + cross_normal)
-                        .draw_styled(&SPARKLE_STYLE, self.0).unwrap();
+                        .draw_styled(&SPARKLE_STYLE, sampler).unwrap();
                 }
             },
             Screen::Home => {
                 // Status bar
                 // Sensor indicators
-                let gps_img = if ui_state.gps_online {
+                let gps_img = if state.ui.gps_online {
                     &images::GPS_ON
                 } else {
                     &images::GPS_OFF
                 };
-                let imu_img = if ui_state.imu_online {
+                let imu_img = if state.ui.imu_online {
                     &images::IMU_ON
                 } else {
                     &images::IMU_OFF
                 };
-                Image::new(gps_img, Point::zero()).draw(self.0).unwrap();
-                Image::new(imu_img, Point::new((gps_img.size().width + 2) as i32, 0)).draw(self.0).unwrap();
+                Image::new(gps_img, Point::zero()).draw(sampler).unwrap();
+                Image::new(imu_img, Point::new((gps_img.size().width + 2) as i32, 0)).draw(sampler).unwrap();
 
                 #[cfg(feature="demo")]
                 Text::with_alignment("Demo", Point::new(128, 10), TEXT_STYLE, Alignment::Right)
-                    .draw(self.0)
+                    .draw(sampler)
                     .unwrap();
 
                 #[cfg(feature="simulation")]
                 Text::with_alignment("Sim", Point::new(128, 10), TEXT_STYLE, Alignment::Right)
-                    .draw(self.0)
+                    .draw(sampler)
                     .unwrap();
 
                 // Separates the status bar from the UI
-                Line::new(Point::new(0, 18), Point::new(128, 18)).draw_styled(&INACTIVE_STYLE, self.0).unwrap();
+                Line::new(Point::new(0, 18), Point::new(128, 18)).draw_styled(&INACTIVE_STYLE, sampler).unwrap();
 
                 // Speed display at the top
-                Text::with_alignment(&format!("{}", ui_state.velocity), Point::new(128 / 2, 12), SPEED_STYLE, Alignment::Center)
-                    .draw(self.0)
+                Text::with_alignment(&format!("{}", state.ui.velocity), Point::new(128 / 2, 12), SPEED_STYLE, Alignment::Center)
+                    .draw(sampler)
                     .unwrap();
 
                 // The main UI content
-                Image::new(&images::BIKE, Point::new((128 / 2 - images::BIKE.size().width / 2) as i32, 24)).draw(self.0).unwrap();
+                Image::new(&images::BIKE, Point::new((128 / 2 - images::BIKE.size().width / 2) as i32, 24)).draw(sampler).unwrap();
 
-                let headlight_img = if ui_state.headlight {
+                let headlight_img = if state.ui.headlight {
                     &images::HEADLIGHT_ON
                 } else {
                     &images::HEADLIGHT_OFF
                 };
-                let brakelight_img = if ui_state.brakelight {
+                let brakelight_img = if state.ui.brakelight {
                     &images::BRAKELIGHT_ON
                 } else {
                     &images::BRAKELIGHT_OFF
                 };
-                Image::new(headlight_img, Point::new(((128 / 2 - images::BIKE.size().width / 2) - 18) as i32, 28)).draw(self.0).unwrap();
-                Image::new(brakelight_img, Point::new(((128 / 2 + images::BIKE.size().width / 2) + 2) as i32, 28)).draw(self.0).unwrap();
+                Image::new(headlight_img, Point::new(((128 / 2 - images::BIKE.size().width / 2) - 18) as i32, 28)).draw(sampler).unwrap();
+                Image::new(brakelight_img, Point::new(((128 / 2 + images::BIKE.size().width / 2) + 2) as i32, 28)).draw(sampler).unwrap();
 
                 // TODO: Replace the state texts with cute animations or smth
                 // Current prediction from the motion engine
-                Text::with_alignment(&format!("{:?}", ui_state.motion), Point::new(128 / 2, 64 - 3), TEXT_STYLE, Alignment::Center)
-                    .draw(self.0)
+                Text::with_alignment(&format!("{:?}", state.ui.motion), Point::new(128 / 2, 64 - 3), TEXT_STYLE, Alignment::Center)
+                    .draw(sampler)
                     .unwrap();
                 
                 // Current scene in the UI
-                Text::with_alignment(&format!("{:?}", ui_state.scene), Point::new(128 / 2, 64 - 13), TEXT_STYLE, Alignment::Center)
-                    .draw(self.0)
+                Text::with_alignment(&format!("{:?}", state.ui.scene), Point::new(128 / 2, 64 - 13), TEXT_STYLE, Alignment::Center)
+                    .draw(sampler)
                     .unwrap();
             },
             Screen::Sleeping => {
                 Text::with_alignment("so eepy Zzz...", Point::new(128 / 2, 64 / 2), LOGO_STYLE, Alignment::Center)
-                    .draw(self.0)
+                    .draw(sampler)
                     .unwrap();
             },
             Screen::Waking => {
                 Text::with_alignment("OwO hewwo again :D", Point::new(128 / 2, 64 / 2), LOGO_STYLE, Alignment::Center)
-                    .draw(self.0)
+                    .draw(sampler)
                     .unwrap();
             }
         }
     }
+}
 
-    pub fn blank(&mut self) {
-        self.0.clear_buffer();
-    }
-
-    pub async fn commit(&mut self) {
-        self.0.flush().await.expect("Could not commit frame");
+impl RenderSource<OledUniforms, Matrix2DSpace, BinaryColor, SsdPixel> for Screen {
+    fn render_to<'a, Smp>(&'a self, output: &'a mut Smp, uniforms: &OledUniforms)
+        where 
+            Smp: Sample<'a, Matrix2DSpace, Output = SsdPixel> {
+        let mut sampler = EmbeddedGraphicsSampler(output, embedded_graphics::primitives::Rectangle::new(Point::new(0, 0), Size::new(128, 64)));
+        self.draw_screen(&mut sampler, uniforms);
     }
 }
 
@@ -208,53 +194,75 @@ const INACTIVE_STYLE: PrimitiveStyle<BinaryColor> = PrimitiveStyleBuilder::new()
     .build();
 
 #[embassy_executor::task]
-pub async fn oled_task(i2c: I2c<'static, Async>,  mut reset_pin: Output<'static>, mut events: DynSubscriber<'static, Telemetry>) {
-    let interface = RefCell::new(Some(I2CDisplayInterface::new(i2c)));
-    let mut ui_state = OledUI::default();
+async fn oled_render(mut output: SsdOutput, surfaces: BufferedSurfacePool<OledUniforms, Matrix2DSpace, BinaryColor>, uniforms: Arc<Mutex<NoopRawMutex, OledUniforms>>) {
+    Backoff::from_secs(1).forever().attempt::<_, (), DisplayError>(async || {
+        const FPS: u64 = 30;
+        const RENDER_BUDGET: Duration = Duration::from_millis(1000 / FPS);
 
-    // initialize the display
-    let mut display = Backoff::from_secs(3).forever().attempt(async || {
-        info!("Setting up OLED display");
-        Timer::after_millis(10).await;
-        reset_pin.set_high();
-        Timer::after_millis(10).await;
-        reset_pin.set_low();
-        Timer::after_millis(10).await;
-        reset_pin.set_high();
-        let mut display = Ssd1306Async::new(interface.replace(None).unwrap(), DisplaySize128x64, DisplayRotation::Rotate0)
-            .into_buffered_graphics_mode();
-        if let Err(e) = display.init().await {
-            warn!("Failed to set up OLED display: {e:?}");
-            interface.replace(Some(display.release()));
-            Err(())
-        } else {
-            Ok(display)
+        const ANIMATION_TPS: u64 = 30;
+        const ANIMATION_FRAME_TIME: Duration = Duration::from_millis(1000 / ANIMATION_TPS);
+        info!("Starting Oled renderer");
+        loop {
+            let start = Instant::now();
+            {
+                let mut locked = uniforms.lock().await;
+                output.clear(BinaryColor::Off).unwrap();
+                locked.frame = (Instant::now().as_millis() / ANIMATION_FRAME_TIME.as_millis()) as usize;
+                locked.current_screen.render_to(&mut output, &locked);
+                surfaces.render_to(&mut output, &locked);
+            }
+            output.commit_async().await?;
+            let frame_time = Instant::now() - start;
+            if frame_time < RENDER_BUDGET {
+                Timer::after(RENDER_BUDGET - frame_time).await;
+            } else {
+                //warn!("OLED Frame took too long to render! {}ms", frame_time.as_millis());
+            }
         }
     }).await.unwrap();
+}
 
-    display.clear(BinaryColor::Off).unwrap();
-    display.flush().await.unwrap();
-
-    info!("Running boot splash animation");
-    let mut painter = UiPainter(&mut display);
-    painter.screen_transition(Screen::Blank, Screen::Bootsplash, &ui_state).await;
-    for i in 0..255 {
-        painter.blank();
-        painter.draw_screen(Screen::Bootsplash, i, &ui_state);
-        painter.commit().await;
+async fn screen_transition(overlay: &mut BufferedSurface<OledUniforms, Matrix2DSpace, BinaryColor>, uniforms: &mut Arc<Mutex<NoopRawMutex, OledUniforms>>, next_screen: Screen) {
+    const FADE_IN: Animation = Animation::new().from(0).to(255).duration(Duration::from_millis(300));
+    const FADE_OUT: Animation = Animation::new().from(255).to(0).duration(Duration::from_millis(300));
+    info!("Fading in to screen {next_screen:?}");
+    FADE_IN.apply(overlay).await;
+    {
+        let mut locked = uniforms.lock().await;
+        locked.current_screen = next_screen;
     }
-    painter.screen_transition(Screen::Bootsplash, Screen::Home, &ui_state).await;
+    FADE_OUT.apply(overlay).await;
+}
+
+#[embassy_executor::task]
+async fn oled_ui(mut events: DynSubscriber<'static, Telemetry>, mut overlay: BufferedSurface<OledUniforms, Matrix2DSpace, BinaryColor>, mut uniforms: Arc<Mutex<NoopRawMutex, OledUniforms>>, mut controls: SsdControls) {
+    
+    screen_transition(&mut overlay, &mut uniforms, Screen::Bootsplash).await;
+    Timer::after_secs(3).await;
+    screen_transition(&mut overlay, &mut uniforms, Screen::Home).await;
 
-    info!("OLED display is ready!");
     loop {
-        // FIXME: Need to implement sleep handling to electrically turn the display on/off
-        if !ui_state.sleep {
-            painter.blank();
-            painter.draw_screen(Screen::Home, 0, &ui_state);
-            painter.commit().await;
-        }
-
         let evt = events.next_message_pure().await;
+        //info!("Oled UI event: {evt:?}");
+        match evt {
+            Telemetry::Notification(Notification::Sleep) => {
+                screen_transition(&mut overlay, &mut uniforms, Screen::Sleeping).await;
+                Timer::after_secs(1).await;
+                screen_transition(&mut overlay, &mut uniforms, Screen::Blank).await;
+                controls.set_on(false);
+                //ui_state.sleep = true
+            },
+            Telemetry::Notification(Notification::WakeUp) => {
+                controls.set_on(true);
+                screen_transition(&mut overlay, &mut uniforms, Screen::Waking).await;
+                Timer::after_secs(1).await;
+                screen_transition(&mut overlay, &mut uniforms, Screen::Home).await;
+                //ui_state.sleep = false
+            },
+            _ => ()
+        }
+        let mut locked = uniforms.lock().await;
+        let ui_state = &mut locked.ui;
         match evt {
             Telemetry::Prediction(Prediction::Velocity(v)) => ui_state.velocity = v,
             Telemetry::Prediction(Prediction::Location(loc)) => ui_state.location = loc,
@@ -266,19 +274,39 @@ pub async fn oled_task(i2c: I2c<'static, Async>,  mut reset_pin: Output<'static>
             Telemetry::Notification(Notification::SensorOnline(SensorSource::IMU)) => ui_state.imu_online = true,
             Telemetry::Notification(Notification::SensorOffline(SensorSource::GPS)) => ui_state.gps_online = false,
             Telemetry::Notification(Notification::SensorOnline(SensorSource::GPS)) => ui_state.gps_online = true,
-            Telemetry::Notification(Notification::Sleep) => {
-                painter.screen_transition(Screen::Home, Screen::Sleeping, &ui_state).await;
-                Timer::after_secs(1).await;
-                painter.screen_transition(Screen::Sleeping, Screen::Blank, &ui_state).await;
-                ui_state.sleep = true
-            },
-            Telemetry::Notification(Notification::WakeUp) => {
-                painter.screen_transition(Screen::Blank, Screen::Waking, &ui_state).await;
-                Timer::after_secs(1).await;
-                painter.screen_transition(Screen::Waking, Screen::Home, &ui_state).await;
-                ui_state.sleep = false
-            },
             _ => ()
         }
     }
+    /*for (coords, pix) in output.sample(&figments::prelude::Rectangle::everything()) {
+        pix.set(&BinaryColor::Off);
+    }
+    Text::with_alignment("Foo", Point::new(128 / 2, 64 / 2), LOGO_STYLE, Alignment::Center).draw(&mut output).unwrap();
+    output.commit_async().await.unwrap();
+    Timer::after_secs(1).await;
+
+    let painter = UiPainter { current_screen: Screen::Bootsplash, next_screen: Screen::Bootsplash };
+    painter.render_to(&mut output, &OledUniforms { frame: 0, ui: ui_state });*/
+
+}
+
+#[embassy_executor::task]
+pub async fn oled_task(i2c: I2c<'static, Async>,  mut reset_pin: Output<'static>, events: DynSubscriber<'static, Telemetry>) {
+    let interface = I2CDisplayInterface::new(i2c);
+    let mut display = Ssd1306Async::new(interface, DisplaySize128x64, DisplayRotation::Rotate0);
+    display.reset(&mut reset_pin, &mut Delay).await.unwrap();
+    display.init_with_addr_mode(ssd1306::command::AddrMode::Horizontal).await.unwrap();
+
+    let output = SsdOutput::new(display);
+    let mut surfaces = BufferedSurfacePool::default();
+    let uniforms = Default::default();
+
+    let sfc = SurfaceBuilder::build(&mut surfaces).shader(|coords: &Coordinates<Matrix2DSpace>, uniforms: &OledUniforms| {
+        BinaryColor::On
+    }).opacity(0).finish().unwrap();
+
+    info!("Starting OLED display tasks!");
+    let spawner = Spawner::for_current_executor().await;
+    let ui_controls = output.controls().unwrap().clone();
+    spawner.must_spawn(oled_render(output, surfaces, Arc::clone(&uniforms)));
+    spawner.must_spawn(oled_ui(events, sfc, uniforms, ui_controls));
 }