gaia-assembler 0.1.1

//! WASI (WebAssembly System Interface) backend compiler

use super::{Backend, GeneratedFiles};
use crate::{
    config::GaiaConfig,
    instruction::{CmpCondition, CoreInstruction, GaiaInstruction, ManagedInstruction},
    program::{GaiaConstant, GaiaFunction, GaiaModule},
    types::GaiaType,
};
use gaia_types::{
    helpers::{AbiCompatible, ApiCompatible, Architecture, ArtifactType, CompilationTarget},
    Result,
};
use std::collections::HashMap;

#[cfg(feature = "wasi-assembler")]
use wasi_assembler::*;

#[cfg(not(feature = "wasi-assembler"))]
mod wasi_stub {
    pub struct WasiProgram;
    impl WasiProgram {
        pub fn to_wasm(&self) -> gaia_types::Result<Vec<u8>> {
            Err(gaia_types::errors::GaiaError::custom_error("WASI backend not enabled"))
        }
    }
    #[allow(dead_code)]
    pub enum WasmValueType {
        I32,
        I64,
        F32,
        F64,
    }
}
#[cfg(not(feature = "wasi-assembler"))]
use wasi_stub::*;

/// WASI Backend implementation
#[derive(Default)]
pub struct WasiBackend {}

impl Backend for WasiBackend {
    fn name(&self) -> &'static str {
        "WASI"
    }

    fn primary_target(&self) -> CompilationTarget {
        CompilationTarget { build: Architecture::WASM32, host: AbiCompatible::WebAssemblyBinary, target: ApiCompatible::WASI }
    }

    fn artifact_type(&self) -> ArtifactType {
        ArtifactType::Executable
    }

    fn match_score(&self, target: &CompilationTarget) -> f32 {
        match target.host {
            AbiCompatible::WebAssemblyBinary => 10.0,
            AbiCompatible::WebAssemblyTextFormat => 5.0,
            AbiCompatible::Unknown => match target.build {
                // wat output, 5% support
                Architecture::WASM32 => 5.0,
                Architecture::WASM64 => 0.0,
                _ => -100.0,
            },
            _ => -100.0,
        }
    }

    fn generate(&self, program: &GaiaModule, config: &GaiaConfig) -> Result<GeneratedFiles> {
        #[cfg(feature = "wasi-assembler")]
        {
            let mut context = create_wasi_context()?;
            compile_program(&mut context, program)?;

            let wasi_program = context.program;

            let mut files = HashMap::new();

            // 1. Generate .wasm
            files.insert("main.wasm".to_string(), wasi_program.to_wasm()?);

            Ok(GeneratedFiles { artifact_type: self.artifact_type(), files, custom: None, diagnostics: vec![] })
        }
        #[cfg(not(feature = "wasi-assembler"))]
        {
            let _ = program;
            Err(gaia_types::errors::GaiaError::custom_error("WASI backend not enabled"))
        }
    }
}

impl WasiBackend {
    /// Generate WASI WebAssembly bytecode from Gaia program
    pub fn generate(program: &GaiaModule) -> Result<Vec<u8>> {
        #[cfg(feature = "wasi-assembler")]
        {
            let mut context = create_wasi_context()?;
            compile_program(&mut context, program)?;
            context.program.to_wasm()
        }
        #[cfg(not(feature = "wasi-assembler"))]
        {
            let _ = program;
            Err(gaia_types::errors::GaiaError::custom_error("WASI backend not enabled"))
        }
    }
}

/// Compile Gaia program to WASI WebAssembly
pub fn compile(program: &GaiaModule) -> Result<Vec<u8>> {
    WasiBackend::generate(program)
}

#[cfg(feature = "wasi-assembler")]
struct WasiContext {
    program: WasiProgram,
    /// Import map
    import_map: HashMap<String, u32>,
}

#[cfg(feature = "wasi-assembler")]
impl WasiContext {
    fn new() -> Self {
        WasiContext { program: WasiProgram::new_core_module(), import_map: HashMap::new() }
    }

    fn add_import(&mut self, module: &str, field: &str, params: Vec<WasmValueType>, results: Vec<WasmValueType>) -> u32 {
        let key = format!("{}.{}", module, field);
        if let Some(&index) = self.import_map.get(&key) {
            index
        }
        else {
            let index = self.import_map.len() as u32;
            let func_type = WasiFunctionType { params, results };
            let type_index = self.program.add_function_type(func_type);

            self.program.add_import(wasi_assembler::WasiImport {
                module: module.to_string(),
                field: field.to_string(),
                import_type: wasi_assembler::WasmImportType::Function { type_index },
            });

            self.import_map.insert(key, index);
            index
        }
    }

