jvm_assembler/formats/class/writer/

mod.rs

#![doc = include_str!("readme.md")]

mod attributes;
mod cp;
mod entities;
mod instructions;
mod pool;
mod utils;

use crate::program::*;
use cp::CpEntry;
use gaia_binary::{BigEndian, BinaryWriter, Fixed};
use gaia_types::{GaiaDiagnostics, Result};
use std::{collections::HashMap, io::Write};

/// JVM Class file writer
pub struct ClassWriter<W: Write> {
    /// Binary writer
    writer: BinaryWriter<W, Fixed<BigEndian>>,
    /// Constant pool entries
    cp_entries: Vec<CpEntry>,
    /// Constant pool lookup table
    cp_map: HashMap<CpEntry, u16>,
}

impl<W: Write> ClassWriter<W> {
    /// Create a new Class writer
    pub fn new(writer: W) -> Self {
        Self { writer: BinaryWriter::new(writer), cp_entries: Vec::new(), cp_map: HashMap::new() }
    }

    /// Finish writing and return the underlying writer
    pub fn finish(self) -> W {
        self.writer.into_inner()
    }

    /// Write ClassView as binary Class format
    pub fn write(mut self, program: &JvmProgram) -> GaiaDiagnostics<W> {
        match self.write_class_file(program) {
            Ok(_) => GaiaDiagnostics::success(self.finish()),
            Err(error) => GaiaDiagnostics::failure(error),
        }
    }

