synth-backend 0.11.8

//! Memory Layout Analyzer
//!
//! Analyzes WebAssembly modules and generates memory layouts for embedded targets

use synth_core::{Component, Error, HardwareCapabilities, Result};

/// Memory section type
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SectionType {
    /// Executable code (.text)
    Text,
    /// Read-only data (.rodata)
    ReadOnlyData,
    /// Initialized data (.data)
    Data,
    /// Uninitialized data (.bss)
    Bss,
    /// Stack
    Stack,
    /// Heap
    Heap,
}

/// Memory section in the layout
#[derive(Debug, Clone)]
pub struct MemorySection {
    /// Section type
    pub section_type: SectionType,

    /// Section name
    pub name: String,

    /// Base address
    pub base_address: u32,

    /// Size in bytes
    pub size: u32,

    /// Alignment requirement
    pub alignment: u32,

    /// Whether section is in flash (XIP) or RAM
    pub in_flash: bool,
}

impl MemorySection {
    /// Get end address (exclusive)
    pub fn end_address(&self) -> u32 {
        self.base_address + self.size
    }

    /// Check if this section overlaps with another
    pub fn overlaps(&self, other: &MemorySection) -> bool {
        let self_end = self.end_address();
        let other_end = other.end_address();

        !(self_end <= other.base_address || other_end <= self.base_address)
    }
}

/// Memory layout for a WebAssembly module
#[derive(Debug, Clone)]
pub struct MemoryLayout {
    /// Hardware capabilities
    hw_caps: HardwareCapabilities,

    /// Sections in the layout
    sections: Vec<MemorySection>,

    /// Total flash usage
    flash_usage: u32,

    /// Total RAM usage
    ram_usage: u32,
}

impl MemoryLayout {
    /// Create a new memory layout
    pub fn new(hw_caps: HardwareCapabilities) -> Self {
        Self {
            hw_caps,
            sections: Vec::new(),
            flash_usage: 0,
            ram_usage: 0,
        }
    }

    /// Add a section to the layout
    pub fn add_section(&mut self, section: MemorySection) -> Result<()> {
        // Check for overlaps
        for existing in &self.sections {
            if section.overlaps(existing) {
                return Err(Error::MemoryLayoutError(format!(
                    "Section '{}' at 0x{:08X} overlaps with '{}' at 0x{:08X}",
                    section.name, section.base_address, existing.name, existing.base_address
                )));
            }
        }

        // Update usage counters
        if section.in_flash {
            self.flash_usage += section.size;
        } else {
            self.ram_usage += section.size;
        }

        self.sections.push(section);
        Ok(())
    }

    /// Get all sections
    pub fn sections(&self) -> &[MemorySection] {
        &self.sections
    }

    /// Get flash usage
    pub fn flash_usage(&self) -> u32 {
        self.flash_usage
    }

    /// Get RAM usage
    pub fn ram_usage(&self) -> u32 {
        self.ram_usage
    }

    /// Validate layout against hardware capabilities
    pub fn validate(&self) -> Result<()> {
        // Check flash capacity
        if self.flash_usage as u64 > self.hw_caps.flash_size {
            return Err(Error::MemoryLayoutError(format!(
                "Flash usage {} bytes exceeds capacity {} bytes",
                self.flash_usage, self.hw_caps.flash_size
            )));
        }

        // Check RAM capacity
        if self.ram_usage as u64 > self.hw_caps.ram_size {
            return Err(Error::MemoryLayoutError(format!(
                "RAM usage {} bytes exceeds capacity {} bytes",
                self.ram_usage, self.hw_caps.ram_size
            )));
        }

        Ok(())
    }

    /// Get section by type
    pub fn get_section(&self, section_type: SectionType) -> Option<&MemorySection> {
        self.sections
            .iter()
            .find(|s| s.section_type == section_type)
    }

