rld/src/common/loadable.rs

use std::iter::once;

use crate::common::{Relocatable, Section, SectionInfo, Symbol};
use crate::error::Error;

pub enum SegmentType {
    Text,
    Data,
    Bss,
}

#[derive(Default)]
pub struct SegmentSections {
    sections_info: Vec<SectionInfo>,
    data_size: u64,
}

impl SegmentSections {
    pub fn push(&mut self, si: SectionInfo) {
        self.data_size += si.data_size;
        self.sections_info.push(si);
    }

    pub fn iter(&self) -> impl Iterator<Item = &SectionInfo> {
        self.sections_info.iter()
    }

    pub fn data_size(&self) -> u64 {
        self.data_size
    }
}

#[derive(Default)]
pub struct Loadable {
    start_offset: Option<u64>,

    text: SegmentSections,
    rodata: SegmentSections,
    data: SegmentSections,
    bss: SegmentSections,
    // symbol_map: HashMap<Rc<str>, Symbol>,
}

impl Loadable {
    pub fn process_section(&mut self, section: Section) -> Result<(), Error> {
        match section {
            Section::Text(si) => self.text.push(si),
            Section::Data(si, true) => self.rodata.push(si),
            Section::Data(si, false) => self.data.push(si),
            Section::Bss(si) => self.bss.push(si),
        }

        Ok(())
    }

    pub fn process_symbol<R>(&mut self, _symbol: Symbol, _objects: &[R]) -> Result<(), Error>
    where
        R: Relocatable,
    {
        Ok(())
        // if symbol.binding == Binding::Local {
        //     return Ok(()); // local symbols are not stored
        // }

        // let key = Rc::from(symbol.name.as_str());
        // let object_index = symbol.object_index;

        // if let Some(prev) = self.symbol_map.insert(key, symbol) {
        //     let prev_trace = Trace {
        //         origin: objects[prev.object_index].origin().into(),
        //         ..Default::default()
        //     };
        //     let symbol_trace = Trace {
        //         origin: objects[object_index].origin().into(),
        //         ..Default::default()
        //     };
        //     let link_error = LinkError {
        //         message: format!("Duplicate global symbol found: {}", prev.name),
        //         code: LE_GLOBAL_SYMBOL_DUPLICATE,
        //         traces: vec![prev_trace, symbol_trace],
        //     };
        //     Err(Error::LinkingError(link_error))
        // } else {
        //     Ok(())
        // }
    }

    pub fn segment_sections(&self) -> impl Iterator<Item = &SegmentSections> {
        let text = once(&self.text);
        let rodata = once(&self.rodata);
        let data = once(&self.data);
        let bss = once(&self.bss);

        text.chain(rodata).chain(data).chain(bss)
    }

    pub fn program_bytes<'l, R: Relocatable>(
        &'l self,
        objects: &'l [R],
    ) -> impl Iterator<Item = Result<&'l [u8], Error>> {
        let text_iter = self.text.iter();
        let rodata_iter = self.rodata.iter();

        let data1 = text_iter.filter_map(move |si| match si.file_size {
            0 => None,
            _ => Some(objects[si.object_index].bytes(si.offset, si.file_size)),
        });
        let data2 = rodata_iter.filter_map(move |si| match si.data_size {
            0 => None,
            _ => Some(objects[si.object_index].bytes(si.offset, si.file_size)),
        });

        data1.chain(data2)
    }

    pub fn data_bytes<'l, R: Relocatable>(
        &'l self,
        objects: &'l [R],
    ) -> impl Iterator<Item = Result<&'l [u8], Error>> {
        let iter = self.data.iter().filter_map(move |si| match si.file_size {
            0 => None,
            _ => Some(objects[si.object_index].bytes(si.offset, si.file_size)),
        });

        iter
    }

    // .text + .rodata
    pub fn program_size(&self) -> u64 {
        let text_iter = self.text.iter();
        let rodata_iter = self.rodata.iter();

        let mut result = 0u64;
        for section in text_iter.chain(rodata_iter) {
            result += section.data_size
        }

        result
    }

    // data
    pub fn data_size(&self) -> u64 {
        let mut result = 0u64;
        for section in self.data.iter() {
            result += section.data_size
        }

        result
    }

    // bss
    pub fn bss_size(&self) -> u64 {
        let mut result = 0u64;
        for section in self.bss.iter() {
            result += section.data_size
        }

        result
    }

    pub fn set_start_offset(&mut self, offset: u64) {
        self.start_offset = Some(offset);
    }

    pub fn start_offset(&self) -> Result<u64, Error> {
        use crate::error::LinkError;

        self.start_offset.ok_or_else(|| {
            let link_error = LinkError {
                message: "Program entrypoint not found".into(),
                ..Default::default()
            };

            Error::LinkingError(link_error)
        })
    }
}