Skip to content

Commit 704b157

Browse files
committed
Add optimize flag to ignore assert statements.
1 parent 8c1c090 commit 704b157

File tree

11 files changed

+90
-56
lines changed

11 files changed

+90
-56
lines changed

compiler/src/compile.rs

Lines changed: 53 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -20,32 +20,39 @@ struct Compiler {
2020
current_qualified_path: Option<String>,
2121
in_loop: bool,
2222
in_function_def: bool,
23+
optimize: bool,
2324
}
2425

2526
/// Compile a given sourcecode into a bytecode object.
26-
pub fn compile(source: &str, mode: &Mode, source_path: String) -> Result<CodeObject, CompileError> {
27+
pub fn compile(
28+
source: &str,
29+
mode: &Mode,
30+
source_path: String,
31+
optimize: bool,
32+
) -> Result<CodeObject, CompileError> {
2733
match mode {
2834
Mode::Exec => {
2935
let ast = parser::parse_program(source)?;
30-
compile_program(ast, source_path)
36+
compile_program(ast, source_path, optimize)
3137
}
3238
Mode::Eval => {
3339
let statement = parser::parse_statement(source)?;
34-
compile_statement_eval(statement, source_path)
40+
compile_statement_eval(statement, source_path, optimize)
3541
}
3642
Mode::Single => {
3743
let ast = parser::parse_program(source)?;
38-
compile_program_single(ast, source_path)
44+
compile_program_single(ast, source_path, optimize)
3945
}
4046
}
4147
}
4248

4349
/// A helper function for the shared code of the different compile functions
4450
fn with_compiler(
4551
source_path: String,
52+
optimize: bool,
4653
f: impl FnOnce(&mut Compiler) -> Result<(), CompileError>,
4754
) -> Result<CodeObject, CompileError> {
48-
let mut compiler = Compiler::new();
55+
let mut compiler = Compiler::new(optimize);
4956
compiler.source_path = Some(source_path);
5057
compiler.push_new_code_object("<module>".to_string());
5158
f(&mut compiler)?;
@@ -55,8 +62,12 @@ fn with_compiler(
5562
}
5663

5764
/// Compile a standard Python program to bytecode
58-
pub fn compile_program(ast: ast::Program, source_path: String) -> Result<CodeObject, CompileError> {
59-
with_compiler(source_path, |compiler| {
65+
pub fn compile_program(
66+
ast: ast::Program,
67+
source_path: String,
68+
optimize: bool,
69+
) -> Result<CodeObject, CompileError> {
70+
with_compiler(source_path, optimize, |compiler| {
6071
let symbol_table = make_symbol_table(&ast)?;
6172
compiler.compile_program(&ast, symbol_table)
6273
})
@@ -66,8 +77,9 @@ pub fn compile_program(ast: ast::Program, source_path: String) -> Result<CodeObj
6677
pub fn compile_statement_eval(
6778
statement: Vec<ast::LocatedStatement>,
6879
source_path: String,
80+
optimize: bool,
6981
) -> Result<CodeObject, CompileError> {
70-
with_compiler(source_path, |compiler| {
82+
with_compiler(source_path, optimize, |compiler| {
7183
let symbol_table = statements_to_symbol_table(&statement)?;
7284
compiler.compile_statement_eval(&statement, symbol_table)
7385
})
@@ -77,8 +89,9 @@ pub fn compile_statement_eval(
7789
pub fn compile_program_single(
7890
ast: ast::Program,
7991
source_path: String,
92+
optimize: bool,
8093
) -> Result<CodeObject, CompileError> {
81-
with_compiler(source_path, |compiler| {
94+
with_compiler(source_path, optimize, |compiler| {
8295
let symbol_table = make_symbol_table(&ast)?;
8396
compiler.compile_program_single(&ast, symbol_table)
8497
})
@@ -98,8 +111,14 @@ enum EvalContext {
98111

99112
type Label = usize;
100113

114+
impl Default for Compiler {
115+
fn default() -> Self {
116+
Compiler::new(false)
117+
}
118+
}
119+
101120
impl Compiler {
102-
fn new() -> Self {
121+
fn new(optimize: bool) -> Self {
103122
Compiler {
104123
code_object_stack: Vec::new(),
105124
scope_stack: Vec::new(),
@@ -109,6 +128,7 @@ impl Compiler {
109128
current_qualified_path: None,
110129
in_loop: false,
111130
in_function_def: false,
131+
optimize,
112132
}
113133
}
114134

@@ -434,29 +454,30 @@ impl Compiler {
434454
decorator_list,
435455
} => self.compile_class_def(name, body, bases, keywords, decorator_list)?,
436456
ast::Statement::Assert { test, msg } => {
437-
// TODO: if some flag, ignore all assert statements!
438-
439-
let end_label = self.new_label();
440-
self.compile_test(test, Some(end_label), None, EvalContext::Statement)?;
441-
self.emit(Instruction::LoadName {
442-
name: String::from("AssertionError"),
443-
scope: bytecode::NameScope::Local,
444-
});
445-
match msg {
446-
Some(e) => {
447-
self.compile_expression(e)?;
448-
self.emit(Instruction::CallFunction {
449-
typ: CallType::Positional(1),
450-
});
451-
}
452-
None => {
453-
self.emit(Instruction::CallFunction {
454-
typ: CallType::Positional(0),
455-
});
457+
// if some flag, ignore all assert statements!
458+
if !self.optimize {
459+
let end_label = self.new_label();
460+
self.compile_test(test, Some(end_label), None, EvalContext::Statement)?;
461+
self.emit(Instruction::LoadName {
462+
name: String::from("AssertionError"),
463+
scope: bytecode::NameScope::Local,
464+
});
465+
match msg {
466+
Some(e) => {
467+
self.compile_expression(e)?;
468+
self.emit(Instruction::CallFunction {
469+
typ: CallType::Positional(1),
470+
});
471+
}
472+
None => {
473+
self.emit(Instruction::CallFunction {
474+
typ: CallType::Positional(0),
475+
});
476+
}
456477
}
478+
self.emit(Instruction::Raise { argc: 1 });
479+
self.set_label(end_label);
457480
}
458-
self.emit(Instruction::Raise { argc: 1 });
459-
self.set_label(end_label);
460481
}
461482
ast::Statement::Break => {
462483
if !self.in_loop {
@@ -1868,7 +1889,7 @@ mod tests {
18681889
use rustpython_parser::parser;
18691890

18701891
fn compile_exec(source: &str) -> CodeObject {
1871-
let mut compiler = Compiler::new();
1892+
let mut compiler: Compiler = Default::default();
18721893
compiler.source_path = Some("source_path".to_string());
18731894
compiler.push_new_code_object("<module>".to_string());
18741895
let ast = parser::parse_program(&source.to_string()).unwrap();

derive/src/compile_bytecode.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ struct CompilationSource {
3838
impl CompilationSource {
3939
fn compile(self, mode: &compile::Mode, module_name: String) -> Result<CodeObject, Diagnostic> {
4040
let compile = |source| {
41-
compile::compile(source, mode, module_name).map_err(|err| {
41+
compile::compile(source, mode, module_name, false).map_err(|err| {
4242
Diagnostic::spans_error(self.span, format!("Compile error: {}", err))
4343
})
4444
};

src/main.rs

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,12 @@ fn main() {
3030
.author(crate_authors!())
3131
.about("Rust implementation of the Python language")
3232
.arg(Arg::with_name("script").required(false).index(1))
33+
.arg(
34+
Arg::with_name("optimize")
35+
.short("O")
36+
.multiple(true)
37+
.help("Optimize. Set __debug__ to false. Remove debug statements."),
38+
)
3339
.arg(
3440
Arg::with_name("v")
3541
.short("v")
@@ -65,8 +71,11 @@ fn main() {
6571
);
6672
let matches = app.get_matches();
6773

74+
let opt_level = matches.occurrences_of("optimize");
75+
let optimize = opt_level > 0;
76+
debug!("Optimize: {}", optimize);
6877
// Construct vm:
69-
let vm = VirtualMachine::new();
78+
let vm = VirtualMachine::new(optimize);
7079

7180
let res = import::init_importlib(&vm, true);
7281
handle_exception(&vm, res);
@@ -215,7 +224,7 @@ fn run_script(vm: &VirtualMachine, script_file: &str) -> PyResult {
215224

216225
#[test]
217226
fn test_run_script() {
218-
let vm = VirtualMachine::new();
227+
let vm: VirtualMachine = Default::default();
219228

220229
// test file run
221230
let r = run_script(&vm, "tests/snippets/dir_main/__main__.py");

vm/src/builtins.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -798,7 +798,9 @@ pub fn make_module(vm: &VirtualMachine, module: PyObjectRef) {
798798
});
799799
}
800800

801+
let debug_mode = !vm.optimize;
801802
extend_module!(vm, module, {
803+
"__debug__" => ctx.new_bool(debug_mode),
802804
//set __name__ fixes: https://github.com/RustPython/RustPython/issues/146
803805
"__name__" => ctx.new_str(String::from("__main__")),
804806

vm/src/dictdatatype.rs

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -290,29 +290,29 @@ mod tests {
290290

291291
#[test]
292292
fn test_insert() {
293-
let mut vm = VirtualMachine::new();
293+
let vm: VirtualMachine = Default::default();
294294
let mut dict = Dict::default();
295295
assert_eq!(0, dict.len());
296296

297297
let key1 = vm.new_bool(true);
298298
let value1 = vm.new_str("abc".to_string());
299-
dict.insert(&mut vm, &key1, value1.clone()).unwrap();
299+
dict.insert(&vm, &key1, value1.clone()).unwrap();
300300
assert_eq!(1, dict.len());
301301

302302
let key2 = vm.new_str("x".to_string());
303303
let value2 = vm.new_str("def".to_string());
304-
dict.insert(&mut vm, &key2, value2.clone()).unwrap();
304+
dict.insert(&vm, &key2, value2.clone()).unwrap();
305305
assert_eq!(2, dict.len());
306306

307-
dict.insert(&mut vm, &key1, value2.clone()).unwrap();
307+
dict.insert(&vm, &key1, value2.clone()).unwrap();
308308
assert_eq!(2, dict.len());
309309

310-
dict.delete(&mut vm, &key1).unwrap();
310+
dict.delete(&vm, &key1).unwrap();
311311
assert_eq!(1, dict.len());
312312

313-
dict.insert(&mut vm, &key1, value2).unwrap();
313+
dict.insert(&vm, &key1, value2).unwrap();
314314
assert_eq!(2, dict.len());
315315

316-
assert_eq!(true, dict.contains(&mut vm, &key1).unwrap());
316+
assert_eq!(true, dict.contains(&vm, &key1).unwrap());
317317
}
318318
}

vm/src/eval.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,9 @@ mod tests {
2121
#[test]
2222
fn test_print_42() {
2323
let source = String::from("print('Hello world')\n");
24-
let mut vm = VirtualMachine::new();
24+
let vm: VirtualMachine = Default::default();
2525
let vars = vm.new_scope_with_builtins();
26-
let _result = eval(&mut vm, &source, vars, "<unittest>");
26+
let _result = eval(&vm, &source, vars, "<unittest>");
2727

2828
// TODO: check result?
2929
//assert_eq!(

vm/src/import.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ pub fn import_file(
5555
file_path: String,
5656
content: String,
5757
) -> PyResult {
58-
let code_obj = compile::compile(&content, &compile::Mode::Exec, file_path)
58+
let code_obj = compile::compile(&content, &compile::Mode::Exec, file_path, vm.optimize)
5959
.map_err(|err| vm.new_syntax_error(&err))?;
6060
import_codeobj(vm, module_name, code_obj, true)
6161
}

vm/src/macros.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -188,7 +188,7 @@ macro_rules! py_namespace {
188188
/// use rustpython_vm::obj::objint::PyInt;
189189
/// use rustpython_vm::pyobject::PyValue;
190190
///
191-
/// let vm = VirtualMachine::new();
191+
/// let vm: VirtualMachine = Default::default();
192192
/// let obj = PyInt::new(0).into_ref(&vm).into_object();
193193
/// assert_eq!(
194194
/// "int",
@@ -213,7 +213,7 @@ macro_rules! py_namespace {
213213
/// use rustpython_vm::obj::objint::PyInt;
214214
/// use rustpython_vm::pyobject::PyValue;
215215
///
216-
/// let vm = VirtualMachine::new();
216+
/// let vm: VirtualMachine = Default::default();
217217
/// let obj = PyInt::new(0).into_ref(&vm).into_object();
218218
///
219219
/// let int_value = match_class!(obj,

vm/src/obj/objstr.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1467,7 +1467,7 @@ mod tests {
14671467

14681468
#[test]
14691469
fn str_title() {
1470-
let vm = VirtualMachine::new();
1470+
let vm: VirtualMachine = Default::default();
14711471

14721472
let tests = vec![
14731473
(" Hello ", " hello "),
@@ -1486,7 +1486,7 @@ mod tests {
14861486

14871487
#[test]
14881488
fn str_istitle() {
1489-
let vm = VirtualMachine::new();
1489+
let vm: VirtualMachine = Default::default();
14901490

14911491
let pos = vec![
14921492
"A",
@@ -1517,7 +1517,7 @@ mod tests {
15171517

15181518
#[test]
15191519
fn str_maketrans_and_translate() {
1520-
let vm = VirtualMachine::new();
1520+
let vm: VirtualMachine = Default::default();
15211521

15221522
let table = vm.context().new_dict();
15231523
table

vm/src/vm.rs

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -56,11 +56,12 @@ pub struct VirtualMachine {
5656
pub exceptions: RefCell<Vec<PyObjectRef>>,
5757
pub frozen: RefCell<HashMap<String, bytecode::CodeObject>>,
5858
pub import_func: RefCell<PyObjectRef>,
59+
pub optimize: bool,
5960
}
6061

6162
impl VirtualMachine {
6263
/// Create a new `VirtualMachine` structure.
63-
pub fn new() -> VirtualMachine {
64+
pub fn new(optimize: bool) -> VirtualMachine {
6465
flame_guard!("init VirtualMachine");
6566
let ctx = PyContext::new();
6667

@@ -81,6 +82,7 @@ impl VirtualMachine {
8182
exceptions: RefCell::new(vec![]),
8283
frozen,
8384
import_func,
85+
optimize,
8486
};
8587

8688
builtins::make_module(&vm, builtins.clone());
@@ -749,7 +751,7 @@ impl VirtualMachine {
749751
mode: &compile::Mode,
750752
source_path: String,
751753
) -> Result<PyCodeRef, CompileError> {
752-
compile::compile(source, mode, source_path)
754+
compile::compile(source, mode, source_path, self.optimize)
753755
.map(|codeobj| PyCode::new(codeobj).into_ref(self))
754756
}
755757

@@ -1022,7 +1024,7 @@ impl VirtualMachine {
10221024

10231025
impl Default for VirtualMachine {
10241026
fn default() -> Self {
1025-
VirtualMachine::new()
1027+
VirtualMachine::new(false)
10261028
}
10271029
}
10281030

@@ -1070,7 +1072,7 @@ mod tests {
10701072

10711073
#[test]
10721074
fn test_add_py_integers() {
1073-
let vm = VirtualMachine::new();
1075+
let vm: VirtualMachine = Default::default();
10741076
let a = vm.ctx.new_int(33_i32);
10751077
let b = vm.ctx.new_int(12_i32);
10761078
let res = vm._add(a, b).unwrap();
@@ -1080,7 +1082,7 @@ mod tests {
10801082

10811083
#[test]
10821084
fn test_multiply_str() {
1083-
let vm = VirtualMachine::new();
1085+
let vm: VirtualMachine = Default::default();
10841086
let a = vm.ctx.new_str(String::from("Hello "));
10851087
let b = vm.ctx.new_int(4_i32);
10861088
let res = vm._mul(a, b).unwrap();

0 commit comments

Comments
 (0)