    /// Generate GNU LD linker script for ARM Cortex-M
    pub fn generate_linker_script(&self) -> String {
        let mut script = String::new();

        script.push_str("/* Linker script for ARM Cortex-M */\n");
        script.push_str("/* Generated by Synth WebAssembly Component Synthesizer */\n\n");

        // Memory regions
        script.push_str("MEMORY\n{\n");
        script.push_str(&format!(
            "  FLASH (rx) : ORIGIN = 0x00000000, LENGTH = {}K\n",
            self.hw_caps.flash_size / 1024
        ));
        script.push_str(&format!(
            "  RAM (rwx)  : ORIGIN = 0x20000000, LENGTH = {}K\n",
            self.hw_caps.ram_size / 1024
        ));
        script.push_str("}\n\n");

        // Entry point
        script.push_str("ENTRY(Reset_Handler)\n\n");

        // Stack size
        let stack_section = self.get_section(SectionType::Stack);
        let stack_size = stack_section.map(|s| s.size).unwrap_or(4096);
        script.push_str(&format!("_stack_size = {};\n\n", stack_size));

        // Heap size
        let heap_section = self.get_section(SectionType::Heap);
        let heap_size = heap_section.map(|s| s.size).unwrap_or(8192);
        script.push_str(&format!("_heap_size = {};\n\n", heap_size));

        // Sections
        script.push_str("SECTIONS\n{\n");

        // .text section
        script.push_str("  .text :\n");
        script.push_str("  {\n");
        script.push_str("    KEEP(*(.isr_vector))\n");
        script.push_str("    *(.text*)\n");
        script.push_str("    *(.rodata*)\n");
        script.push_str("    KEEP(*(.init))\n");
        script.push_str("    KEEP(*(.fini))\n");
        script.push_str("    . = ALIGN(4);\n");
        script.push_str("    _etext = .;\n");
        script.push_str("  } > FLASH\n\n");

        // .ARM.exidx section (for exception handling)
        script.push_str("  .ARM.exidx :\n");
        script.push_str("  {\n");
        script.push_str("    __exidx_start = .;\n");
        script.push_str("    *(.ARM.exidx* .gnu.linkonce.armexidx.*)\n");
        script.push_str("    __exidx_end = .;\n");
        script.push_str("  } > FLASH\n\n");

        // .data section
        script.push_str("  .data :\n");
        script.push_str("  {\n");
        script.push_str("    _sdata = .;\n");
        script.push_str("    *(.data*)\n");
        script.push_str("    . = ALIGN(4);\n");
        script.push_str("    _edata = .;\n");
        script.push_str("  } > RAM AT> FLASH\n\n");
        script.push_str("  _sidata = LOADADDR(.data);\n\n");

        // .bss section
        script.push_str("  .bss :\n");
        script.push_str("  {\n");
        script.push_str("    _sbss = .;\n");
        script.push_str("    *(.bss*)\n");
        script.push_str("    *(COMMON)\n");
        script.push_str("    . = ALIGN(4);\n");
        script.push_str("    _ebss = .;\n");
        script.push_str("  } > RAM\n\n");

        // .heap section
        script.push_str("  .heap :\n");
        script.push_str("  {\n");
        script.push_str("    _heap_start = .;\n");
        script.push_str("    . = . + _heap_size;\n");
        script.push_str("    _heap_end = .;\n");
        script.push_str("  } > RAM\n\n");

        // .stack section
        script.push_str("  .stack :\n");
        script.push_str("  {\n");
        script.push_str("    . = . + _stack_size;\n");
        script.push_str("    _stack_top = .;\n");
        script.push_str("  } > RAM\n\n");

        script.push_str("  /* Remove information from the standard libraries */\n");
        script.push_str("  /DISCARD/ :\n");
        script.push_str("  {\n");
        script.push_str("    libc.a ( * )\n");
        script.push_str("    libm.a ( * )\n");
        script.push_str("    libgcc.a ( * )\n");
        script.push_str("  }\n");

        script.push_str("}\n");

        script
    }
}

/// Memory layout analyzer
pub struct MemoryLayoutAnalyzer {
    hw_caps: HardwareCapabilities,
}

impl MemoryLayoutAnalyzer {
    /// Create a new analyzer
    pub fn new(hw_caps: HardwareCapabilities) -> Self {
        Self { hw_caps }
    }