    fn add_function(&mut self, body: Vec<WasiInstruction>, params: Vec<WasmValueType>, results: Vec<WasmValueType>) {
        let func_type = WasiFunctionType { params, results };
        let type_index = self.program.add_function_type(func_type);
        let func = wasi_assembler::WasiFunction { type_index, locals: vec![], body };
        self.program.add_function(func);
    }
}

#[cfg(feature = "wasi-assembler")]
fn create_wasi_context() -> Result<WasiContext> {
    Ok(WasiContext::new())
}

#[cfg(feature = "wasi-assembler")]
fn compile_program(context: &mut WasiContext, program: &GaiaModule) -> Result<()> {
    // 2. Compile all functions
    for function in &program.functions {
        compile_function(context, function)?;
    }

    // 3. Export main function
    if !context.program.functions.is_empty() {
        let main_index = (context.program.imports.len() + context.program.functions.len() - 1) as u32;
        context.program.add_export(wasi_assembler::WasiExport {
            name: "main".to_string(),
            export_type: wasi_assembler::WasmExportType::Function { function_index: main_index },
        });
    }

    Ok(())
}

#[cfg(feature = "wasi-assembler")]
fn compile_function(context: &mut WasiContext, function: &GaiaFunction) -> Result<()> {
    let mut instrs = Vec::new();

    // Compile code blocks
    for block in &function.blocks {
        for instruction in &block.instructions {
            compile_instruction(context, &mut instrs, instruction)?;
        }

        // Compile terminators
        match &block.terminator {
            crate::program::GaiaTerminator::Jump(_label) => instrs.push(WasiInstruction::Br { label_index: 0 }), // FIXME
            crate::program::GaiaTerminator::Branch { .. } => {
                instrs.push(WasiInstruction::BrIf { label_index: 0 }); // FIXME
            }
            crate::program::GaiaTerminator::Return => instrs.push(WasiInstruction::Return),
            crate::program::GaiaTerminator::Call { .. } => {
                // TODO: Find function index
                instrs.push(WasiInstruction::Call { function_index: 0 });
            }
            crate::program::GaiaTerminator::Halt => {
                // WASI: exit(0)
                instrs.push(WasiInstruction::I32Const { value: 0 });
                // TODO: Find proc_exit index
                instrs.push(WasiInstruction::Call { function_index: 0 });
            }
        }
    }

    instrs.push(WasiInstruction::End);

    // Map parameters and return types
    let params = function.signature.params.iter().map(|p| map_type(p)).collect();
    let results = match function.signature.return_type {
        GaiaType::Void => vec![],
        _ => vec![map_type(&function.signature.return_type)],
    };

    // Add function to program
    context.add_function(instrs, params, results);

    Ok(())
}

#[cfg(feature = "wasi-assembler")]
fn map_type(ty: &GaiaType) -> WasmValueType {
    match ty {
        GaiaType::I32 => WasmValueType::I32,
        GaiaType::I64 => WasmValueType::I64,
        GaiaType::F32 => WasmValueType::F32,
        GaiaType::F64 => WasmValueType::F64,
        _ => WasmValueType::I32, // Fallback
    }
}

