firmware for my Touchscreen E-Paper Input Module for Framework Laptop 16
1use core::sync::atomic::Ordering;
2use core::fmt::Write;
3use portable_atomic::{AtomicBool, AtomicU8};
4use usb_device::device::UsbDevice;
5use usbd_serial::SerialPort;
6use eepy_serial::{Response, SerialCommand};
7use eepy_sys::flash::erase_and_program;
8use eepy_sys::header::{slot, slot_ptr, Programs};
9use eepy_sys::image::refresh;
10use eepy_sys::{header, IMAGE_BYTES};
11use eepy_sys::input::{next_event, set_touch_enabled};
12use eepy_sys::misc::{debug, info, trace};
13use eepy_sys::usb::UsbBus;
14use crate::{delete_program, USB_DEVICE, USB_SERIAL};
15use crate::ui::flashing::draw_flashing_ui;
16
17#[derive(Copy, Clone, Debug)]
18enum SerialState {
19 ReadyForCommand,
20
21 ReceivingImage {
22 fast_refresh: bool,
23 index: usize,
24 },
25
26 FlashingProgram {
27 index: usize,
28 page: usize,
29 num_pages: Option<usize>,
30 remainder: Option<usize>,
31 },
32}
33
34fn erase_cycles(slot: u8) -> u32 {
35 let c = unsafe { u32::from_ne_bytes(*slot_ptr(slot).cast()) };
36 if c == u32::MAX {
37 0
38 } else {
39 c
40 }
41}
42
43fn best_slot() -> Option<u8> {
44 (1u8..=31)
45 .filter(|s| unsafe { !(*slot(*s)).is_valid() })
46 .map(|s| (s, erase_cycles(s)))
47 .min_by_key(|(_s, e)| *e)
48 .map(|(s, _e)| s)
49}
50
51unsafe fn write_flash(buf: &[u8], slot: u8, page: usize) {
52 erase_and_program((slot as u32) * 512 * 1024 + (page as u32) * 4096, buf);
53}
54
55fn write_all(serial: &mut SerialPort<UsbBus>, mut buf: &[u8]) {
56 while !buf.is_empty() {
57 let _ = serial.write(buf).map(|len| buf = &buf[len..]);
58 }
59}
60
61pub(crate) static NEEDS_REFRESH: AtomicBool = AtomicBool::new(false);
62pub(crate) static NEEDS_REFRESH_PROGRAMS: AtomicBool = AtomicBool::new(false);
63pub(crate) static HOST_APP: AtomicBool = AtomicBool::new(false);
64
65static PROG_SLOT: AtomicU8 = AtomicU8::new(0);
66
67pub(crate) extern "C" fn usb_handler() {
68 trace("USB handler");
69
70 static mut STATE: SerialState = SerialState::ReadyForCommand;
71 #[allow(static_mut_refs)]
72 let state = unsafe { &mut STATE };
73
74 #[allow(static_mut_refs)]
75 let dev: &mut UsbDevice<UsbBus> = unsafe { USB_DEVICE.as_mut().unwrap() };
76 #[allow(static_mut_refs)]
77 let serial: &mut SerialPort<UsbBus> = unsafe { USB_SERIAL.as_mut().unwrap() };
78
79 // Receive buffer. Size equal to IMAGE_BYTES so it can store an entire frame; also used for
80 // receiving flash applications.
81 static mut BUF: [u8; IMAGE_BYTES] = [0; IMAGE_BYTES];
82 #[allow(static_mut_refs)]
83 let buf = unsafe { &mut BUF };
84
85 if dev.poll(&mut [serial]) {
86 let mut s = heapless::String::<100>::new();
87 write!(s, "{state:?}").unwrap();
88 debug(&s);
89
90 match state {
91 SerialState::ReadyForCommand => {
92 let mut cmd_buf = [0u8];
93 if let Ok(count) = serial.read(&mut cmd_buf) {
94 if count == 0 {
95 return;
96 }
97
98 if HOST_APP.load(Ordering::Relaxed) {
99 match SerialCommand::try_from(cmd_buf[0]) {
100 Ok(SerialCommand::RefreshNormal) => *state = SerialState::ReceivingImage { fast_refresh: false, index: 0 },
101 Ok(SerialCommand::RefreshFast) => *state = SerialState::ReceivingImage { fast_refresh: true, index: 0 },
102 Ok(SerialCommand::ExitHostApp) => {
103 set_touch_enabled(true);
104 HOST_APP.store(false, Ordering::Relaxed);
105 NEEDS_REFRESH.store(true, Ordering::Relaxed);
106 write_all(serial, &[Response::Ack as u8]);
107 },
108 Ok(SerialCommand::NextEvent) => {
109 write_all(serial, &[Response::Ack as u8]);
110 write_all(serial, &postcard::to_vec::<_, 32>(&next_event()).unwrap());
111 },
112 Ok(SerialCommand::EnableTouch) => {
113 set_touch_enabled(true);
114 write_all(serial, &[Response::Ack as u8]);
115 },
116 Ok(SerialCommand::DisableTouch) => {
117 set_touch_enabled(false);
118 write_all(serial, &[Response::Ack as u8]);
119 },
120 Ok(SerialCommand::EnterHostApp | SerialCommand::GetProgramSlot | SerialCommand::UploadProgram) => {
121 write_all(serial, &[Response::IncorrectMode as u8]);
122 }
123 Err(_) => write_all(serial, &[Response::UnknownCommand as u8]),
124 }
125 } else {
126 match SerialCommand::try_from(cmd_buf[0]) {
127 Ok(SerialCommand::GetProgramSlot) => {
128 if let Some(slot) = best_slot() {
129 write_all(serial, &[Response::Ack as u8, slot]);
130 PROG_SLOT.store(slot, Ordering::Relaxed);
131 } else {
132 write_all(serial, &[Response::ProgramSlotsFull as u8]);
133 }
134 },
135 Ok(SerialCommand::UploadProgram) => {
136 if PROG_SLOT.load(Ordering::Relaxed) == 0 {
137 write_all(serial, &[Response::NoProgramSlot as u8]);
138 } else {
139 set_touch_enabled(false);
140 *state = SerialState::FlashingProgram { index: 0, page: 0, num_pages: None, remainder: None };
141 write_all(serial, &[Response::Ack as u8]);
142 }
143 },
144 Ok(SerialCommand::EnterHostApp) => {
145 HOST_APP.store(true, Ordering::Relaxed);
146 refresh(&[0u8; IMAGE_BYTES], false);
147 set_touch_enabled(false);
148 write_all(serial, &[Response::Ack as u8]);
149 },
150 Ok(
151 SerialCommand::RefreshNormal
152 | SerialCommand::RefreshFast
153 | SerialCommand::ExitHostApp
154 | SerialCommand::NextEvent
155 | SerialCommand::DisableTouch
156 | SerialCommand::EnableTouch
157 ) => write_all(serial, &[Response::IncorrectMode as u8]),
158 Err(_) => write_all(serial, &[Response::UnknownCommand as u8]),
159 }
160 }
161 }
162 }
163
164 SerialState::ReceivingImage { fast_refresh, index } => {
165 if let Ok(count) = serial.read(&mut buf[*index..]) {
166 *index += count;
167 if *index == IMAGE_BYTES {
168 refresh(buf, *fast_refresh);
169 write_all(serial, &[Response::Ack as u8]);
170 *state = SerialState::ReadyForCommand;
171 }
172 }
173 }
174
175 SerialState::FlashingProgram { index, page, num_pages, remainder } => {
176 let slot = PROG_SLOT.load(Ordering::Relaxed);
177
178 // Write page 0 last - this is the header, so we only want to write it once everything
179 // else is written successfully
180 // Keep page 0 in the first 4096 bytes of BUF for the end
181 if *page == 0 {
182 draw_flashing_ui(*page, None);
183 debug("receiving page 0");
184 if let Ok(count) = serial.read(&mut buf[*index..4096]) {
185 *index += count;
186
187 if num_pages.is_none() && *index >= 12 {
188 let mut b = [0u8; 4];
189 b.copy_from_slice(&buf[8..12]);
190 let num_bytes = usize::from_le_bytes(b);
191 *num_pages = Some(num_bytes.div_ceil(4096));
192 *remainder = Some(num_bytes % 4096);
193 }
194
195 if *index == 4096 {
196 *index = 0;
197 *page += 1;
198 }
199 }
200 } else {
201 draw_flashing_ui(*page, *num_pages);
202 if let Ok(count) = serial.read(&mut buf[(4096 + *index)..8192]) {
203 let mut message = heapless::String::<32>::new();
204 write!(message, "receiving page {page}").unwrap();
205 debug(&message);
206
207 *index += count;
208
209 let num_pages = num_pages.unwrap();
210 let remainder = remainder.unwrap();
211 if *index == 4096 || (*page == num_pages - 1 && *index == remainder) {
212 *index = 0;
213
214 // Actually write the flash page
215 debug("writing page");
216 unsafe { write_flash(&buf[4096..8192], slot, *page) };
217
218 *page += 1;
219 // If this is the last page, also flash the first page which we didn't
220 // do at the start
221 if *page == num_pages {
222 debug("finalising");
223
224 let erase_cycles = erase_cycles(slot);
225 buf[0..4].copy_from_slice(&(erase_cycles + 1).to_ne_bytes());
226
227 unsafe { write_flash(&buf[0..4096], slot, 0) };
228
229 let this_header = unsafe { header::slot(slot) };
230 let this_name = unsafe { core::slice::from_raw_parts((*this_header).name_ptr, (*this_header).name_len) };
231
232 // If there is an old program with the same name, delete it
233 Programs::new()
234 .filter_map(|prog| unsafe {
235 let name = core::slice::from_raw_parts((*prog).name_ptr, (*prog).name_len);
236 if (*prog).slot() != slot && name == this_name {
237 Some((*prog).slot())
238 } else {
239 None
240 }
241 })
242 .for_each(|slot| unsafe { delete_program(slot) });
243
244 PROG_SLOT.store(0, Ordering::Relaxed);
245
246 NEEDS_REFRESH_PROGRAMS.store(true, Ordering::Relaxed);
247 set_touch_enabled(true);
248 info("Finished writing program");
249
250 *state = SerialState::ReadyForCommand;
251 }
252 }
253 }
254 }
255 }
256 }
257 }
258}