    /// Analyze a component and generate memory layout
    pub fn analyze(&self, component: &Component) -> Result<MemoryLayout> {
        let mut layout = MemoryLayout::new(self.hw_caps.clone());

        // Calculate sizes for each section
        let text_size = self.estimate_text_size(component);
        let rodata_size = self.estimate_rodata_size(component);
        let data_size = self.estimate_data_size(component);
        let bss_size = self.estimate_bss_size(component);
        let stack_size = self.estimate_stack_size(component);
        let heap_size = self.estimate_heap_size(component);

        // Allocate sections
        // For XIP (Execute In Place), put .text and .rodata in flash
        let flash_base = 0x00000000;
        let ram_base = 0x20000000;

        let mut current_flash = flash_base;
        let mut current_ram = ram_base;

        // .text section in flash
        if text_size > 0 {
            let text_section = MemorySection {
                section_type: SectionType::Text,
                name: ".text".to_string(),
                base_address: current_flash,
                size: text_size,
                alignment: 4,
                in_flash: true,
            };
            current_flash = align_up(text_section.end_address(), 4);
            layout.add_section(text_section)?;
        }

        // .rodata section in flash
        if rodata_size > 0 {
            let rodata_section = MemorySection {
                section_type: SectionType::ReadOnlyData,
                name: ".rodata".to_string(),
                base_address: current_flash,
                size: rodata_size,
                alignment: 4,
                in_flash: true,
            };
            layout.add_section(rodata_section)?;
        }

        // .data section in RAM (but stored in flash, copied at startup)
        if data_size > 0 {
            let data_section = MemorySection {
                section_type: SectionType::Data,
                name: ".data".to_string(),
                base_address: current_ram,
                size: data_size,
                alignment: 4,
                in_flash: false,
            };
            current_ram = align_up(data_section.end_address(), 4);
            layout.add_section(data_section)?;
        }

        // .bss section in RAM
        if bss_size > 0 {
            let bss_section = MemorySection {
                section_type: SectionType::Bss,
                name: ".bss".to_string(),
                base_address: current_ram,
                size: bss_size,
                alignment: 4,
                in_flash: false,
            };
            current_ram = align_up(bss_section.end_address(), 4);
            layout.add_section(bss_section)?;
        }

        // Heap in RAM
        if heap_size > 0 {
            let heap_section = MemorySection {
                section_type: SectionType::Heap,
                name: ".heap".to_string(),
                base_address: current_ram,
                size: heap_size,
                alignment: 8,
                in_flash: false,
            };
            layout.add_section(heap_section)?;
        }

        // Stack in RAM (grows downward, so we place it at the end)
        if stack_size > 0 {
            let stack_base = self.hw_caps.ram_size as u32 - stack_size;
            let stack_section = MemorySection {
                section_type: SectionType::Stack,
                name: ".stack".to_string(),
                base_address: stack_base,
                size: stack_size,
                alignment: 8,
                in_flash: false,
            };
            layout.add_section(stack_section)?;
        }

        // Validate the layout
        layout.validate()?;

        Ok(layout)
    }

    /// Estimate .text section size
    fn estimate_text_size(&self, component: &Component) -> u32 {
        // Rough estimate: assume 1 WASM instruction = 2-4 ARM instructions
        // and each ARM instruction is 2 or 4 bytes (Thumb-2)
        // For now, use a conservative estimate of 8 bytes per function
        let mut size = 0u32;

        for module in &component.modules {
            // Estimate based on number of functions
            size += (module.functions.len() as u32) * 128; // 128 bytes per function average
        }

        // Round up to alignment
        align_up(size, 4)
    }

    /// Estimate .rodata section size
    fn estimate_rodata_size(&self, component: &Component) -> u32 {
        let mut size = 0u32;

        for module in &component.modules {
            // Estimate from global data
            size += (module.globals.len() as u32) * 4;
        }

        align_up(size, 4)
    }

    /// Estimate .data section size
    fn estimate_data_size(&self, component: &Component) -> u32 {
        let mut size = 0u32;

        for module in &component.modules {
            // Data segments from WebAssembly linear memory
            for memory in &module.memories {
                size += memory.initial * 65536; // Pages to bytes
            }
        }

        align_up(size, 4)
    }

    /// Estimate .bss section size
    fn estimate_bss_size(&self, _component: &Component) -> u32 {
        // For now, allocate a fixed amount for uninitialized data
        align_up(4096, 4)
    }

    /// Estimate stack size
    fn estimate_stack_size(&self, _component: &Component) -> u32 {
        // Conservative estimate: 256 bytes per stack frame, max depth of 16
        let stack_size = 256 * 16;

        align_up(stack_size, 8)
    }

    /// Estimate heap size
    fn estimate_heap_size(&self, _component: &Component) -> u32 {
        // Allocate remaining RAM after other sections
        // For now, use a conservative 8KB
        align_up(8192, 8)
    }
}