#[cfg(feature = "wasi-assembler")]
fn compile_instruction(
    _context: &mut WasiContext,
    instrs: &mut Vec<WasiInstruction>,
    instruction: &GaiaInstruction,
) -> Result<()> {
    match instruction {
        GaiaInstruction::Core(core) => match core {
            CoreInstruction::PushConstant(constant) => compile_load_constant(instrs, constant),
            CoreInstruction::Load(gaia_type) => compile_load_indirect(instrs, gaia_type),
            CoreInstruction::Store(gaia_type) => compile_store_indirect(instrs, gaia_type),
            CoreInstruction::Add(_) => {
                instrs.push(WasiInstruction::I32Add);
                Ok(())
            }
            CoreInstruction::Sub(_) => {
                instrs.push(WasiInstruction::I32Sub);
                Ok(())
            }
            CoreInstruction::Mul(_) => {
                instrs.push(WasiInstruction::I32Mul);
                Ok(())
            }
            CoreInstruction::Div(_) => {
                instrs.push(WasiInstruction::I32DivS);
                Ok(())
            }
            CoreInstruction::Rem(_) => {
                instrs.push(WasiInstruction::I32RemS);
                Ok(())
            }
            CoreInstruction::And(_) => {
                instrs.push(WasiInstruction::I32And);
                Ok(())
            }
            CoreInstruction::Or(_) => {
                instrs.push(WasiInstruction::I32Or);
                Ok(())
            }
            CoreInstruction::Xor(_) => {
                instrs.push(WasiInstruction::I32Xor);
                Ok(())
            }
            CoreInstruction::Shl(_) => {
                instrs.push(WasiInstruction::I32Shl);
                Ok(())
            }
            CoreInstruction::Shr(_) => {
                instrs.push(WasiInstruction::I32ShrS);
                Ok(())
            }
            CoreInstruction::Pop => {
                instrs.push(WasiInstruction::Drop);
                Ok(())
            }
            CoreInstruction::LoadLocal(index, _) => {
                instrs.push(WasiInstruction::LocalGet { local_index: *index });
                Ok(())
            }
            CoreInstruction::StoreLocal(index, _) => {
                instrs.push(WasiInstruction::LocalSet { local_index: *index });
                Ok(())
            }
            CoreInstruction::LoadArg(index, _) => {
                instrs.push(WasiInstruction::LocalGet { local_index: *index });
                Ok(())
            }
            CoreInstruction::StoreArg(index, _) => {
                instrs.push(WasiInstruction::LocalSet { local_index: *index });
                Ok(())
            }
            CoreInstruction::Ret => {
                instrs.push(WasiInstruction::Return);
                Ok(())
            }
            CoreInstruction::Call(_, _) => {
                // TODO: Find function index
                instrs.push(WasiInstruction::Call { function_index: 0 });
                Ok(())
            }
            CoreInstruction::Cmp(cond, _) => compile_compare(instrs, cond),
            CoreInstruction::StructNew(_) => {
                // TODO: Find type index
                instrs.push(WasiInstruction::StructNew { type_index: 0 });
                Ok(())
            }
            CoreInstruction::StructGet { field_index, .. } => {
                // TODO: Find type index
                instrs.push(WasiInstruction::StructGet { type_index: 0, field_index: *field_index as u32 });
                Ok(())
            }
            CoreInstruction::StructSet { field_index, .. } => {
                // TODO: Find type index
                instrs.push(WasiInstruction::StructSet { type_index: 0, field_index: *field_index as u32 });
                Ok(())
            }
            _ => Ok(()),
        },
        GaiaInstruction::Managed(managed) => match managed {
            ManagedInstruction::CallStatic { .. } => {
                // AOT mode maps static method calls to regular calls
                instrs.push(WasiInstruction::Call { function_index: 0 });
                Ok(())
            }
            _ => Ok(()),
        },
        _ => Ok(()),
    }
}

#[cfg(feature = "wasi-assembler")]
fn compile_compare(instrs: &mut Vec<WasiInstruction>, cond: &CmpCondition) -> Result<()> {
    match cond {
        CmpCondition::Eq => instrs.push(WasiInstruction::I32Eq),
        CmpCondition::Ne => instrs.push(WasiInstruction::I32Ne),
        CmpCondition::Lt => instrs.push(WasiInstruction::I32LtS),
        CmpCondition::Le => instrs.push(WasiInstruction::I32LeS),
        CmpCondition::Gt => instrs.push(WasiInstruction::I32GtS),
        CmpCondition::Ge => instrs.push(WasiInstruction::I32GeS),
    }
    Ok(())
}

#[cfg(feature = "wasi-assembler")]
fn compile_load_indirect(instrs: &mut Vec<WasiInstruction>, gaia_type: &GaiaType) -> Result<()> {
    // Assume memory offset is already on top of the stack
    match gaia_type {
        GaiaType::I32 => instrs.push(WasiInstruction::I32Load { offset: 0, align: 0 }),
        GaiaType::I64 => instrs.push(WasiInstruction::I64Load { offset: 0, align: 0 }),
        GaiaType::F32 => instrs.push(WasiInstruction::F32Load { offset: 0, align: 0 }),
        GaiaType::F64 => instrs.push(WasiInstruction::F64Load { offset: 0, align: 0 }),
        _ => {}
    }
    Ok(())
}

#[cfg(feature = "wasi-assembler")]
fn compile_store_indirect(instrs: &mut Vec<WasiInstruction>, gaia_type: &GaiaType) -> Result<()> {
    // Assume [ptr, value] are already on the stack
    match gaia_type {
        GaiaType::I32 => instrs.push(WasiInstruction::I32Store { offset: 0, align: 0 }),
        GaiaType::I64 => instrs.push(WasiInstruction::I64Store { offset: 0, align: 0 }),
        GaiaType::F32 => instrs.push(WasiInstruction::F32Store { offset: 0, align: 0 }),
        GaiaType::F64 => instrs.push(WasiInstruction::F64Store { offset: 0, align: 0 }),
        _ => {}
    }
    Ok(())
}

#[cfg(feature = "wasi-assembler")]
fn compile_load_constant(instrs: &mut Vec<WasiInstruction>, constant: &GaiaConstant) -> Result<()> {
    match constant {
        GaiaConstant::I32(value) => instrs.push(WasiInstruction::I32Const { value: *value }),
        GaiaConstant::I64(value) => instrs.push(WasiInstruction::I64Const { value: *value }),
        GaiaConstant::F32(value) => instrs.push(WasiInstruction::F32Const { value: *value }),
        GaiaConstant::F64(value) => instrs.push(WasiInstruction::F64Const { value: *value }),
        _ => {}
    }
    Ok(())
}