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