/// Align value up to alignment
fn align_up(value: u32, alignment: u32) -> u32 {
    (value + alignment - 1) & !(alignment - 1)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;
    use synth_core::{CoreModule, Function, FunctionSignature, Global, Memory, ValueType};

    fn test_component() -> Component {
        Component {
            name: "test".to_string(),
            modules: vec![CoreModule {
                id: "test_module".to_string(),
                binary: vec![],
                functions: vec![Function {
                    index: 0,
                    name: Some("add".to_string()),
                    signature: FunctionSignature {
                        params: vec![ValueType::I32, ValueType::I32],
                        results: vec![ValueType::I32],
                    },
                    exported: true,
                    imported: false,
                }],
                memories: vec![Memory {
                    index: 0,
                    initial: 1, // 64KB
                    maximum: None,
                    shared: false,
                    memory64: false,
                }],
                tables: vec![],
                globals: vec![Global {
                    index: 0,
                    value_type: ValueType::I32,
                    mutable: false,
                }],
            }],
            components: vec![],
            instances: vec![],
            interfaces: HashMap::new(),
            imports: vec![],
            exports: vec![],
        }
    }

    #[test]
    fn test_memory_section_overlap() {
        let section1 = MemorySection {
            section_type: SectionType::Text,
            name: ".text".to_string(),
            base_address: 0x00000000,
            size: 1024,
            alignment: 4,
            in_flash: true,
        };

        let section2 = MemorySection {
            section_type: SectionType::ReadOnlyData,
            name: ".rodata".to_string(),
            base_address: 0x00000400,
            size: 512,
            alignment: 4,
            in_flash: true,
        };

        let section3 = MemorySection {
            section_type: SectionType::Data,
            name: ".data".to_string(),
            base_address: 0x00000200, // Overlaps with section1
            size: 512,
            alignment: 4,
            in_flash: false,
        };

        assert!(!section1.overlaps(&section2));
        assert!(section1.overlaps(&section3));
    }

    #[test]
    fn test_memory_layout_creation() {
        let hw_caps = HardwareCapabilities::nrf52840();
        let analyzer = MemoryLayoutAnalyzer::new(hw_caps);
        let component = test_component();

        let layout = analyzer.analyze(&component).unwrap();

        // Should have sections allocated
        assert!(!layout.sections().is_empty());

        // Should have flash usage (text + rodata)
        assert!(layout.flash_usage() > 0);

        // Should have RAM usage (data + bss + stack + heap)
        assert!(layout.ram_usage() > 0);

        // Validate against hardware
        assert!(layout.validate().is_ok());
    }

    #[test]
    fn test_memory_layout_validation() {
        let hw_caps = HardwareCapabilities::nrf52840();
        let analyzer = MemoryLayoutAnalyzer::new(hw_caps);
        let component = test_component();

        let layout = analyzer.analyze(&component).unwrap();

        // Print layout for inspection
        println!("\nMemory Layout:");
        println!(
            "Flash usage: {} / {} bytes",
            layout.flash_usage(),
            layout.flash_usage
        );
        println!(
            "RAM usage: {} / {} bytes",
            layout.ram_usage(),
            layout.ram_usage
        );
        println!("\nSections:");
        for section in layout.sections() {
            println!(
                "  {} ({:?}): 0x{:08X} - 0x{:08X} ({} bytes, {})",
                section.name,
                section.section_type,
                section.base_address,
                section.end_address(),
                section.size,
                if section.in_flash { "flash" } else { "RAM" }
            );
        }

        assert!(layout.validate().is_ok());
    }

    #[test]
    fn test_linker_script_generation() {
        let hw_caps = HardwareCapabilities::nrf52840();
        let analyzer = MemoryLayoutAnalyzer::new(hw_caps);
        let component = test_component();

        let layout = analyzer.analyze(&component).unwrap();
        let linker_script = layout.generate_linker_script();

        // Print linker script for inspection
        println!("\nGenerated Linker Script:");
        println!("{}", linker_script);

        // Verify key elements are present
        assert!(linker_script.contains("MEMORY"));
        assert!(linker_script.contains("FLASH"));
        assert!(linker_script.contains("RAM"));
        assert!(linker_script.contains("ENTRY(Reset_Handler)"));
        assert!(linker_script.contains(".text"));
        assert!(linker_script.contains(".data"));
        assert!(linker_script.contains(".bss"));
        assert!(linker_script.contains(".heap"));
        assert!(linker_script.contains(".stack"));
        assert!(linker_script.contains("_sdata"));
        assert!(linker_script.contains("_edata"));
        assert!(linker_script.contains("_sbss"));
        assert!(linker_script.contains("_ebss"));
        assert!(linker_script.contains("_stack_top"));
    }
}