acid/firmware2/src/main.rs

#![no_std]
#![no_main]
#![feature(macro_metavar_expr)]

extern crate alloc;

use core::alloc::Layout;
use core::cell::{OnceCell, RefCell};
use core::fmt::Write;
use core::time::Duration;

use alloc::boxed::Box;
use alloc::rc::Rc;
use alloc::vec;
use shadow_rs::shadow;
use bt_hci::controller::ExternalController;
use critical_section::Mutex;
use embassy_embedded_hal::adapter::BlockingAsync;
use embassy_executor::Spawner;
use embassy_time::Timer;
use esp_alloc::{HeapRegion, MemoryCapability};
use esp_hal::clock::CpuClock;
use esp_hal::dma::{BurstConfig, DmaDescriptor, DmaTxBuf, ExternalBurstConfig};
use esp_hal::gpio::{Flex, Input, InputConfig, Io, Level, Output, OutputConfig, Pull};
use esp_hal::handler;
use esp_hal::i2c::master::{I2c, I2cAddress};
use esp_hal::interrupt::Priority;
use esp_hal::interrupt::software::SoftwareInterruptControl;
use esp_hal::lcd_cam::LcdCam;
use esp_hal::lcd_cam::lcd::dpi::Dpi;
use esp_hal::mcpwm::{McPwm, PeripheralClockConfig};
use esp_hal::peripherals::Peripherals;
use esp_hal::psram::{FlashFreq, PsramConfig, PsramSize, SpiRamFreq, SpiTimingConfigCoreClock};
use esp_hal::rng::TrngSource;
use esp_hal::system::Stack;
use esp_hal::time::Instant;
use esp_hal::timer::timg::TimerGroup;
use esp_hal::uart::Uart;
use esp_hal::usb_serial_jtag::{UsbSerialJtag, UsbSerialJtagTx};
use esp_hal::{Blocking, ram};
use esp_radio::Controller;
use esp_radio::ble::controller::BleConnector;
use esp_rtos::embassy::{Executor, InterruptExecutor};
use esp_storage::FlashStorage;
use itertools::chain;
use log::{LevelFilter, Log, debug, error, info};
use rmk::channel::EVENT_CHANNEL;
use rmk::config::{BehaviorConfig, PositionalConfig, RmkConfig, StorageConfig, VialConfig};
use rmk::debounce::default_debouncer::DefaultDebouncer;
use rmk::embassy_futures::yield_now;
use rmk::input_device::Runnable;
use rmk::{join_all};
use rmk::keyboard::Keyboard;
use rmk::storage::async_flash_wrapper;
use rmk::{initialize_keymap_and_storage, run_devices, run_rmk};
use slint::platform::software_renderer::{Rgb565Pixel, TargetPixel};
use slint::{ComponentHandle, PhysicalSize, WindowSize};
use slint::platform::Platform;
use static_cell::StaticCell;
use ui::AppWindow;
use {esp_alloc as _, esp_backtrace as _};

use crate::matrix::IoeMatrix;
use crate::peripherals::st7701s::St7701s;
use crate::vial::{VIAL_KEYBOARD_DEF, VIAL_KEYBOARD_ID};

mod keymap;
mod matrix;
mod peripherals;
mod vial;
mod ui;
mod logging;
mod console;

shadow!(build);

// This creates a default app-descriptor required by the esp-idf bootloader.
// For more information see: <https://docs.espressif.com/projects/esp-idf/en/stable/esp32/api-reference/system/app_image_format.html#application-description>
esp_bootloader_esp_idf::esp_app_desc!();

const LOG_LEVEL_FILTER: LevelFilter = LevelFilter::Info;

static PSRAM_ALLOCATOR: esp_alloc::EspHeap = esp_alloc::EspHeap::empty();

