storage.rs

use core::slice::from_raw_parts;

use heapless::{String, Vec};
use rp2040_flash::flash;

const SECTOR_SIZE: usize = 4096;
const FLASH_START: usize = 0x10000000;
// First 6MB is for code
const DATA_START: usize = FLASH_START + 6144 * 1024;
// 2MB for data
const DATA_END: usize = DATA_START + 2048 * 1024;

const CONTACT_HEADER_START: usize = DATA_START;
const CONTACT_HEADER_SIZE: usize = 3;

pub const MAX_CONTACTS: usize = 100;
pub const MAX_CONTACT_NAME_LENGTH: usize = 20;
pub const MAX_CONTACT_CALLSIGN_LENGTH: usize = 10;
// 1 byte for each length, 1 byte for SSID, 2 bytes for checksum
const CONTACT_LENGTH: usize = MAX_CONTACT_NAME_LENGTH + 1 + 1 + MAX_CONTACT_CALLSIGN_LENGTH + 1 + 2;

const X25: crc::Crc<u16> = crc::Crc::<u16>::new(&crc::CRC_16_IBM_SDLC);

pub struct Storage {
    contact_header: ContactHeader,
}

impl Storage {
    pub fn new() -> Self {
        Self {
            // TODO need to invalidate when writing contacts
            contact_header: ContactHeader::read(),
        }
    }

    pub fn save_contact(&mut self, contact: Contact) {
        let len = self.contact_header.num_contacts;

        // if id isn't less than the len, we're inserting a new contact
        if contact.id >= usize::from(len) {
            assert!(usize::from(len) < MAX_CONTACTS);
            self.contact_header.num_contacts += 1;
            let ch = self.contact_header;
            ch.write(self);
        }

        contact.write(self);
    }

    pub fn delete_contact(&mut self, contact_id: usize) {
        assert!(contact_id < self.contacts_len());

        // Shift every contact down
        for mut c in self.contacts().skip(contact_id + 1) {
            c.id -= 1;
            c.write(self);
        }

        self.contact_header.num_contacts -= 1;
        let ch = self.contact_header;
        ch.write(self);
    }

    pub fn contacts(&self) -> ContactIter {
        ContactIter::new(self.contact_header.num_contacts.into())
    }

    pub fn contacts_len(&self) -> usize {
        self.contact_header.num_contacts.into()
    }

    // Note - calling this twice without saving one
    // will cause a duplicate ID
    // We don't increment the len until the contact is actually saved
    pub fn new_contact(&self) -> Contact {
        Contact {
            callsign: String::new(),
            ssid: 0,
            name: String::new(),
            id: usize::from(self.contact_header.num_contacts),
        }
    }

    fn write_data(&mut self, start_addr: usize, buf: &[u8]) {
        let start_sector = start_addr / SECTOR_SIZE;
        let num_sectors = buf.len() / SECTOR_SIZE + 1;
        let diff = start_addr % SECTOR_SIZE;

        for s in 0..num_sectors {
            let existing_start = if s == 0 { diff } else { 0 };
            let existing_end = (diff + (s + 1) * SECTOR_SIZE).min(existing_start + buf.len());
            let mut existing = self.read_sector((start_sector + s) * SECTOR_SIZE);
            let buf_start = (s * SECTOR_SIZE).min(buf.len());
            let buf_end = (buf_start + SECTOR_SIZE).min(buf.len());
            // only write if there's a difference
            if existing[existing_start..existing_end] != buf[buf_start..buf_end] {
                existing[existing_start..existing_end].copy_from_slice(&buf[buf_start..buf_end]);
                self.write_sector((start_sector + s) * SECTOR_SIZE, &existing);
            }
        }
    }

    fn read_sector(&self, addr: usize) -> [u8; SECTOR_SIZE] {
        assert!(addr >= DATA_START);
        assert!(addr + SECTOR_SIZE < DATA_END);

        let sector = unsafe { from_raw_parts(addr as *mut u8, SECTOR_SIZE) };

        let mut buf = [0; SECTOR_SIZE];

        buf.copy_from_slice(sector);

        buf
    }

    fn write_sector(&mut self, start_addr: usize, buf: &[u8; 4096]) {
        let addr: u32 = (start_addr - FLASH_START).try_into().unwrap();
        let sector_size = u32::try_from(SECTOR_SIZE).unwrap();
        assert!(addr % sector_size == 0);

        cortex_m::interrupt::free(|_| unsafe {
            flash::flash_range_erase_and_program(addr, buf, true);
        });
    }
}

