python_ir/
lib.rs

1//! Python intermediate representation for Rusty Python
2//!
3//! This crate provides the intermediate representation (IR) for Python code,
4//! which is used for optimization and code generation.
5
6#![warn(missing_docs)]
7
8use oak_python::ast::{BinaryOperator, Expression, Literal, Statement};
9use python_types::{PythonError, PythonResult, PythonValue};
10
11/// IR instruction
12#[derive(Debug, Clone, PartialEq)]
13pub enum Instruction {
14    // 加载指令
15    LoadConst(PythonValue), // 加载常量到寄存器
16    LoadLocal(usize),       // 加载局部变量到寄存器
17    LoadGlobal(String),     // 加载全局变量到寄存器
18    LoadAttr(String),       // 加载属性到寄存器
19    LoadIndex(usize),       // 加载索引到寄存器
20
21    // 存储指令
22    StoreLocal(usize),   // 存储寄存器值到局部变量
23    StoreGlobal(String), // 存储寄存器值到全局变量
24    StoreAttr(String),   // 存储寄存器值到属性
25    StoreIndex(usize),   // 存储寄存器值到索引
26
27    // 算术指令
28    Add,      // 加法
29    Sub,      // 减法
30    Mul,      // 乘法
31    Div,      // 除法
32    Mod,      // 取模
33    Exp,      // 幂运算
34    FloorDiv, // 地板除
35
36    // 比较指令
37    Eq,  // 等于
38    Neq, // 不等于
39    Lt,  // 小于
40    Lte, // 小于等于
41    Gt,  // 大于
42    Gte, // 大于等于
43
44    // 逻辑指令
45    And, // 逻辑与
46    Or,  // 逻辑或
47    Not, // 逻辑非
48
49    // 控制流指令
50    Jump(i32),          // 无条件跳转
51    JumpIfFalse(i32),   // 条件跳转（如果为假）
52    JumpIfTrue(i32),    // 条件跳转（如果为真）
53    JumpIfNone(i32),    // 条件跳转（如果为 None）
54    JumpIfNotNone(i32), // 条件跳转（如果不为 None）
55
56    // 方法调用指令
57    Call(usize), // 调用函数，参数为参数个数
58    Return,      // 返回
59
60    // 数组和字典指令
61    NewList(usize), // 创建新列表，参数为元素个数
62    NewDict(usize), // 创建新字典，参数为键值对个数
63
64    // 特殊指令
65    NoneOp,  // 加载 None
66    TrueOp,  // 加载 True
67    FalseOp, // 加载 False
68
69    // 异常处理指令
70    Raise,    // 抛出异常
71    TryBegin, // 开始 try 块
72    TryEnd,   // 结束 try 块
73    Except,   // 异常处理
74}
75
76/// IR function
77#[derive(Debug, Clone, PartialEq)]
78pub struct Function {
79    /// Function name
80    pub name: String,
81    /// Parameters
82    pub params: Vec<String>,
83    /// Instructions
84    pub instructions: Vec<Instruction>,
85    /// Local variables
86    pub locals: Vec<String>,
87}
88
89/// IR module
90#[derive(Debug, Clone, PartialEq)]
91pub struct Module {
92    /// Module name
93    pub name: String,
94    /// Functions
95    pub functions: Vec<Function>,
96    /// Global variables
97    pub globals: Vec<String>,
98}
99
100/// Convert AST to IR
101pub fn ast_to_ir(ast: &oak_python::PythonRoot) -> PythonResult<Module> {
102    let mut module = Module { name: "__main__".to_string(), functions: Vec::new(), globals: Vec::new() };
103
104    // Create a default main function if no functions are defined
105    if module.functions.is_empty() {
106        let main_function =
107            Function { name: "__main__".to_string(), params: Vec::new(), instructions: Vec::new(), locals: Vec::new() };
108        module.functions.push(main_function);
109    }
110
111    // Process each statement in the program
112    for statement in &ast.program.statements {
113        match statement {
114            Statement::FunctionDef { name, parameters, body, .. } => {
115                // Create a new function
116                let mut function = Function {
117                    name: name.clone(),
118                    params: parameters.iter().map(|p| p.name.clone()).collect(),
119                    instructions: Vec::new(),
120                    locals: parameters.iter().map(|p| p.name.clone()).collect(),
121                };
122
123                // Process function body
124                for stmt in body {
125                    process_statement(stmt, &mut function.instructions, &mut function.locals)?;
126                }
127
128                // Add return statement if not present
129                if !function.instructions.iter().any(|inst| matches!(inst, Instruction::Return)) {
130                    function.instructions.push(Instruction::NoneOp);
131                    function.instructions.push(Instruction::Return);
132                }
133
134                // Replace the default main function if this is the first function
135                if module.functions.len() == 1 && module.functions[0].name == "__main__" {
136                    module.functions[0] = function;
137                }
138                else {
139                    module.functions.push(function);
140                }
141            }
142            Statement::Assignment { target, value } => {
143                // Process assignment
144                process_expression(value, &mut module.functions.last_mut().unwrap().instructions)?;
145
146                // Handle target
147                if let Expression::Name(name) = target {
148                    // Check if it's a local variable
149                    if let Some(func) = module.functions.last_mut() {
150                        if func.locals.contains(name) {
151                            let index = func.locals.iter().position(|x| x == name).unwrap();
152                            func.instructions.push(Instruction::StoreLocal(index));
153                        }
154                        else {
155                            func.instructions.push(Instruction::StoreGlobal(name.clone()));
156                            if !module.globals.contains(name) {
157                                module.globals.push(name.clone());
158                            }
159                        }
160                    }
161                }
162            }
163            Statement::Expression(expr) => {
164                // Process expression
165                process_expression(expr, &mut module.functions.last_mut().unwrap().instructions)?;
166            }
167            Statement::Return(expr) => {
168                // Process return value
169                if let Some(expr) = expr {
170                    process_expression(expr, &mut module.functions.last_mut().unwrap().instructions)?;
171                }
172                else {
173                    module.functions.last_mut().unwrap().instructions.push(Instruction::NoneOp);
174                }
175                module.functions.last_mut().unwrap().instructions.push(Instruction::Return);
176            }
177            _ => {
178                // TODO: Handle other statement types
179            }
180        }
181    }
182
183    // Add return statement to main function if not present
184    if let Some(func) = module.functions.last_mut() {
185        if !func.instructions.iter().any(|inst| matches!(inst, Instruction::Return)) {
186            func.instructions.push(Instruction::NoneOp);
187            func.instructions.push(Instruction::Return);
188        }
189    }
190
191    Ok(module)
192}
193
194/// Process a statement and generate IR instructions
195fn process_statement(stmt: &Statement, instructions: &mut Vec<Instruction>, locals: &mut Vec<String>) -> PythonResult<()> {
196    match stmt {
197        Statement::Assignment { target, value } => {
198            // Process value expression
199            process_expression(value, instructions)?;
200
201            // Handle target
202            if let Expression::Name(name) = target {
203                if !locals.contains(name) {
204                    locals.push(name.clone());
205                }
206                let index = locals.iter().position(|x| x == name).unwrap();
207                instructions.push(Instruction::StoreLocal(index));
208            }
209        }
210        Statement::Expression(expr) => {
211            process_expression(expr, instructions)?;
212        }
213        Statement::Return(expr) => {
214            if let Some(expr) = expr {
215                process_expression(expr, instructions)?;
216            }
217            else {
218                instructions.push(Instruction::NoneOp);
219            }
220            instructions.push(Instruction::Return);
221        }
222        _ => {
223            // TODO: Handle other statement types
224        }
225    }
226    Ok(())
227}
228
229/// Process an expression and generate IR instructions
230fn process_expression(expr: &Expression, instructions: &mut Vec<Instruction>) -> PythonResult<()> {
231    match expr {
232        Expression::Literal(lit) => {
233            match lit {
234                Literal::Integer(i) => {
235                    instructions.push(Instruction::LoadConst(PythonValue::Integer(*i)));
236                }
237                Literal::Float(f) => {
238                    instructions.push(Instruction::LoadConst(PythonValue::Float(*f)));
239                }
240                Literal::String(s) => {
241                    instructions.push(Instruction::LoadConst(PythonValue::String(s.clone())));
242                }
243                Literal::Boolean(b) => {
244                    if *b {
245                        instructions.push(Instruction::TrueOp);
246                    }
247                    else {
248                        instructions.push(Instruction::FalseOp);
249                    }
250                }
251                Literal::None => {
252                    instructions.push(Instruction::NoneOp);
253                }
254                _ => {
255                    // TODO: Handle other literals
256                }
257            }
258        }
259        Expression::Name(name) => {
260            // TODO: Check if it's a local or global variable
261            instructions.push(Instruction::LoadGlobal(name.clone()));
262        }
263        Expression::BinaryOp { left, operator, right } => {
264            process_expression(left, instructions)?;
265            process_expression(right, instructions)?;
266
267            match operator {
268                BinaryOperator::Add => instructions.push(Instruction::Add),
269                BinaryOperator::Sub => instructions.push(Instruction::Sub),
270                BinaryOperator::Mult => instructions.push(Instruction::Mul),
271                BinaryOperator::Div => instructions.push(Instruction::Div),
272                _ => {
273                    // TODO: Handle other binary operators
274                }
275            }
276        }
277        Expression::Call { func, args, .. } => {
278            // Process arguments
279            for arg in args {
280                process_expression(arg, instructions)?;
281            }
282
283            // Process function
284            process_expression(func, instructions)?;
285
286            // Call function
287            instructions.push(Instruction::Call(args.len() as usize));
288        }
289        _ => {
290            // TODO: Handle other expression types
291        }
292    }
293    Ok(())
294}
295
296/// Optimize IR
297pub fn optimize_ir(ir: &Module) -> PythonResult<Module> {
298    let mut optimized_module = Module { name: ir.name.clone(), functions: Vec::new(), globals: ir.globals.clone() };
299
300    // Optimize each function
301    for func in &ir.functions {
302        let mut optimized_func = Function {
303            name: func.name.clone(),
304            params: func.params.clone(),
305            instructions: optimize_instructions(&func.instructions),
306            locals: func.locals.clone(),
307        };
308        optimized_module.functions.push(optimized_func);
309    }
310
311    Ok(optimized_module)
312}
313
314/// Optimize a list of instructions
315fn optimize_instructions(instructions: &[Instruction]) -> Vec<Instruction> {
316    let mut optimized = Vec::new();
317    let mut i = 0;
318
319    while i < instructions.len() {
320        // TODO: Implement optimization passes
321        // For now, just copy the instructions
322        optimized.push(instructions[i].clone());
323        i += 1;
324    }
325
326    optimized
327}
python_ir/lib.rs

python_ir/
lib.rs
