Skip to main content

x86_64_assembler/builder/
mod.rs

1#![doc = include_str!("readme.md")]
2
3use crate::instruction::{Instruction, Operand, Register};
4use gaia_types::{helpers::Architecture, Result};
5
6/// 通用程序构建器,用于生成汇编指令序列
7#[derive(Debug, Clone)]
8pub struct ProgramBuilder {
9    architecture: Architecture,
10    instructions: Vec<Instruction>,
11    data_sections: Vec<DataSection>,
12}
13
14/// 数据段定义
15#[derive(Debug, Clone)]
16pub struct DataSection {
17    /// 数据段名称
18    pub name: String,
19    /// 数据内容
20    pub data: Vec<u8>,
21    /// 对齐要求(可选)
22    pub alignment: Option<u32>,
23}
24
25impl ProgramBuilder {
26    /// 创建新的程序构建器
27    pub fn new(architecture: Architecture) -> Self {
28        Self { architecture, instructions: Vec::new(), data_sections: Vec::new() }
29    }
30
31    /// 获取当前架构
32    pub fn architecture(&self) -> &Architecture {
33        &self.architecture
34    }
35
36    /// 添加指令
37    pub fn add_instruction(&mut self, instruction: Instruction) -> &mut Self {
38        self.instructions.push(instruction);
39        self
40    }
41
42    /// 批量添加指令
43    pub fn add_instructions(&mut self, instructions: Vec<Instruction>) -> &mut Self {
44        self.instructions.extend(instructions);
45        self
46    }
47
48    /// 添加数据段
49    pub fn add_data_section(&mut self, name: String, data: Vec<u8>) -> &mut Self {
50        self.data_sections.push(DataSection { name, data, alignment: None });
51        self
52    }
53
54    /// 添加带对齐的数据段
55    pub fn add_data_section_with_alignment(&mut self, name: String, data: Vec<u8>, alignment: u32) -> &mut Self {
56        self.data_sections.push(DataSection { name, data, alignment: Some(alignment) });
57        self
58    }
59
60    /// 获取所有指令
61    pub fn instructions(&self) -> &[Instruction] {
62        &self.instructions
63    }
64
65    /// 获取所有数据段
66    pub fn data_sections(&self) -> &[DataSection] {
67        &self.data_sections
68    }
69
70    /// 清空所有指令
71    pub fn clear_instructions(&mut self) -> &mut Self {
72        self.instructions.clear();
73        self
74    }
75
76    /// 清空所有数据段
77    pub fn clear_data_sections(&mut self) -> &mut Self {
78        self.data_sections.clear();
79        self
80    }
81
82    /// 编译指令为字节码
83    pub fn compile_instructions(&self) -> Result<Vec<u8>> {
84        let assembler = crate::X86_64Assembler::new(self.architecture.clone())?;
85        let mut code = Vec::new();
86        let mut labels = std::collections::HashMap::new();
87
88        // 第一遍:记录标签位置
89        let mut current_pos = 0;
90        for instruction in &self.instructions {
91            match instruction {
92                Instruction::Label(name) => {
93                    labels.insert(name.clone(), current_pos);
94                }
95                _ => {
96                    let bytes = assembler.encode(instruction)?;
97                    current_pos += bytes.len();
98                }
99            }
100        }
101
102        // 第二遍:生成最终代码并解析标签偏移
103        current_pos = 0;
104        for instruction in &self.instructions {
105            match instruction {
106                Instruction::Label(_) => continue,
107                Instruction::Jmp { target }
108                | Instruction::Jz { target }
109                | Instruction::Jnz { target }
110                | Instruction::Call { target } => {
111                    if let Operand::Label(name) = target {
112                        if let Some(&label_pos) = labels.get(name) {
113                            let mut bytes = assembler.encode(instruction)?;
114                            // 简单的短跳转解析 (假设是 2 字节指令,第 2 字节是偏移)
115                            if bytes.len() == 2 {
116                                let offset = (label_pos as i32) - (current_pos as i32 + 2);
117                                bytes[1] = offset as u8;
118                            }
119                            // 处理近调用 (E8 disp32)
120                            else if bytes.len() == 5 && bytes[0] == 0xE8 {
121                                let offset = (label_pos as i32) - (current_pos as i32 + 5);
122                                bytes[1..5].copy_from_slice(&offset.to_le_bytes());
123                            }
124                            // 处理近跳转 (E9 disp32)
125                            else if bytes.len() == 2 && bytes[0] == 0xEB {
126                                let offset = (label_pos as i32) - (current_pos as i32 + 2);
127                                if offset >= -128 && offset <= 127 {
128                                    bytes[1] = offset as u8;
129                                }
130                                else {
131                                    // 自动升级为长跳转 (E9 disp32)
132                                    bytes = vec![0xE9, 0, 0, 0, 0];
133                                    let offset = (label_pos as i32) - (current_pos as i32 + 5);
134                                    bytes[1..5].copy_from_slice(&offset.to_le_bytes());
135                                }
136                            }
137                            code.extend(bytes.clone());
138                            current_pos += bytes.len();
139                        }
140                        else {
141                            return Err(gaia_types::GaiaError::not_implemented(format!("Label not found: {}", name)));
142                        }
143                    }
144                    else {
145                        let bytes = assembler.encode(instruction)?;
146                        current_pos += bytes.len();
147                        code.extend(bytes);
148                    }
149                }
150                _ => {
151                    let bytes = assembler.encode(instruction)?;
152                    current_pos += bytes.len();
153                    code.extend(bytes);
154                }
155            }
156        }
157
158        Ok(code)
159    }
160
161    // === 便捷方法:创建常用指令 ===
162
163    /// 创建 MOV 寄存器到寄存器指令
164    pub fn mov_reg_reg(&mut self, dest: Register, src: Register) -> Result<&mut Self> {
165        self.add_instruction(Instruction::Mov { dst: Operand::reg(dest), src: Operand::reg(src) });
166        Ok(self)
167    }
168
169    /// 创建 MOV 立即数到寄存器指令
170    pub fn mov_reg_imm(&mut self, dest: Register, value: i64) -> Result<&mut Self> {
171        // 根据架构和寄存器类型确定立即数大小
172        let size = match self.architecture {
173            Architecture::X86 => 32,
174            Architecture::X86_64 => {
175                // 对于 x64,检查寄存器类型
176                match dest {
177                    Register::EAX
178                    | Register::EBX
179                    | Register::ECX
180                    | Register::EDX
181                    | Register::ESI
182                    | Register::EDI
183                    | Register::ESP
184                    | Register::EBP
185                    | Register::R8D
186                    | Register::R9D
187                    | Register::R10D
188                    | Register::R11D
189                    | Register::R12D
190                    | Register::R13D
191                    | Register::R14D
192                    | Register::R15D => 32,
193                    _ => 64,
194                }
195            }
196            _ => 32,
197        };
198        self.add_instruction(Instruction::Mov { dst: Operand::reg(dest), src: Operand::imm(value, size) });
199        Ok(self)
200    }
201
202    /// 创建 MOV 标签地址到寄存器指令
203    pub fn mov_reg_label(&mut self, dest: Register, label: String) -> Result<&mut Self> {
204        self.add_instruction(Instruction::Mov { dst: Operand::reg(dest), src: Operand::label(label) });
205        Ok(self)
206    }
207
208    /// 创建 PUSH 立即数指令
209    pub fn push_imm(&mut self, value: i64) -> Result<&mut Self> {
210        self.add_instruction(Instruction::Push { op: Operand::imm(value, 32) });
211        Ok(self)
212    }
213
214    /// 创建 PUSH 寄存器指令
215    pub fn push_reg(&mut self, reg: Register) -> Result<&mut Self> {
216        self.add_instruction(Instruction::Push { op: Operand::reg(reg) });
217        Ok(self)
218    }
219
220    /// 创建 PUSH 标签地址指令
221    pub fn push_label(&mut self, label: String) -> Result<&mut Self> {
222        self.add_instruction(Instruction::Push { op: Operand::label(label) });
223        Ok(self)
224    }
225
226    /// 创建 POP 寄存器指令
227    pub fn pop_reg(&mut self, reg: Register) -> Result<&mut Self> {
228        self.add_instruction(Instruction::Pop { dst: Operand::reg(reg) });
229        Ok(self)
230    }
231
232    /// 创建 CALL 指令
233    pub fn call(&mut self, target: String) -> Result<&mut Self> {
234        self.add_instruction(Instruction::Call { target: Operand::label(target) });
235        Ok(self)
236    }
237
238    /// 创建间接 CALL 指令 - CALL [address]
239    pub fn call_indirect(&mut self, address: i32) -> Result<&mut Self> {
240        // CALL [address]
241        self.add_instruction(Instruction::Call { target: Operand::mem(None, None, 1, address) });
242        Ok(self)
243    }
244
245    /// 创建 RET 指令
246    pub fn ret(&mut self) -> Result<&mut Self> {
247        self.add_instruction(Instruction::Ret);
248        Ok(self)
249    }
250
251    /// 创建 NOP 指令
252    pub fn nop(&mut self) -> Result<&mut Self> {
253        self.add_instruction(Instruction::Nop);
254        Ok(self)
255    }
256
257    /// 创建 ADD 指令
258    pub fn add_reg_imm(&mut self, dest: Register, value: i64) -> Result<&mut Self> {
259        self.add_instruction(Instruction::Add { dst: Operand::reg(dest), src: Operand::imm(value, 32) });
260        Ok(self)
261    }
262
263    /// 创建 SUB 指令
264    pub fn sub_reg_imm(&mut self, dest: Register, value: i64) -> Result<&mut Self> {
265        self.add_instruction(Instruction::Sub { dst: Operand::reg(dest), src: Operand::imm(value, 32) });
266        Ok(self)
267    }
268}
269
270impl Default for ProgramBuilder {
271    fn default() -> Self {
272        Self::new(Architecture::X86_64)
273    }
274}