// Up to 100 contacts
// 2 byte CRC checksum
#[derive(Copy, Clone)]
struct ContactHeader {
    num_contacts: u8,
}

impl ContactHeader {
    fn read() -> Self {
        let bytes = unsafe { from_raw_parts(CONTACT_HEADER_START as *mut u8, CONTACT_HEADER_SIZE) };
        let num_contacts = bytes[0];
        let checksum_expected = u16::from_le_bytes([bytes[1], bytes[2]]);
        // CRC checksum is very overkill right now
        // But we may want to add more fields in the future
        let checksum_actual = X25.checksum(&[num_contacts]);

        if checksum_expected == checksum_actual && usize::from(num_contacts) < MAX_CONTACTS {
            Self { num_contacts }
        } else {
            Self { num_contacts: 0 }
        }
    }

    fn write(&self, storage: &mut Storage) {
        let mut buf = [0; CONTACT_HEADER_SIZE];
        buf[0] = self.num_contacts;
        let checksum = X25.checksum(&[self.num_contacts]);
        buf[1..].copy_from_slice(&checksum.to_le_bytes());

        storage.write_data(CONTACT_HEADER_START, &buf);
    }
}

#[derive(Clone)]
pub struct Contact {
    pub callsign: String<MAX_CONTACT_CALLSIGN_LENGTH>,
    pub ssid: u8,
    pub name: String<MAX_CONTACT_NAME_LENGTH>,
    id: usize,
}

impl Contact {
    fn read(idx: usize) -> Option<Self> {
        let start_addr =
            (CONTACT_HEADER_START + CONTACT_HEADER_SIZE + idx * CONTACT_LENGTH) as *mut u8;

        let bytes = unsafe { from_raw_parts(start_addr, CONTACT_LENGTH) };
        let checksum_expected =
            u16::from_le_bytes([bytes[CONTACT_LENGTH - 2], bytes[CONTACT_LENGTH - 1]]);
        let checksum_actual = X25.checksum(&bytes[..(CONTACT_LENGTH - 2)]);

        if checksum_expected != checksum_actual {
            return None;
        }

        let callsign_len: usize = bytes[0].into();
        let name_len: usize = bytes[1].into();
        let ssid = bytes[2];

        if callsign_len > MAX_CONTACT_CALLSIGN_LENGTH || name_len > MAX_CONTACT_NAME_LENGTH {
            return None;
        }

        let callsign_bytes = &bytes[3..(callsign_len + 3)];
        let name_bytes =
            &bytes[(3 + MAX_CONTACT_NAME_LENGTH)..(3 + MAX_CONTACT_NAME_LENGTH + name_len)];

        let callsign = String::from_utf8(Vec::from_slice(callsign_bytes).ok()?).ok()?;
        let name = String::from_utf8(Vec::from_slice(name_bytes).ok()?).ok()?;

        Some(Self {
            callsign,
            ssid,
            name,
            id: idx,
        })
    }

    fn write(&self, storage: &mut Storage) {
        let start_addr = CONTACT_HEADER_START + CONTACT_HEADER_SIZE + self.id * CONTACT_LENGTH;

        let mut buf = [0; CONTACT_LENGTH];
        buf[0] = self.callsign.len().try_into().unwrap();
        buf[1] = self.name.len().try_into().unwrap();
        buf[2] = self.ssid;

        let callsign_bytes = self.callsign.as_bytes();
        let name_bytes = self.name.as_bytes();
        buf[3..(3 + callsign_bytes.len())].copy_from_slice(callsign_bytes);
        buf[(3 + MAX_CONTACT_NAME_LENGTH)..(3 + MAX_CONTACT_NAME_LENGTH + name_bytes.len())]
            .copy_from_slice(name_bytes);
        let checksum = X25.checksum(&buf[..(CONTACT_LENGTH - 2)]);
        buf[(CONTACT_LENGTH - 2)..].copy_from_slice(&checksum.to_le_bytes());

        storage.write_data(start_addr, &buf);
    }
}

#[derive(Copy, Clone)]
pub struct ContactIter {
    i: usize,
    len: usize,
}

impl ContactIter {
    fn new(len: usize) -> Self {
        Self { i: 0, len }
    }
}

impl Iterator for ContactIter {
    type Item = Contact;

    fn next(&mut self) -> Option<Self::Item> {
        while self.i < self.len {
            let c = Contact::read(self.i);

            self.i += 1;

            if let Some(c) = c {
                return Some(c);
            }
        }

        None
    }
}