#[esp_rtos::main]
async fn main(_spawner: Spawner) {
    #[cfg(not(feature = "alt-log"))]
    let alt_uart_rx_task = {
        esp_println::logger::init_logger(LOG_LEVEL_FILTER);
        info!("Logger initialized!");
        async {}
    };

    let config = esp_hal::Config::default()
        .with_cpu_clock(CpuClock::max())
        .with_psram(PsramConfig {
            size: PsramSize::AutoDetect,
            core_clock: Some(SpiTimingConfigCoreClock::SpiTimingConfigCoreClock80m),
            flash_frequency: FlashFreq::default(),
            ram_frequency: SpiRamFreq::Freq80m,
        });
    let peripherals: esp_hal::peripherals::Peripherals = esp_hal::init(config);
    info!("System initialized!");

    #[cfg(feature = "alt-log")]
    let alt_uart_rx_task = {
        let (uart_rx, uart_tx) = Uart::new(peripherals.UART2, Default::default()).unwrap().with_tx(peripherals.GPIO12).with_rx(peripherals.GPIO5).split();
        logging::setup_alternative_logging(uart_tx, LOG_LEVEL_FILTER);
        info!("Logger initialized!");
        console::run_console(uart_rx.into_async())
    };

    // Use the internal DRAM as the heap.
    esp_alloc::heap_allocator!(#[unsafe(link_section = ".dram2_uninit")] size: 64 * 1024);
    info!("Heap initialized! {:#?}", esp_alloc::HEAP.stats());

    // Initialize the PSRAM allocator.
    {
        let (psram_offset, psram_size) = esp_hal::psram::psram_raw_parts(&peripherals.PSRAM);
        unsafe {
            PSRAM_ALLOCATOR.add_region(HeapRegion::new(
                psram_offset,
                psram_size,
                MemoryCapability::External.into(),
            ));
        }
    }

    let mut io = Io::new(peripherals.IO_MUX);
    io.set_interrupt_handler(interrupt_handler);

    // Enable pull-up on GPIO0 to prevent booting into download mode.
    let gpio0 = Output::new(
        peripherals.GPIO0,
        Level::High,
        OutputConfig::default().with_pull(Pull::Up),
    );

    // Enable LDO2
    let _ = Output::new(peripherals.GPIO17, Level::High, OutputConfig::default());

    // Enable antenna
    let _ = Output::new(peripherals.GPIO11, Level::Low, OutputConfig::default());

    // TODO: Use PWM to control the pwm_pin.
    let mut _pwm = McPwm::new(peripherals.MCPWM0, PeripheralClockConfig::with_prescaler(1));
    let mut _pwm_pin = Output::new(peripherals.GPIO21, Level::High, OutputConfig::default());

    let timg0 = TimerGroup::new(peripherals.TIMG0);
    let software_interrupt = SoftwareInterruptControl::new(peripherals.SW_INTERRUPT);
    esp_rtos::start(
        timg0.timer0, /*, software_interrupt.software_interrupt0 */
    );
    // Enable the TRNG source, so `Trng` can be constructed.
    let _trng_source = TrngSource::new(peripherals.RNG, peripherals.ADC1);
    let mut rng = esp_hal::rng::Trng::try_new().unwrap();
    static RADIO: StaticCell<Controller<'static>> = StaticCell::new();
    let radio = RADIO.init(esp_radio::init().unwrap());
    let bluetooth = peripherals.BT;
    let connector = BleConnector::new(radio, bluetooth, Default::default()).unwrap();
    let controller: ExternalController<_, 20> = ExternalController::new(connector);
    let central_addr = [0x18, 0xe2, 0x21, 0x80, 0xc0, 0xc7];
    #[cfg(feature = "ble")]
    let mut host_resources = rmk::HostResources::new();
    #[cfg(feature = "ble")]
    let stack =
        rmk::ble::build_ble_stack(controller, central_addr, &mut rng, &mut host_resources).await;

    // Initialize USB
    #[cfg(not(feature = "no-usb"))]
    let usb_driver = {
        use core::ptr::addr_of_mut;
        use esp_hal::otg_fs::Usb;
        use esp_hal::otg_fs::asynch::{Config, Driver};

        static mut EP_MEMORY: [u8; 1024] = [0; 1024];
        let usb = Usb::new(peripherals.USB0, peripherals.GPIO20, peripherals.GPIO19);
        // Create the driver, from the HAL.
        let config = Config::default();
        Driver::new(usb, unsafe { &mut *addr_of_mut!(EP_MEMORY) }, config)
    };

    // Initialize the flash
    let flash = FlashStorage::new(peripherals.FLASH);
    let flash = async_flash_wrapper(flash);

    let sck = Output::new(peripherals.GPIO36, Level::High, OutputConfig::default());
    let mosi = Flex::new(peripherals.GPIO35);
    let cs = Output::new(peripherals.GPIO6, Level::High, OutputConfig::default());

    let lcd = LcdCam::new(peripherals.LCD_CAM).lcd;
    let unconfigured_dpi = Dpi::new(lcd, peripherals.DMA_CH2, Default::default())
        .unwrap()
        .with_de(peripherals.GPIO37)
        .with_pclk(peripherals.GPIO34)
        .with_hsync(peripherals.GPIO44)
        .with_vsync(peripherals.GPIO43)
        // Blue
        .with_data0(peripherals.GPIO38)
        .with_data1(peripherals.GPIO39)
        .with_data2(peripherals.GPIO40)
        .with_data3(peripherals.GPIO41)
        .with_data4(peripherals.GPIO42)
        // Green
        // .with_data5(peripherals.GPIO5) TODO
        // .with_data6(peripherals.GPIO12)
        .with_data7(peripherals.GPIO13)
        .with_data8(peripherals.GPIO14)
        .with_data9(peripherals.GPIO15)
        .with_data10(peripherals.GPIO16)
        // Red
        .with_data11(gpio0)
        .with_data12(peripherals.GPIO1)
        .with_data13(peripherals.GPIO2)
        .with_data14(peripherals.GPIO3)
        .with_data15(peripherals.GPIO4);

    let st7701s = St7701s::new(sck, mosi, cs, unconfigured_dpi).await;

    // RMK config
    let vial_config = VialConfig::new(VIAL_KEYBOARD_ID, VIAL_KEYBOARD_DEF, &[(0, 0), (1, 1)]);
    let storage_config = StorageConfig {
        start_addr: 0x3f0000,
        num_sectors: 16,
        ..Default::default()
    };
    let rmk_config = RmkConfig {
        vial_config,
        storage_config,
        ..Default::default()
    };

    // Initialze keyboard stuffs
    // Initialize the storage and keymap
    let mut default_keymap = keymap::get_default_keymap();
    let mut behavior_config = BehaviorConfig::default();
    let mut per_key_config = PositionalConfig::default();
    let (keymap, mut storage) = initialize_keymap_and_storage(
        &mut default_keymap,
        flash,
        &storage_config,
        &mut behavior_config,
        &mut per_key_config,
    )
    .await;

    // Initialize the matrix and keyboard
    const I2C_ADDR_MATRIX_LEFT: I2cAddress = I2cAddress::SevenBit(0b0100000);
    const I2C_ADDR_MATRIX_RIGHT: I2cAddress = I2cAddress::SevenBit(0b0100001);

    let i2c = I2c::new(peripherals.I2C0, Default::default())
        .unwrap()
        .with_sda(peripherals.GPIO8)
        .with_scl(peripherals.GPIO9);

    let matrix_interrupt_low = Input::new(peripherals.GPIO7, InputConfig::default());

    let mut matrix = IoeMatrix::new(
        matrix_interrupt_low,
        i2c.into_async(),
        DefaultDebouncer::new(),
        [I2C_ADDR_MATRIX_LEFT, I2C_ADDR_MATRIX_RIGHT],
    )
    .await;
    let mut keyboard = Keyboard::new(&keymap); // Initialize the light controller

    static FRAMEBUFFER: StaticCell<Framebuffer> = StaticCell::new();
    let framebuffer = FRAMEBUFFER.init(Framebuffer::new(
        360 + /* TODO: Figure out why more bytes are needed: */ 8,
        960,
    ));

    let window_size = [framebuffer.width, framebuffer.height];
    let framebuffer_ptr = FramebufferPtr(framebuffer.as_target_pixels() as _);

    static SECOND_CORE_STACK: StaticCell<Stack<8192>> = StaticCell::new();
    let second_core_stack = SECOND_CORE_STACK.init(Stack::new());
    esp_rtos::start_second_core(
        peripherals.CPU_CTRL,
        software_interrupt.software_interrupt0,
        software_interrupt.software_interrupt1,
        second_core_stack,
        move || {
            // static EXECUTOR: StaticCell<InterruptExecutor<2>> = StaticCell::new();
            // let exec = EXECUTOR.init(InterruptExecutor::new(
            //     software_interrupt.software_interrupt2,
            // ));
            // let spawner = exec.start(Priority::Priority3);
            // spawner.must_spawn(run_renderer_task());
            static EXECUTOR: StaticCell<Executor> = StaticCell::new();
            let executor: &mut Executor = EXECUTOR.init(Executor::new());
            executor.run(|spawner| {
                let slint_backend = SlintBackend {
                    // peripherals: RefCell::new(Some(peripherals)),
                    window_size,
                    window: RefCell::new(None),
                    framebuffer: framebuffer_ptr,
                };
                spawner.must_spawn(run_renderer_task(slint_backend));
            });
        },
    );

    join_all![
        // We currently send the framebuffer data using the main core, which does not seem to slow
        // down the rest of the tasks too much.
        run_lcd(st7701s, framebuffer),
        run_devices! (
            (matrix) => EVENT_CHANNEL,
        ),
        keyboard.run(), // Keyboard is special
        run_rmk(
            &keymap,
            #[cfg(not(feature = "no-usb"))]
            usb_driver,
            #[cfg(feature = "ble")]
            &stack,
            &mut storage,
            rmk_config,
        ),
        alt_uart_rx_task
    ]
    .await;
}

#[embassy_executor::task]
async fn run_renderer_task(backend: SlintBackend) {
    slint::platform::set_platform(Box::new(backend))
    .expect("backend already initialized");

    loop {
        let main = AppWindow::new().unwrap();

        main.run().unwrap();
    }
}

struct Framebuffer {
    width: u32,
    height: u32,
    dma_buf: Option<DmaTxBuf>,
}

impl Framebuffer {
    pub fn new(width: u32, height: u32) -> Self {
        let buffer_len = width as usize * height as usize * core::mem::size_of::<u16>();
        // Allocate the framebuffer in the external PSRAM memory.
        // Note: We just leak this buffer.
        let buffer_ptr = unsafe {
            // ⚠️ Note: For chips that support DMA to/from PSRAM (ESP32-S3) DMA transfers to/from PSRAM
            // have extra alignment requirements. The address and size of the buffer pointed to by each
            // descriptor must be a multiple of the cache line (block) size. This is 32 bytes on ESP32-S3.
            PSRAM_ALLOCATOR.alloc_caps(
                MemoryCapability::External.into(),
                Layout::from_size_align(buffer_len, 32).unwrap(),
            )
        };
        let buffer = unsafe { core::slice::from_raw_parts_mut(buffer_ptr, buffer_len) };
        let burst_config: BurstConfig = ExternalBurstConfig::Size16.into();

        info!(
            "PSRAM SPI burst config: max_compatible_chunk_size={}",
            burst_config.max_compatible_chunk_size()
        );
        let dma_buf_descs_len = esp_hal::dma::descriptor_count(
            buffer_len,
            burst_config.max_compatible_chunk_size(),
            false,
        );
        // Descriptors are initialized by `DmaTxBuf::new`.
        let dma_buf_descs =
            Box::leak(vec![DmaDescriptor::EMPTY; dma_buf_descs_len].into_boxed_slice());
        let dma_buf = DmaTxBuf::new(dma_buf_descs, buffer).unwrap();

        Self {
            width,
            height,
            dma_buf: Some(dma_buf),
        }
    }

    pub fn as_target_pixels(&mut self) -> &mut [Rgb565Pixel] {
        bytemuck::cast_slice_mut::<_, Rgb565Pixel>(self.dma_buf.as_mut().unwrap().as_mut_slice())
    }
}

struct SlintBackend {
    window_size: [u32; 2],
    window: RefCell<Option<Rc<slint::platform::software_renderer::MinimalSoftwareWindow>>>,
    framebuffer: FramebufferPtr,
    // peripherals: RefCell<Option<Peripherals>>,
}

struct FramebufferPtr(*mut [Rgb565Pixel]);

unsafe impl Send for FramebufferPtr {}

impl slint::platform::Platform for SlintBackend {
    fn create_window_adapter(&self) -> Result<Rc<dyn slint::platform::WindowAdapter>, slint::PlatformError> {
        let window = slint::platform::software_renderer::MinimalSoftwareWindow::new(
            slint::platform::software_renderer::RepaintBufferType::ReusedBuffer,
        );
        window.set_size(WindowSize::Physical(PhysicalSize::new(self.window_size[0], self.window_size[1])));
        self.window.replace(Some(window.clone()));
        Ok(window)
    }

    fn duration_since_start(&self) -> Duration {
        Duration::from_millis(Instant::now().duration_since_epoch().as_millis())
    }

    fn run_event_loop(&self) -> Result<(), slint::PlatformError> {
        loop {
            slint::platform::update_timers_and_animations();
            
            if let Some(window) = self.window.borrow().clone() {
                // window.try_dispatch_event(todo!())?;

                window.draw_if_needed(|renderer| {
                    // TODO: Proper synchronization.
                    let framebuffer = unsafe { &mut *self.framebuffer.0 };
                    // TODO: Try using height to see if it rotates the screen correctly. Might need
                    // to swap dimensions elsewhere.
                    renderer.render(framebuffer, self.window_size[0] as usize);
                });
            }
        }
    }
}

// impl DrawTarget for Framebuffer {
//     type Color = Rgb565;
//     type Error = ();

//     fn draw_iter<I>(&mut self, pixels: I) -> Result<(), Self::Error>
//     where
//         I: IntoIterator<Item = Pixel<Self::Color>>,
//     {
//         let buffer = bytemuck::try_cast_slice_mut::<_, u16>(self.dma_buf.as_mut_slice()).unwrap();

//         for Pixel(coord, color) in pixels.into_iter() {
//             // Check if the pixel coordinates are out of bounds.
//             // `DrawTarget` implementation are required to discard any out of bounds
//             // pixels without returning an error or causing a panic.
//             if coord.x >= 0
//                 && coord.x < self.width as i32
//                 && coord.y >= 0
//                 && coord.y < self.height as i32
//             {
//                 let index = coord.x as usize + coord.y as usize * self.width as usize;
//                 buffer[index] = RawU16::from(color).into_inner();
//             }
//         }

//         Ok(())
//     }
// }

#[embassy_executor::task]
async fn run_lcd_task(st7701s: St7701s<'static, Blocking>, framebuffer: &'static mut Framebuffer) {
    run_lcd(st7701s, framebuffer).await
}

async fn run_lcd(mut st7701s: St7701s<'static, Blocking>, framebuffer: &'static mut Framebuffer) {
    // const PADDING_LEFT: usize = 121;
    // const PADDING_RIGHT: usize = 7;
    // const COLORS_WIDTH: usize = 240;
    // const COLORS_HEIGHT: usize = 960;

    // const MAX_RED: u8 = (1 << 5) - 1;
    // const MAX_GREEN: u8 = (1 << 6) - 1;
    // const MAX_BLUE: u8 = (1 << 5) - 1;

    // fn rgb(r: u8, g: u8, b: u8) -> u16 {
    //     (((r & MAX_RED) as u16) << 11) | (((g & MAX_GREEN) as u16) << 5) | ((b & MAX_BLUE) as u16)
    // }

    // fn row(edge: u16) -> impl Iterator<Item = u16> + Clone {
    //     chain![
    //         core::iter::repeat_n(rgb(0, 0, 0xFF), PADDING_LEFT),
    //         core::iter::once(rgb(0xFF, 0xFF, 0xFF)),
    //         core::iter::repeat_n(edge, COLORS_WIDTH - 2),
    //         core::iter::once(rgb(0xFF, 0xFF, 0xFF)),
    //         core::iter::repeat_n(rgb(0, 0, 0xFF), PADDING_RIGHT),
    //     ]
    // }

    // let mut colors = chain![
    //     row(rgb(0xFF, 0, 0)),
    //     core::iter::repeat_n(row(0), COLORS_HEIGHT - 2).flatten(),
    //     row(rgb(0xFF, 0xFF, 0)),
    // ];

    // const BUFFER_LEN: usize = core::mem::size_of::<u16>()
    //     * (360 + /* TODO: Figure out why more bytes are needed: */ 8)
    //     * 960;

    // {
    //     for (chunk, color) in dma_buf.as_mut_slice().chunks_mut(2).zip(&mut colors) {
    //         chunk.copy_from_slice(&color.to_le_bytes());
    //     }

    //     info!("chunk addr: {}", dma_buf.as_slice().as_ptr() as usize);
    //     // colors.next(); // Shift colors
    // }

    loop {
        // Timer::after(Duration::from_millis(100)).await;
        yield_now().await;

        // TODO: Use bounce buffers:
        // https://docs.espressif.com/projects/esp-idf/en/v5.0/esp32s3/api-reference/peripherals/lcd.html#bounce-buffer-with-single-psram-frame-buffer
        // They need to be implemented in esp-hal.
        let transfer = match st7701s.dpi.send(false, framebuffer.dma_buf.take().unwrap()) {
            Err((error, result_dpi, result_dma_buf)) => {
                error!(
                    "An error occurred while initiating transfer of the framebuffer to the LCD display: {error:?}"
                );
                st7701s.dpi = result_dpi;
                framebuffer.dma_buf = Some(result_dma_buf);
                continue;
            }
            Ok(transfer) => transfer,
        };

        while !transfer.is_done() {
            // Timer::after_millis(1).await;
            yield_now().await;
        }

        let result;
        let dma_buf;
        (result, st7701s.dpi, dma_buf) = transfer.wait();
        framebuffer.dma_buf = Some(dma_buf);

        if let Err(error) = result {
            error!(
                "An error occurred while transferring framebuffer to the LCD display: {error:?}"
            );
        }
    }
}

#[handler]
#[ram] // TODO: Is this necessary?
fn interrupt_handler() {
    esp_println::println!(
        "GPIO Interrupt with priority {}",
        esp_hal::xtensa_lx::interrupt::get_level()
    );
}