rld/src/common/loadable.rs

use std::collections::HashMap;
use std::iter::once;

use crate::error::Error;
use crate::{
    common::{Relocatable, Section, SectionInfo, SymbolIndex, Symbols},
    error::LinkError,
};

use super::SymbolValue;

pub enum SegmentType {
    Text,
    Data,
    Bss,
}

pub struct SegmentSection {
    si: SectionInfo,
    output_offset: usize,
}

impl SegmentSection {
    // TODO: refactor this into newtype swap with compile-time checking
    pub fn section_info(&mut self, output_offset: usize) -> Result<&SectionInfo, Error> {
        if self.output_offset > 0 {
            Err(Error::InvalidSectionIndex) // duplicate access??
        } else {
            self.output_offset = output_offset;
            Ok(&mut self.si)
        }
    }
}

impl From<SectionInfo> for SegmentSection {
    fn from(si: SectionInfo) -> Self {
        Self {
            si,
            output_offset: 0,
        }
    }
}

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

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

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

    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut SegmentSection> {
        self.sections_info.iter_mut()
    }

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

    pub fn unresolved_offsets<'s>(
        &'s self,
        symbols: &'s Symbols,
    ) -> impl Iterator<Item = usize> + 's {
        symbols
            .unresolved()
            .map(move |si| {
                self.sections_info.iter().filter_map(move |section| {
                    match section.si.section_index == si.section_index {
                        true => Some(section.output_offset + si.offsets.symbol_offset),
                        false => None,
                    }
                })
            })
            .flatten()
    }
}

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

    text: SegmentSections,
    rodata: SegmentSections,
    data: SegmentSections,
    bss: SegmentSections,
    symbol_map: HashMap<String, Symbols>,
}

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, si: SymbolIndex, objects: &[R]) -> Result<(), Error>
    where
        R: Relocatable,
    {
        let object = objects
            .get(si.object_index)
            .ok_or(Error::InvalidSectionIndex)?;

        if !object.symbol_needs_resolving(si.symbol_index) {
            return Ok(()); // local symbols are not stored
        }

        let name = object.symbol_name(si.symbol_index)?;
        eprint!("OBJECT: {:?}\nNAME: {}/", object.origin(), name);
        eprintln!("{:?}", object.symbol_value(si.symbol_index)?);

        if let Some(existing) = self.symbol_map.get_mut(name) {
            existing.push(si, objects)
        } else {
            let mut symbols = Symbols::default();
            symbols.push(si, objects)?;
            self.symbol_map.insert(name.to_owned(), symbols);

            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_sections(&mut self) -> impl Iterator<Item = &mut SegmentSection> {
        let text_iter = self.text.iter_mut();
        let rodata_iter = self.rodata.iter_mut();

        let data1 = text_iter.filter_map(move |ss| match ss.si.file_size {
            0 => None,
            _ => Some(ss),
        });
        let data2 = rodata_iter.filter_map(move |ss| match ss.si.data_size {
            0 => None,
            _ => Some(ss),
        });

        data1.chain(data2)
    }

    pub fn data_sections(&mut self) -> impl Iterator<Item = &mut SegmentSection> {
        let iter = self
            .data
            .iter_mut()
            .filter_map(move |ss| match ss.si.file_size {
                0 => None,
                _ => Some(ss),
            });

        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 ss in text_iter.chain(rodata_iter) {
            result += ss.si.data_size
        }

        result
    }

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

        result
    }

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

        result
    }

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

    pub fn start_offset(&self) -> Result<usize, Error> {
        self.start_offset.ok_or_else(|| {
            let link_error = LinkError {
                message: "Program entrypoint not found".into(),
                ..Default::default()
            };

            Error::LinkingError(link_error)
        })
    }

    pub fn symbol_resolutions<'o, R>(
        &'o self,
        objects: &'o [R],
    ) -> impl Iterator<Item = SymbolResolution<'o>>
    where
        R: Relocatable,
    {
        self.symbol_map.values().map(move |symbols| {
            let value = symbols.value(objects).expect("value"); // TODO
            let offsets = self.unresolved_offsets(symbols).collect();

            SymbolResolution { value, offsets }
        })
    }

    fn unresolved_offsets<'s>(&'s self, symbols: &'s Symbols) -> impl Iterator<Item = usize> + 's {
        self.text
            .unresolved_offsets(symbols)
            .chain(self.rodata.unresolved_offsets(symbols))
            .chain(self.data.unresolved_offsets(symbols))
            .chain(self.bss.unresolved_offsets(symbols))
    }
}

pub struct SymbolResolution<'b> {
    pub value: SymbolValue<'b>,
    pub offsets: Vec<usize>, // TODO: box iterator instead
}