diff --git a/tests/snippets/dismod.py b/tests/snippets/dismod.py new file mode 100644 index 0000000000..cc9c1f63b8 --- /dev/null +++ b/tests/snippets/dismod.py @@ -0,0 +1,10 @@ +import dis + +dis.disassemble(compile("5 + x + 5 or 2", "", "eval")) +print("\n") +dis.disassemble(compile("def f(x):\n return 1", "", "exec")) +print("\n") +dis.disassemble(compile("if a:\n 1 or 2\nelif x == 'hello':\n 3\nelse:\n 4", "", "exec")) +print("\n") +dis.disassemble(compile("f(x=1, y=2)", "", "eval")) +print("\n") diff --git a/vm/src/bytecode.rs b/vm/src/bytecode.rs index 951659f066..1e0b85f133 100644 --- a/vm/src/bytecode.rs +++ b/vm/src/bytecode.rs @@ -8,7 +8,7 @@ use num_bigint::BigInt; use num_complex::Complex64; use rustpython_parser::ast; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::fmt; /// Primary container of a single code object. Each python function has @@ -173,6 +173,8 @@ pub enum Instruction { }, } +use self::Instruction::*; + #[derive(Debug, Clone, PartialEq)] pub enum CallType { Positional(usize), @@ -277,17 +279,123 @@ impl CodeObject { } } +impl fmt::Display for CodeObject { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let label_targets: HashSet<&usize> = self.label_map.values().collect(); + for (offset, instruction) in self.instructions.iter().enumerate() { + let arrow = if label_targets.contains(&offset) { + ">>" + } else { + " " + }; + write!(f, " {} {:5} ", arrow, offset)?; + instruction.fmt_dis(f, &self.label_map)?; + } + Ok(()) + } +} + +impl Instruction { + fn fmt_dis(&self, f: &mut fmt::Formatter, label_map: &HashMap) -> fmt::Result { + macro_rules! w { + ($variant:ident) => { + write!(f, "{:20}\n", stringify!($variant)) + }; + ($variant:ident, $var:expr) => { + write!(f, "{:20} ({})\n", stringify!($variant), $var) + }; + ($variant:ident, $var1:expr, $var2:expr) => { + write!(f, "{:20} ({}, {})\n", stringify!($variant), $var1, $var2) + }; + } + + match self { + Import { name, symbol } => w!(Import, name, format!("{:?}", symbol)), + ImportStar { name } => w!(ImportStar, name), + LoadName { name } => w!(LoadName, name), + StoreName { name } => w!(StoreName, name), + DeleteName { name } => w!(DeleteName, name), + StoreSubscript => w!(StoreSubscript), + DeleteSubscript => w!(DeleteSubscript), + StoreAttr { name } => w!(StoreAttr, name), + DeleteAttr { name } => w!(DeleteAttr, name), + LoadConst { value } => w!(LoadConst, value), + UnaryOperation { op } => w!(UnaryOperation, format!("{:?}", op)), + BinaryOperation { op, inplace } => w!(BinaryOperation, format!("{:?}", op), inplace), + LoadAttr { name } => w!(LoadAttr, name), + CompareOperation { op } => w!(CompareOperation, format!("{:?}", op)), + Pop => w!(Pop), + Rotate { amount } => w!(Rotate, amount), + Duplicate => w!(Duplicate), + GetIter => w!(GetIter), + Pass => w!(Pass), + Continue => w!(Continue), + Break => w!(Break), + Jump { target } => w!(Jump, label_map[target]), + JumpIf { target } => w!(JumpIf, label_map[target]), + JumpIfFalse { target } => w!(JumpIfFalse, label_map[target]), + MakeFunction { flags } => w!(MakeFunction, format!("{:?}", flags)), + CallFunction { typ } => w!(CallFunction, format!("{:?}", typ)), + ForIter { target } => w!(ForIter, label_map[target]), + ReturnValue => w!(ReturnValue), + YieldValue => w!(YieldValue), + YieldFrom => w!(YieldFrom), + SetupLoop { start, end } => w!(SetupLoop, label_map[start], label_map[end]), + SetupExcept { handler } => w!(SetupExcept, handler), + SetupWith { end } => w!(SetupWith, end), + CleanupWith { end } => w!(CleanupWith, end), + PopBlock => w!(PopBlock), + Raise { argc } => w!(Raise, argc), + BuildString { size } => w!(BuildString, size), + BuildTuple { size, unpack } => w!(BuildTuple, size, unpack), + BuildList { size, unpack } => w!(BuildList, size, unpack), + BuildSet { size, unpack } => w!(BuildSet, size, unpack), + BuildMap { size, unpack } => w!(BuildMap, size, unpack), + BuildSlice { size } => w!(BuildSlice, size), + ListAppend { i } => w!(ListAppend, i), + SetAdd { i } => w!(SetAdd, i), + MapAdd { i } => w!(MapAdd, i), + PrintExpr => w!(PrintExpr), + LoadBuildClass => w!(LoadBuildClass), + StoreLocals => w!(StoreLocals), + UnpackSequence { size } => w!(UnpackSequence, size), + UnpackEx { before, after } => w!(UnpackEx, before, after), + Unpack => w!(Unpack), + FormatValue { spec } => w!(FormatValue, spec), + } + } +} + +impl fmt::Display for Constant { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + Constant::Integer { value } => write!(f, "{}", value), + Constant::Float { value } => write!(f, "{}", value), + Constant::Complex { value } => write!(f, "{}", value), + Constant::Boolean { value } => write!(f, "{}", value), + Constant::String { value } => write!(f, "{:?}", value), + Constant::Bytes { value } => write!(f, "{:?}", value), + Constant::Code { code } => write!(f, "{:?}", code), + Constant::Tuple { elements } => write!( + f, + "({})", + elements + .iter() + .map(|e| format!("{}", e)) + .collect::>() + .join(", ") + ), + Constant::None => write!(f, "None"), + } + } +} + impl fmt::Debug for CodeObject { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - let inst_str = self - .instructions - .iter() - .zip(self.locations.iter()) - .enumerate() - .map(|(i, inst)| format!("Inst {}: {:?}", i, inst)) - .collect::>() - .join("\n"); - let labelmap_str = format!("label_map: {:?}", self.label_map); - write!(f, "Code Object {{ \n{}\n{} }}", inst_str, labelmap_str) + write!( + f, + "", + self.obj_name, self.source_path, self.first_line_number + ) } } diff --git a/vm/src/stdlib/dis.rs b/vm/src/stdlib/dis.rs new file mode 100644 index 0000000000..ec95725e77 --- /dev/null +++ b/vm/src/stdlib/dis.rs @@ -0,0 +1,18 @@ +use super::super::obj::objcode; +use super::super::obj::objtype; +use super::super::pyobject::{PyContext, PyFuncArgs, PyObjectRef, PyResult, TypeProtocol}; +use super::super::vm::VirtualMachine; + +fn dis_disassemble(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { + arg_check!(vm, args, required = [(co, Some(vm.ctx.code_type()))]); + + let code = objcode::get_value(co); + print!("{}", code); + Ok(vm.get_none()) +} + +pub fn mk_module(ctx: &PyContext) -> PyObjectRef { + let py_mod = ctx.new_module("dis", ctx.new_scope(None)); + ctx.set_attr(&py_mod, "disassemble", ctx.new_rustfunc(dis_disassemble)); + py_mod +} diff --git a/vm/src/stdlib/mod.rs b/vm/src/stdlib/mod.rs index 159d1cd53d..9b6158d210 100644 --- a/vm/src/stdlib/mod.rs +++ b/vm/src/stdlib/mod.rs @@ -1,4 +1,5 @@ mod ast; +mod dis; pub mod io; mod json; mod keyword; @@ -21,6 +22,7 @@ pub type StdlibInitFunc = fn(&PyContext) -> PyObjectRef; pub fn get_module_inits() -> HashMap { let mut modules = HashMap::new(); modules.insert("ast".to_string(), ast::mk_module as StdlibInitFunc); + modules.insert("dis".to_string(), dis::mk_module as StdlibInitFunc); modules.insert("io".to_string(), io::mk_module as StdlibInitFunc); modules.insert("json".to_string(), json::mk_module as StdlibInitFunc); modules.insert("keyword".to_string(), keyword::mk_module as StdlibInitFunc);