    /// Write Class file
    fn write_class_file(&mut self, program: &JvmProgram) -> Result<()> {
        // 1. Pre-collect all constants
        let this_class_idx = self.add_class(program.name.clone());
        let super_class_idx = if let Some(super_name) = &program.super_class {
            self.add_class(super_name.clone())
        }
        else {
            self.add_class("java/lang/Object".to_string())
        };

        // Collect class attribute constants
        for attr in &program.attributes {
            self.collect_attribute_constants(attr);
        }

        // Collect inner classes constants
        if !program.inner_classes.is_empty() {
            let inner_classes_attr = JvmAttribute::InnerClasses { classes: program.inner_classes.clone() };
            self.collect_attribute_constants(&inner_classes_attr);
        }

        // Collect enclosing method constants
        if let Some((class_name, method_name, method_descriptor)) = &program.enclosing_method {
            let enclosing_method_attr = JvmAttribute::EnclosingMethod {
                class_name: class_name.clone(),
                method_name: method_name.clone(),
                method_descriptor: method_descriptor.clone(),
            };
            self.collect_attribute_constants(&enclosing_method_attr);
        }

        // Collect field constants
        for field in &program.fields {
            self.add_utf8(field.name.clone());
            self.add_utf8(field.descriptor.clone());
            for attr in &field.attributes {
                self.collect_attribute_constants(attr);
            }
            if field.constant_value.is_some() {
                self.add_utf8("ConstantValue".to_string());
            }
        }

        // Collect method constants
        let code_utf8_idx = self.add_utf8("Code".to_string());
        for method in &program.methods {
            self.add_utf8(method.name.clone());
            self.add_utf8(method.descriptor.clone());

            for handler in &method.exception_handlers {
                if let Some(catch_type) = &handler.catch_type {
                    self.add_class(catch_type.clone());
                }
            }

            for attr in &method.attributes {
                self.collect_attribute_constants(attr);
            }
            if !method.exceptions.is_empty() {
                self.add_utf8("Exceptions".to_string());
            }

            // Pre-generate bytecode to get label positions for StackMapTable generation
            let (_, label_positions) = self.generate_method_bytecode(method);

            if program.version.major >= 50 {
                self.add_utf8("StackMapTable".to_string());
                // If automatic StackMapTable generation is needed, pre-collect its constants
                let has_stack_map = method.attributes.iter().any(|a| matches!(a, JvmAttribute::StackMapTable { .. }));
                if !has_stack_map {
                    let analyzer = crate::analyzer::StackMapAnalyzer::new(program.name.clone(), method, &label_positions);
                    let frames = analyzer.analyze();
                    for frame in &frames {
                        match frame {
                            JvmStackMapFrame::SameLocals1StackItem { stack, .. }
                            | JvmStackMapFrame::SameLocals1StackItemExtended { stack, .. } => {
                                self.collect_verification_type_constants(stack);
                            }
                            JvmStackMapFrame::Append { locals, .. } => {
                                for vt in locals {
                                    self.collect_verification_type_constants(vt);
                                }
                            }
                            JvmStackMapFrame::Full { locals, stack, .. } => {
                                for vt in locals {
                                    self.collect_verification_type_constants(vt);
                                }
                                for vt in stack {
                                    self.collect_verification_type_constants(vt);
                                }
                            }
                            _ => {}
                        }
                    }
                }
            }

            // Collect constants in instructions
            for inst in &method.instructions {
                match inst {
                    JvmInstruction::Ldc { symbol } | JvmInstruction::LdcW { symbol } | JvmInstruction::Ldc2W { symbol } => {
                        self.add_string(symbol.clone());
                    }
                    JvmInstruction::Getstatic { class_name, field_name, descriptor }
                    | JvmInstruction::Putstatic { class_name, field_name, descriptor }
                    | JvmInstruction::Getfield { class_name, field_name, descriptor }
                    | JvmInstruction::Putfield { class_name, field_name, descriptor } => {
                        self.add_field_ref(class_name.clone(), field_name.clone(), descriptor.clone());
                    }
                    JvmInstruction::Invokevirtual { class_name, method_name, descriptor }
                    | JvmInstruction::Invokespecial { class_name, method_name, descriptor }
                    | JvmInstruction::Invokestatic { class_name, method_name, descriptor }
                    | JvmInstruction::Invokedynamic { class_name, method_name, descriptor } => {
                        self.add_method_ref(class_name.clone(), method_name.clone(), descriptor.clone());
                    }
                    JvmInstruction::Invokeinterface { class_name, method_name, descriptor } => {
                        self.add_interface_method_ref(class_name.clone(), method_name.clone(), descriptor.clone());
                    }
                    JvmInstruction::New { class_name }
                    | JvmInstruction::Anewarray { class_name }
                    | JvmInstruction::Checkcast { class_name }
                    | JvmInstruction::Instanceof { class_name }
                    | JvmInstruction::Multianewarray { class_name, .. } => {
                        self.add_class(class_name.clone());
                    }
                    _ => {}
                }
            }
        }

        // 2. Start writing
        // Write magic number
        self.writer.write_u32(0xCAFEBABE)?;

        // Write version information
        self.writer.write_u16(program.version.minor)?;
        self.writer.write_u16(program.version.major)?;

        // Write constant pool
        self.write_constant_pool()?;

        // Write access flags
        self.writer.write_u16(program.access_flags.to_flags())?;

        // Write this_class index
        self.writer.write_u16(this_class_idx)?;

        // Write super_class index
        self.writer.write_u16(super_class_idx)?;

        // Write interface count (currently 0)
        self.writer.write_u16(0)?;

        // Write fields
        self.write_fields(program)?;

        // Write methods
        self.write_methods(program, code_utf8_idx)?;

        // Write class attributes
        let mut attributes = program.attributes.clone();
        if !program.inner_classes.is_empty() {
            attributes.push(JvmAttribute::InnerClasses { classes: program.inner_classes.clone() });
        }
        if let Some((class_name, method_name, method_descriptor)) = &program.enclosing_method {
            attributes.push(JvmAttribute::EnclosingMethod {
                class_name: class_name.clone(),
                method_name: method_name.clone(),
                method_descriptor: method_descriptor.clone(),
            });
        }
        self.writer.write_u16(attributes.len() as u16)?;
        for attr in &attributes {
            self.write_attribute(attr)?;
        }

        Ok(())
    }
}