espflash 2.0.0-rc.2

use std::{borrow::Cow, io::Write, iter::once, mem::size_of};

use bytemuck::bytes_of;

use super::{
    encode_flash_frequency, update_checksum, EspCommonHeader, ImageFormat, SegmentHeader,
    ESP_CHECKSUM_MAGIC, ESP_MAGIC,
};
use crate::{
    elf::{CodeSegment, FirmwareImage, RomSegment},
    error::{Error, FlashDetectError},
    flasher::{FlashFrequency, FlashMode, FlashSize},
    targets::Chip,
};

/// Image format for flashing to the ESP8266
pub struct Esp8266Format<'a> {
    irom_data: Option<RomSegment<'a>>,
    flash_segment: RomSegment<'a>,
    app_size: u32,
}

impl<'a> Esp8266Format<'a> {
    pub fn new(
        image: &'a dyn FirmwareImage<'a>,
        flash_mode: Option<FlashMode>,
        flash_size: Option<FlashSize>,
        flash_freq: Option<FlashFrequency>,
    ) -> Result<Self, Error> {
        // IROM goes into a separate plain binary
        let irom_data = merge_rom_segments(image.rom_segments(Chip::Esp8266));

        let mut common_data = Vec::with_capacity(
            image
                .ram_segments(Chip::Esp8266)
                .map(|segment| segment.size() as usize)
                .sum(),
        );

        // Common header
        let flash_mode = flash_mode.unwrap_or_default() as u8;
        let flash_freq = flash_freq.unwrap_or_default();
        let flash_size = flash_size.unwrap_or_default();
        let flash_config =
            encode_flash_size(flash_size)? + encode_flash_frequency(Chip::Esp8266, flash_freq)?;
        let segment_count = image.ram_segments(Chip::Esp8266).count() as u8;

        let header = EspCommonHeader {
            magic: ESP_MAGIC,
            segment_count,
            flash_mode,
            flash_config,
            entry: image.entry(),
        };
        common_data.write_all(bytes_of(&header))?;

        let mut total_len = 8;
        let mut checksum = ESP_CHECKSUM_MAGIC;

        for segment in image.ram_segments(Chip::Esp8266) {
            let data = segment.data();
            let padding = 4 - data.len() % 4;

            let segment_header = SegmentHeader {
                addr: segment.addr,
                length: (data.len() + padding) as u32,
            };

            total_len += size_of::<SegmentHeader>() as u32 + segment_header.length;

            common_data.write_all(bytes_of(&segment_header))?;
            common_data.write_all(data)?;

            let padding = &[0u8; 4][0..padding];
            common_data.write_all(padding)?;

            checksum = update_checksum(data, checksum);
        }

        let padding = 15 - (total_len % 16);
        let padding = &[0u8; 16][0..padding as usize];

        common_data.write_all(padding)?;
        common_data.write_all(&[checksum])?;

        let flash_segment = RomSegment {
            addr: 0,
            data: Cow::Owned(common_data),
        };

        let app_size = irom_data
            .clone()
            .map(|d| d.data.len() as u32)
            .unwrap_or_default()
            + flash_segment.data.len() as u32;

        Ok(Self {
            irom_data,
            flash_segment,
            app_size,
        })
    }
}

impl<'a> ImageFormat<'a> for Esp8266Format<'a> {
    fn flash_segments<'b>(&'b self) -> Box<dyn Iterator<Item = RomSegment<'b>> + 'b>
    where
        'a: 'b,
    {
        Box::new(
            self.irom_data
                .iter()
                .map(RomSegment::borrow)
                .chain(once(self.flash_segment.borrow())),
        )
    }

    fn ota_segments<'b>(&'b self) -> Box<dyn Iterator<Item = RomSegment<'b>> + 'b>
    where
        'a: 'b,
    {
        Box::new(
            self.irom_data
                .iter()
                .map(RomSegment::borrow)
                .chain(once(self.flash_segment.borrow())),
        )
    }

    fn app_size(&self) -> u32 {
        self.app_size
    }

    fn part_size(&self) -> Option<u32> {
        None
    }
}

fn merge_rom_segments<'a>(
    mut segments: impl Iterator<Item = CodeSegment<'a>>,
) -> Option<RomSegment<'a>> {
    const IROM_MAP_START: u32 = 0x4020_0000;

    let first = segments.next()?;
    let data = if let Some(second) = segments.next() {
        let mut data = Vec::with_capacity(first.data().len() + second.data().len());
        data.extend_from_slice(first.data());

        for segment in once(second).chain(segments) {
            let padding_size = segment.addr as usize - first.addr as usize - data.len();
            data.resize(data.len() + padding_size, 0);
            data.extend_from_slice(segment.data());
        }

        data
    } else {
        first.data().into()
    };

    Some(RomSegment {
        addr: first.addr - IROM_MAP_START,
        data: Cow::Owned(data),
    })
}

fn encode_flash_size(size: FlashSize) -> Result<u8, FlashDetectError> {
    use FlashSize::*;

    match size {
        Flash256Kb => Ok(0x10),
        Flash512Kb => Ok(0x00),
        Flash1Mb => Ok(0x20),
        Flash2Mb => Ok(0x30),
        Flash4Mb => Ok(0x40),
        Flash8Mb => Ok(0x80),
        Flash16Mb => Ok(0x90),
        _ => Err(FlashDetectError::from(size as u8)),
    }
}

#[cfg(test)]
mod tests {
    use std::fs;

    use super::*;
    use crate::elf::ElfFirmwareImage;

    #[test]
    fn test_esp8266_image_format() {
        let input_bytes = fs::read("tests/resources/esp8266_hal_blinky").unwrap();
        let expected_bin = fs::read("tests/resources/esp8266_hal_blinky.bin").unwrap();

        let image = ElfFirmwareImage::try_from(input_bytes.as_slice()).unwrap();
        let flash_image = Esp8266Format::new(&image, None, None, None).unwrap();

        let segments = flash_image.flash_segments().collect::<Vec<_>>();
        let buf = segments[0].data.as_ref();

        assert_eq!(expected_bin.len(), buf.len());
        assert_eq!(expected_bin.as_slice(), buf);
    }
}