jvm_assembler/formats/class/writer/
attributes.rs

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

jvm_assembler/formats/class/writer/
attributes.rs
