Skip to main content

jvm_assembler/formats/class/writer/
pool.rs

1use super::{ClassWriter, CpEntry};
2use std::io::Write;
3
4impl<W: Write> ClassWriter<W> {
5    /// Adds a Utf8 constant to the constant pool
6    pub fn add_utf8(&mut self, s: String) -> u16 {
7        self.add_cp_entry(CpEntry::Utf8(s))
8    }
9
10    /// Adds a Class constant to the constant pool
11    pub fn add_class(&mut self, name: String) -> u16 {
12        let name_index = self.add_utf8(name);
13        self.add_cp_entry(CpEntry::Class(name_index))
14    }
15
16    /// Adds a String constant to the constant pool
17    pub fn add_string(&mut self, value: String) -> u16 {
18        let utf8_index = self.add_utf8(value);
19        self.add_cp_entry(CpEntry::String(utf8_index))
20    }
21
22    /// Adds a NameAndType constant to the constant pool
23    pub fn add_name_and_type(&mut self, name: String, descriptor: String) -> u16 {
24        let name_index = self.add_utf8(name);
25        let descriptor_index = self.add_utf8(descriptor);
26        self.add_cp_entry(CpEntry::NameAndType(name_index, descriptor_index))
27    }
28
29    /// Adds a Fieldref constant to the constant pool
30    pub fn add_field_ref(&mut self, class_name: String, name: String, descriptor: String) -> u16 {
31        let class_index = self.add_class(class_name);
32        let name_and_type_index = self.add_name_and_type(name, descriptor);
33        self.add_cp_entry(CpEntry::Fieldref(class_index, name_and_type_index))
34    }
35
36    /// Adds a Methodref constant to the constant pool
37    pub fn add_method_ref(&mut self, class_name: String, name: String, descriptor: String) -> u16 {
38        let class_index = self.add_class(class_name);
39        let name_and_type_index = self.add_name_and_type(name, descriptor);
40        self.add_cp_entry(CpEntry::Methodref(class_index, name_and_type_index))
41    }
42
43    /// Adds an InterfaceMethodref constant to the constant pool
44    pub fn add_interface_method_ref(&mut self, class_name: String, name: String, descriptor: String) -> u16 {
45        let class_index = self.add_class(class_name);
46        let name_and_type_index = self.add_name_and_type(name, descriptor);
47        self.add_cp_entry(CpEntry::InterfaceMethodref(class_index, name_and_type_index))
48    }
49
50    /// Adds an Integer constant to the constant pool
51    pub fn add_int(&mut self, val: i32) -> u16 {
52        self.add_cp_entry(CpEntry::Integer(val))
53    }
54
55    /// Adds a Float constant to the constant pool
56    pub fn add_float(&mut self, val: f32) -> u16 {
57        self.add_cp_entry(CpEntry::Float(val.to_bits()))
58    }
59
60    /// Adds a Long constant to the constant pool
61    pub fn add_long(&mut self, val: i64) -> u16 {
62        self.add_cp_entry(CpEntry::Long(val))
63    }
64
65    /// Adds a Double constant to the constant pool
66    pub fn add_double(&mut self, val: f64) -> u16 {
67        self.add_cp_entry(CpEntry::Double(val.to_bits()))
68    }
69
70    /// Adds a MethodHandle constant to the constant pool
71    pub fn add_method_handle(&mut self, kind: u8, ref_idx: u16) -> u16 {
72        self.add_cp_entry(CpEntry::MethodHandle(kind, ref_idx))
73    }
74
75    /// Adds a MethodType constant to the constant pool
76    pub fn add_method_type(&mut self, descriptor: String) -> u16 {
77        let desc_idx = self.add_utf8(descriptor);
78        self.add_cp_entry(CpEntry::MethodType(desc_idx))
79    }
80
81    /// Adds a Dynamic constant to the constant pool
82    pub fn add_dynamic(&mut self, bootstrap_method_attr_idx: u16, name: String, descriptor: String) -> u16 {
83        let nt_idx = self.add_name_and_type(name, descriptor);
84        self.add_cp_entry(CpEntry::Dynamic(bootstrap_method_attr_idx, nt_idx))
85    }
86
87    /// Adds an InvokeDynamic constant to the constant pool
88    pub fn add_invoke_dynamic(&mut self, bootstrap_method_attr_idx: u16, name: String, descriptor: String) -> u16 {
89        let nt_idx = self.add_name_and_type(name, descriptor);
90        self.add_cp_entry(CpEntry::InvokeDynamic(bootstrap_method_attr_idx, nt_idx))
91    }
92
93    /// Adds a Module constant to the constant pool
94    pub fn add_module(&mut self, name: String) -> u16 {
95        let name_idx = self.add_utf8(name);
96        self.add_cp_entry(CpEntry::Module(name_idx))
97    }
98
99    /// Adds a Package constant to the constant pool
100    pub fn add_package(&mut self, name: String) -> u16 {
101        let name_idx = self.add_utf8(name);
102        self.add_cp_entry(CpEntry::Package(name_idx))
103    }
104
105    /// Adds a constant pool entry and returns its index
106    pub fn add_cp_entry(&mut self, entry: CpEntry) -> u16 {
107        if let Some(&index) = self.cp_map.get(&entry) {
108            return index;
109        }
110
111        let index = (self.cp_entries.len() + 1) as u16;
112        self.cp_entries.push(entry.clone());
113        self.cp_map.insert(entry, index);
114
115        // Long 和 Double 占用两个常量池槽位
116        match self.cp_entries.last().unwrap() {
117            CpEntry::Long(_) | CpEntry::Double(_) => {
118                // 占位符,不实际写入,但增加长度
119                self.cp_entries.push(CpEntry::Utf8("Padding".to_string()));
120            }
121            _ => {}
122        }
123
124        index
125    }
126
127    /// Writes the constant pool to the class file
128    pub fn write_constant_pool(&mut self) -> gaia_types::Result<()> {
129        // 写入常量池计数(+1 因为索引从1开始)
130        self.writer.write_u16((self.cp_entries.len() + 1) as u16)?;
131
132        // 写入常量池条目
133        let mut i = 0;
134        while i < self.cp_entries.len() {
135            let entry = &self.cp_entries[i];
136            match entry {
137                CpEntry::Utf8(s) => {
138                    self.writer.write_u8(1)?; // CONSTANT_Utf8 tag
139                    self.writer.write_u16(s.len() as u16)?;
140                    self.writer.write_all(s.as_bytes())?;
141                }
142                CpEntry::Class(name_idx) => {
143                    self.writer.write_u8(7)?; // CONSTANT_Class tag
144                    self.writer.write_u16(*name_idx)?;
145                }
146                CpEntry::String(utf8_idx) => {
147                    self.writer.write_u8(8)?; // CONSTANT_String tag
148                    self.writer.write_u16(*utf8_idx)?;
149                }
150                CpEntry::NameAndType(name_idx, desc_idx) => {
151                    self.writer.write_u8(12)?; // CONSTANT_NameAndType tag
152                    self.writer.write_u16(*name_idx)?;
153                    self.writer.write_u16(*desc_idx)?;
154                }
155                CpEntry::Fieldref(class_idx, nt_idx) => {
156                    self.writer.write_u8(9)?; // CONSTANT_Fieldref tag
157                    self.writer.write_u16(*class_idx)?;
158                    self.writer.write_u16(*nt_idx)?;
159                }
160                CpEntry::Methodref(class_idx, nt_idx) => {
161                    self.writer.write_u8(10)?; // CONSTANT_Methodref tag
162                    self.writer.write_u16(*class_idx)?;
163                    self.writer.write_u16(*nt_idx)?;
164                }
165                CpEntry::InterfaceMethodref(class_idx, nt_idx) => {
166                    self.writer.write_u8(11)?; // CONSTANT_InterfaceMethodref tag
167                    self.writer.write_u16(*class_idx)?;
168                    self.writer.write_u16(*nt_idx)?;
169                }
170                CpEntry::Integer(val) => {
171                    self.writer.write_u8(3)?; // CONSTANT_Integer tag
172                    self.writer.write_i32(*val)?;
173                }
174                CpEntry::Float(bits) => {
175                    self.writer.write_u8(4)?; // CONSTANT_Float tag
176                    self.writer.write_u32(*bits)?;
177                }
178                CpEntry::Long(val) => {
179                    self.writer.write_u8(5)?; // CONSTANT_Long tag
180                    self.writer.write_i64(*val)?;
181                    i += 1; // Long 占用两个槽位
182                }
183                CpEntry::Double(bits) => {
184                    self.writer.write_u8(6)?; // CONSTANT_Double tag
185                    self.writer.write_u64(*bits)?;
186                    i += 1; // Double 占用两个槽位
187                }
188                CpEntry::MethodHandle(kind, ref_idx) => {
189                    self.writer.write_u8(15)?; // CONSTANT_MethodHandle tag
190                    self.writer.write_u8(*kind)?;
191                    self.writer.write_u16(*ref_idx)?;
192                }
193                CpEntry::MethodType(desc_idx) => {
194                    self.writer.write_u8(16)?; // CONSTANT_MethodType tag
195                    self.writer.write_u16(*desc_idx)?;
196                }
197                CpEntry::Dynamic(bootstrap_method_attr_idx, name_and_type_idx) => {
198                    self.writer.write_u8(17)?; // CONSTANT_Dynamic tag
199                    self.writer.write_u16(*bootstrap_method_attr_idx)?;
200                    self.writer.write_u16(*name_and_type_idx)?;
201                }
202                CpEntry::InvokeDynamic(bootstrap_method_attr_idx, name_and_type_idx) => {
203                    self.writer.write_u8(18)?; // CONSTANT_InvokeDynamic tag
204                    self.writer.write_u16(*bootstrap_method_attr_idx)?;
205                    self.writer.write_u16(*name_and_type_idx)?;
206                }
207                CpEntry::Module(name_idx) => {
208                    self.writer.write_u8(19)?; // CONSTANT_Module tag
209                    self.writer.write_u16(*name_idx)?;
210                }
211                CpEntry::Package(name_idx) => {
212                    self.writer.write_u8(20)?; // CONSTANT_Package tag
213                    self.writer.write_u16(*name_idx)?;
214                }
215            }
216            i += 1;
217        }
218
219        Ok(())
220    }
221}