jvm_assembler/formats/class/writer/
attributes.rs

1use super::ClassWriter;
2use crate::program::*;
3use gaia_types::{BinaryWriter, Result};
4use std::io::Write;
5
6/// Raw exception handler entry
7struct RawExceptionHandler {
8    start_pc: u16,
9    end_pc: u16,
10    handler_pc: u16,
11    catch_type_index: u16,
12}
13
14impl<W: Write> ClassWriter<W> {
15    /// Collects constants from an attribute and adds them to the constant pool
16    pub fn collect_attribute_constants(&mut self, attribute: &JvmAttribute) {
17        match attribute {
18            JvmAttribute::SourceFile { filename } => {
19                self.add_utf8("SourceFile".to_string());
20                self.add_utf8(filename.clone());
21            }
22            JvmAttribute::ConstantValue { value } => {
23                self.add_utf8("ConstantValue".to_string());
24                match value {
25                    JvmConstantPoolEntry::Integer { .. } => {}
26                    JvmConstantPoolEntry::Float { .. } => {}
27                    JvmConstantPoolEntry::Long { .. } => {}
28                    JvmConstantPoolEntry::Double { .. } => {}
29                    JvmConstantPoolEntry::String { value } => {
30                        self.add_string(value.clone());
31                    }
32                    _ => {}
33                }
34            }
35            JvmAttribute::Exceptions { exceptions } => {
36                self.add_utf8("Exceptions".to_string());
37                for exc in exceptions {
38                    self.add_class(exc.clone());
39                }
40            }
41            JvmAttribute::Signature { signature } => {
42                self.add_utf8("Signature".to_string());
43                self.add_utf8(signature.clone());
44            }
45            JvmAttribute::StackMapTable { frames } => {
46                self.add_utf8("StackMapTable".to_string());
47                for frame in frames {
48                    match frame {
49                        JvmStackMapFrame::SameLocals1StackItem { stack, .. }
50                        | JvmStackMapFrame::SameLocals1StackItemExtended { stack, .. } => {
51                            self.collect_verification_type_constants(stack);
52                        }
53                        JvmStackMapFrame::Append { locals, .. } => {
54                            for vt in locals {
55                                self.collect_verification_type_constants(vt);
56                            }
57                        }
58                        JvmStackMapFrame::Full { locals, stack, .. } => {
59                            for vt in locals {
60                                self.collect_verification_type_constants(vt);
61                            }
62                            for vt in stack {
63                                self.collect_verification_type_constants(vt);
64                            }
65                        }
66                        _ => {}
67                    }
68                }
69            }
70            JvmAttribute::InnerClasses { classes } => {
71                self.add_utf8("InnerClasses".to_string());
72                for inner in classes {
73                    self.add_class(inner.inner_class.clone());
74                    if let Some(outer) = &inner.outer_class {
75                        self.add_class(outer.clone());
76                    }
77                    if let Some(name) = &inner.inner_name {
78                        self.add_utf8(name.clone());
79                    }
80                }
81            }
82            JvmAttribute::EnclosingMethod { class_name, method_name, method_descriptor } => {
83                self.add_utf8("EnclosingMethod".to_string());
84                self.add_class(class_name.clone());
85                if let (Some(name), Some(desc)) = (method_name, method_descriptor) {
86                    self.add_name_and_type(name.clone(), desc.clone());
87                }
88            }
89            JvmAttribute::Code { attributes, exception_table, .. } => {
90                self.add_utf8("Code".to_string());
91                for attr in attributes {
92                    self.collect_attribute_constants(attr);
93                }
94                for handler in exception_table {
95                    if let Some(catch_type) = &handler.catch_type {
96                        self.add_class(catch_type.clone());
97                    }
98                }
99            }
100            JvmAttribute::LineNumberTable { .. } => {
101                self.add_utf8("LineNumberTable".to_string());
102            }
103            JvmAttribute::LocalVariableTable { entries } => {
104                self.add_utf8("LocalVariableTable".to_string());
105                for entry in entries {
106                    self.add_utf8(entry.name.clone());
107                    self.add_utf8(entry.descriptor.clone());
108                }
109            }
110            JvmAttribute::Unknown { name, .. } => {
111                self.add_utf8(name.clone());
112            }
113        }
114    }
115
116    /// Collects constants from a verification type and adds them to the constant pool
117    pub fn collect_verification_type_constants(&mut self, vt: &JvmVerificationType) {
118        if let JvmVerificationType::Object { class_name } = vt {
119            self.add_class(class_name.clone());
120        }
121    }
122
123    /// Writes an attribute to the class file
124    pub fn write_attribute(&mut self, attribute: &JvmAttribute) -> Result<()> {
125        match attribute {
126            JvmAttribute::SourceFile { filename } => {
127                let name_idx = self.add_utf8("SourceFile".to_string());
128                let file_idx = self.add_utf8(filename.clone());
129                self.writer.write_u16(name_idx)?;
130                self.writer.write_u32(2)?;
131                self.writer.write_u16(file_idx)?;
132            }
133            JvmAttribute::ConstantValue { value } => {
134                let name_idx = self.add_utf8("ConstantValue".to_string());
135                let val_idx = match value {
136                    JvmConstantPoolEntry::Integer { value } => self.add_int(*value),
137                    JvmConstantPoolEntry::Float { value } => self.add_float(*value),
138                    JvmConstantPoolEntry::Long { value } => self.add_long(*value),
139                    JvmConstantPoolEntry::Double { value } => self.add_double(*value),
140                    JvmConstantPoolEntry::String { value } => self.add_string(value.clone()),
141                    _ => 0,
142                };
143                self.writer.write_u16(name_idx)?;
144                self.writer.write_u32(2)?;
145                self.writer.write_u16(val_idx)?;
146            }
147            JvmAttribute::Exceptions { exceptions } => {
148                let name_idx = self.add_utf8("Exceptions".to_string());
149                let attr_len = 2 + exceptions.len() * 2;
150                self.writer.write_u16(name_idx)?;
151                self.writer.write_u32(attr_len as u32)?;
152                self.writer.write_u16(exceptions.len() as u16)?;
153                for exc in exceptions {
154                    let exc_idx = self.add_class(exc.clone());
155                    self.writer.write_u16(exc_idx)?;
156                }
157            }
158            JvmAttribute::Signature { signature } => {
159                let name_idx = self.add_utf8("Signature".to_string());
160                let sig_idx = self.add_utf8(signature.clone());
161                self.writer.write_u16(name_idx)?;
162                self.writer.write_u32(2)?;
163                self.writer.write_u16(sig_idx)?;
164            }
165            JvmAttribute::StackMapTable { frames } => {
166                let name_idx = self.add_utf8("StackMapTable".to_string());
167                let mut buf = Vec::new();
168                buf.extend_from_slice(&(frames.len() as u16).to_be_bytes());
169                for frame in frames {
170                    self.write_stack_map_frame_to_buf(frame, &mut buf)?;
171                }
172                self.writer.write_u16(name_idx)?;
173                self.writer.write_u32(buf.len() as u32)?;
174                self.writer.write_all(&buf)?;
175            }
176            JvmAttribute::InnerClasses { classes } => {
177                let name_idx = self.add_utf8("InnerClasses".to_string());
178                let attr_len = 2 + classes.len() * 8;
179                self.writer.write_u16(name_idx)?;
180                self.writer.write_u32(attr_len as u32)?;
181                self.writer.write_u16(classes.len() as u16)?;
182                for inner in classes {
183                    let inner_idx = self.add_class(inner.inner_class.clone());
184                    let outer_idx = if let Some(outer) = &inner.outer_class { self.add_class(outer.clone()) } else { 0 };
185                    let name_idx = if let Some(name) = &inner.inner_name { self.add_utf8(name.clone()) } else { 0 };
186                    self.writer.write_u16(inner_idx)?;
187                    self.writer.write_u16(outer_idx)?;
188                    self.writer.write_u16(name_idx)?;
189                    self.writer.write_u16(inner.access_flags.to_flags())?;
190                }
191            }
192            JvmAttribute::EnclosingMethod { class_name, method_name, method_descriptor } => {
193                let name_idx = self.add_utf8("EnclosingMethod".to_string());
194                let class_idx = self.add_class(class_name.clone());
195                let nt_idx = if let (Some(name), Some(desc)) = (method_name, method_descriptor) {
196                    self.add_name_and_type(name.clone(), desc.clone())
197                }
198                else {
199                    0
200                };
201                self.writer.write_u16(name_idx)?;
202                self.writer.write_u32(4)?;
203                self.writer.write_u16(class_idx)?;
204                self.writer.write_u16(nt_idx)?;
205            }
206            JvmAttribute::Unknown { name, data } => {
207                let name_idx = self.add_utf8(name.clone());
208                self.writer.write_u16(name_idx)?;
209                self.writer.write_u32(data.len() as u32)?;
210                self.writer.write_all(data)?;
211            }
212            _ => {
213                // Code, LineNumberTable, LocalVariableTable 等需要特殊处理或暂不支持
214            }
215        }
216        Ok(())
217    }
218
219    /// Writes a verification type to a buffer
220    pub fn write_verification_type_to_buf(&mut self, vt: &JvmVerificationType, buf: &mut Vec<u8>) -> Result<()> {
221        match vt {
222            JvmVerificationType::Top => buf.push(0),
223            JvmVerificationType::Integer => buf.push(1),
224            JvmVerificationType::Float => buf.push(2),
225            JvmVerificationType::Double => buf.push(3),
226            JvmVerificationType::Long => buf.push(4),
227            JvmVerificationType::Null => buf.push(5),
228            JvmVerificationType::UninitializedThis => buf.push(6),
229            JvmVerificationType::Object { class_name } => {
230                buf.push(7);
231                let class_idx = self.add_class(class_name.clone());
232                buf.extend_from_slice(&class_idx.to_be_bytes());
233            }
234            JvmVerificationType::Uninitialized { offset } => {
235                buf.push(8);
236                buf.extend_from_slice(&offset.to_be_bytes());
237            }
238        }
239        Ok(())
240    }
241
242    /// Writes a stack map frame to a buffer
243    pub fn write_stack_map_frame_to_buf(&mut self, frame: &JvmStackMapFrame, buf: &mut Vec<u8>) -> Result<()> {
244        match frame {
245            JvmStackMapFrame::Same { offset_delta } => {
246                if *offset_delta <= 63 {
247                    buf.push(*offset_delta as u8);
248                }
249                else {
250                    buf.push(251);
251                    buf.extend_from_slice(&offset_delta.to_be_bytes());
252                }
253            }
254            JvmStackMapFrame::SameLocals1StackItem { offset_delta, stack } => {
255                if *offset_delta <= 63 {
256                    buf.push((*offset_delta + 64) as u8);
257                    self.write_verification_type_to_buf(stack, buf)?;
258                }
259                else {
260                    buf.push(247);
261                    buf.extend_from_slice(&offset_delta.to_be_bytes());
262                    self.write_verification_type_to_buf(stack, buf)?;
263                }
264            }
265            JvmStackMapFrame::SameLocals1StackItemExtended { offset_delta, stack } => {
266                buf.push(247);
267                buf.extend_from_slice(&offset_delta.to_be_bytes());
268                self.write_verification_type_to_buf(stack, buf)?;
269            }
270            JvmStackMapFrame::Chop { offset_delta, k } => {
271                buf.push(251 - k);
272                buf.extend_from_slice(&offset_delta.to_be_bytes());
273            }
274            JvmStackMapFrame::SameExtended { offset_delta } => {
275                buf.push(251);
276                buf.extend_from_slice(&offset_delta.to_be_bytes());
277            }
278            JvmStackMapFrame::Append { offset_delta, locals } => {
279                buf.push((251 + locals.len()) as u8);
280                buf.extend_from_slice(&offset_delta.to_be_bytes());
281                for vt in locals {
282                    self.write_verification_type_to_buf(vt, buf)?;
283                }
284            }
285            JvmStackMapFrame::Full { offset_delta, locals, stack } => {
286                buf.push(255);
287                buf.extend_from_slice(&offset_delta.to_be_bytes());
288                buf.extend_from_slice(&(locals.len() as u16).to_be_bytes());
289                for vt in locals {
290                    self.write_verification_type_to_buf(vt, buf)?;
291                }
292                buf.extend_from_slice(&(stack.len() as u16).to_be_bytes());
293                for vt in stack {
294                    self.write_verification_type_to_buf(vt, buf)?;
295                }
296            }
297        }
298        Ok(())
299    }
300
301    /// Writes a Code attribute to the class file
302    pub fn write_code_attribute(&mut self, program: &JvmProgram, method: &JvmMethod, code_utf8_idx: u16) -> Result<()> {
303        // Code 属性名称索引
304        self.writer.write_u16(code_utf8_idx)?;
305
306        let (bytecode, label_positions) = self.generate_method_bytecode(method);
307
308        // 解析异常表
309        let mut raw_exception_table = Vec::new();
310        for handler in &method.exception_handlers {
311            let start_pc = *label_positions.get(&handler.start_label).unwrap_or(&0) as u16;
312            let end_pc = *label_positions.get(&handler.end_label).unwrap_or(&0) as u16;
313            let handler_pc = *label_positions.get(&handler.handler_label).unwrap_or(&0) as u16;
314            let catch_type_index =
315                if let Some(catch_type) = &handler.catch_type { self.add_class(catch_type.clone()) } else { 0 };
316            raw_exception_table.push(RawExceptionHandler { start_pc, end_pc, handler_pc, catch_type_index });
317        }
318
319        // 收集并准备子属性
320        let mut code_attributes = Vec::new();
321
322        // 自动生成 StackMapTable (如果版本 >= Java 6 且尚未手动提供)
323        if program.version.major >= 50 {
324            let has_stack_map = method.attributes.iter().any(|a| matches!(a, JvmAttribute::StackMapTable { .. }));
325            if !has_stack_map {
326                let analyzer = crate::analyzer::StackMapAnalyzer::new(program.name.clone(), method, &label_positions);
327                let frames = analyzer.analyze();
328                if !frames.is_empty() {
329                    code_attributes.push(JvmAttribute::StackMapTable { frames });
330                }
331            }
332        }
333
334        // 准备子属性数据
335        let mut sub_attr_buf = Vec::new();
336        // 这里需要临时替换 writer 以写入到缓冲区
337        let mut temp_writer = ClassWriter {
338            writer: BinaryWriter::new(&mut sub_attr_buf),
339            cp_entries: self.cp_entries.clone(),
340            cp_map: self.cp_map.clone(),
341        };
342
343        // 写入自动生成的属性
344        for attr in &code_attributes {
345            temp_writer.write_attribute(attr)?;
346        }
347
348        // 同步常量池
349        self.cp_entries = temp_writer.cp_entries;
350        self.cp_map = temp_writer.cp_map;
351
352        // 计算总长度
353        let exception_table_len = raw_exception_table.len() * 8;
354        let attribute_length = 2 + 2 + 4 + bytecode.len() + 2 + exception_table_len + 2 + sub_attr_buf.len();
355        self.writer.write_u32(attribute_length as u32)?;
356
357        // max_stack 和 max_locals
358        self.writer.write_u16(method.max_stack)?;
359        self.writer.write_u16(method.max_locals)?;
360
361        // 字节码长度和字节码
362        self.writer.write_u32(bytecode.len() as u32)?;
363        self.writer.write_all(&bytecode)?;
364
365        // 异常表
366        self.writer.write_u16(raw_exception_table.len() as u16)?;
367        for handler in raw_exception_table {
368            self.writer.write_u16(handler.start_pc)?;
369            self.writer.write_u16(handler.end_pc)?;
370            self.writer.write_u16(handler.handler_pc)?;
371            self.writer.write_u16(handler.catch_type_index)?;
372        }
373
374        // 子属性数量
375        self.writer.write_u16(code_attributes.len() as u16)?;
376        self.writer.write_all(&sub_attr_buf)?;
377
378        Ok(())
379    }
380}
jvm_assembler/formats/class/writer/attributes.rs

jvm_assembler/formats/class/writer/
attributes.rs
