use bitflags::bitflags;
use itertools::Itertools;
use malachite_bigint::BigInt;
use num_complex::Complex64;
use rustpython_parser_core::source_code::{OneIndexed, SourceLocation};
use std::marker::PhantomData;
use std::{collections::BTreeSet, fmt, hash, mem};
pub use rustpython_parser_core::ConversionFlag;
pub trait Constant: Sized {
type Name: AsRef<str>;
fn borrow_constant(&self) -> BorrowedConstant<Self>;
}
impl Constant for ConstantData {
type Name = String;
fn borrow_constant(&self) -> BorrowedConstant<Self> {
use BorrowedConstant::*;
match self {
ConstantData::Integer { value } => Integer { value },
ConstantData::Float { value } => Float { value: *value },
ConstantData::Complex { value } => Complex { value: *value },
ConstantData::Boolean { value } => Boolean { value: *value },
ConstantData::Str { value } => Str { value },
ConstantData::Bytes { value } => Bytes { value },
ConstantData::Code { code } => Code { code },
ConstantData::Tuple { elements } => Tuple { elements },
ConstantData::None => None,
ConstantData::Ellipsis => Ellipsis,
}
}
}
pub trait ConstantBag: Sized + Copy {
type Constant: Constant;
fn make_constant<C: Constant>(&self, constant: BorrowedConstant<C>) -> Self::Constant;
fn make_int(&self, value: BigInt) -> Self::Constant;
fn make_tuple(&self, elements: impl Iterator<Item = Self::Constant>) -> Self::Constant;
fn make_code(&self, code: CodeObject<Self::Constant>) -> Self::Constant;
fn make_name(&self, name: &str) -> <Self::Constant as Constant>::Name;
}
pub trait AsBag {
type Bag: ConstantBag;
#[allow(clippy::wrong_self_convention)]
fn as_bag(self) -> Self::Bag;
}
impl<Bag: ConstantBag> AsBag for Bag {
type Bag = Self;
fn as_bag(self) -> Self {
self
}
}
#[derive(Clone, Copy)]
pub struct BasicBag;
impl ConstantBag for BasicBag {
type Constant = ConstantData;
fn make_constant<C: Constant>(&self, constant: BorrowedConstant<C>) -> Self::Constant {
constant.to_owned()
}
fn make_int(&self, value: BigInt) -> Self::Constant {
ConstantData::Integer { value }
}
fn make_tuple(&self, elements: impl Iterator<Item = Self::Constant>) -> Self::Constant {
ConstantData::Tuple {
elements: elements.collect(),
}
}
fn make_code(&self, code: CodeObject<Self::Constant>) -> Self::Constant {
ConstantData::Code {
code: Box::new(code),
}
}
fn make_name(&self, name: &str) -> <Self::Constant as Constant>::Name {
name.to_owned()
}
}
#[derive(Clone)]
pub struct CodeObject<C: Constant = ConstantData> {
pub instructions: Box<[CodeUnit]>,
pub locations: Box<[SourceLocation]>,
pub flags: CodeFlags,
pub posonlyarg_count: u32,
pub arg_count: u32,
pub kwonlyarg_count: u32,
pub source_path: C::Name,
pub first_line_number: Option<OneIndexed>,
pub max_stackdepth: u32,
pub obj_name: C::Name,
pub cell2arg: Option<Box<[i32]>>,
pub constants: Box<[C]>,
pub names: Box<[C::Name]>,
pub varnames: Box<[C::Name]>,
pub cellvars: Box<[C::Name]>,
pub freevars: Box<[C::Name]>,
}
bitflags! {
#[derive(Copy, Clone, Debug, PartialEq)]
pub struct CodeFlags: u16 {
const NEW_LOCALS = 0x01;
const IS_GENERATOR = 0x02;
const IS_COROUTINE = 0x04;
const HAS_VARARGS = 0x08;
const HAS_VARKEYWORDS = 0x10;
const IS_OPTIMIZED = 0x20;
}
}
impl CodeFlags {
pub const NAME_MAPPING: &'static [(&'static str, CodeFlags)] = &[
("GENERATOR", CodeFlags::IS_GENERATOR),
("COROUTINE", CodeFlags::IS_COROUTINE),
(
"ASYNC_GENERATOR",
Self::from_bits_truncate(Self::IS_GENERATOR.bits() | Self::IS_COROUTINE.bits()),
),
("VARARGS", CodeFlags::HAS_VARARGS),
("VARKEYWORDS", CodeFlags::HAS_VARKEYWORDS),
];
}
#[derive(Copy, Clone, PartialEq, Eq)]
#[repr(transparent)]
pub struct OpArgByte(pub u8);
impl OpArgByte {
pub const fn null() -> Self {
OpArgByte(0)
}
}
impl fmt::Debug for OpArgByte {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(f)
}
}
#[derive(Copy, Clone, Debug)]
#[repr(transparent)]
pub struct OpArg(pub u32);
impl OpArg {
pub const fn null() -> Self {
OpArg(0)
}
#[inline]
pub fn instr_size(self) -> usize {
(self.0 > 0xff) as usize + (self.0 > 0xff_ff) as usize + (self.0 > 0xff_ff_ff) as usize + 1
}
#[inline(always)]
pub fn split(self) -> (impl ExactSizeIterator<Item = OpArgByte>, OpArgByte) {
let mut it = self
.0
.to_le_bytes()
.map(OpArgByte)
.into_iter()
.take(self.instr_size());
let lo = it.next().unwrap();
(it.rev(), lo)
}
}
#[derive(Default, Copy, Clone)]
#[repr(transparent)]
pub struct OpArgState {
state: u32,
}
impl OpArgState {
#[inline(always)]
pub fn get(&mut self, ins: CodeUnit) -> (Instruction, OpArg) {
let arg = self.extend(ins.arg);
if ins.op != Instruction::ExtendedArg {
self.reset();
}
(ins.op, arg)
}
#[inline(always)]
pub fn extend(&mut self, arg: OpArgByte) -> OpArg {
self.state = self.state << 8 | u32::from(arg.0);
OpArg(self.state)
}
#[inline(always)]
pub fn reset(&mut self) {
self.state = 0
}
}
pub trait OpArgType: Copy {
fn from_op_arg(x: u32) -> Option<Self>;
fn to_op_arg(self) -> u32;
}
impl OpArgType for u32 {
#[inline(always)]
fn from_op_arg(x: u32) -> Option<Self> {
Some(x)
}
#[inline(always)]
fn to_op_arg(self) -> u32 {
self
}
}
impl OpArgType for bool {
#[inline(always)]
fn from_op_arg(x: u32) -> Option<Self> {
Some(x != 0)
}
#[inline(always)]
fn to_op_arg(self) -> u32 {
self as u32
}
}
macro_rules! op_arg_enum_impl {
(enum $name:ident { $($(#[$var_attr:meta])* $var:ident = $value:literal,)* }) => {
impl OpArgType for $name {
fn to_op_arg(self) -> u32 {
self as u32
}
fn from_op_arg(x: u32) -> Option<Self> {
Some(match u8::try_from(x).ok()? {
$($value => Self::$var,)*
_ => return None,
})
}
}
};
}
macro_rules! op_arg_enum {
($(#[$attr:meta])* $vis:vis enum $name:ident { $($(#[$var_attr:meta])* $var:ident = $value:literal,)* }) => {
$(#[$attr])*
$vis enum $name {
$($(#[$var_attr])* $var = $value,)*
}
op_arg_enum_impl!(enum $name {
$($(#[$var_attr])* $var = $value,)*
});
};
}
#[derive(Copy, Clone)]
pub struct Arg<T: OpArgType>(PhantomData<T>);
impl<T: OpArgType> Arg<T> {
#[inline]
pub fn marker() -> Self {
Arg(PhantomData)
}
#[inline]
pub fn new(arg: T) -> (Self, OpArg) {
(Self(PhantomData), OpArg(arg.to_op_arg()))
}
#[inline]
pub fn new_single(arg: T) -> (Self, OpArgByte)
where
T: Into<u8>,
{
(Self(PhantomData), OpArgByte(arg.into()))
}
#[inline(always)]
pub fn get(self, arg: OpArg) -> T {
self.try_get(arg).unwrap()
}
#[inline(always)]
pub fn try_get(self, arg: OpArg) -> Option<T> {
T::from_op_arg(arg.0)
}
#[inline(always)]
pub unsafe fn get_unchecked(self, arg: OpArg) -> T {
match T::from_op_arg(arg.0) {
Some(t) => t,
None => std::hint::unreachable_unchecked(),
}
}
}
impl<T: OpArgType> PartialEq for Arg<T> {
fn eq(&self, _: &Self) -> bool {
true
}
}
impl<T: OpArgType> Eq for Arg<T> {}
impl<T: OpArgType> fmt::Debug for Arg<T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "Arg<{}>", std::any::type_name::<T>())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]
#[repr(transparent)]
pub struct Label(pub u32);
impl OpArgType for Label {
#[inline(always)]
fn from_op_arg(x: u32) -> Option<Self> {
Some(Label(x))
}
#[inline(always)]
fn to_op_arg(self) -> u32 {
self.0
}
}
impl fmt::Display for Label {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
self.0.fmt(f)
}
}
impl OpArgType for ConversionFlag {
#[inline]
fn from_op_arg(x: u32) -> Option<Self> {
match x as u8 {
b's' => Some(ConversionFlag::Str),
b'a' => Some(ConversionFlag::Ascii),
b'r' => Some(ConversionFlag::Repr),
std::u8::MAX => Some(ConversionFlag::None),
_ => None,
}
}
#[inline]
fn to_op_arg(self) -> u32 {
self as i8 as u8 as u32
}
}
op_arg_enum!(
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum RaiseKind {
Reraise = 0,
Raise = 1,
RaiseCause = 2,
}
);
pub type NameIdx = u32;
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
#[repr(u8)]
pub enum Instruction {
ImportName {
idx: Arg<NameIdx>,
},
ImportNameless,
ImportStar,
ImportFrom {
idx: Arg<NameIdx>,
},
LoadFast(Arg<NameIdx>),
LoadNameAny(Arg<NameIdx>),
LoadGlobal(Arg<NameIdx>),
LoadDeref(Arg<NameIdx>),
LoadClassDeref(Arg<NameIdx>),
StoreFast(Arg<NameIdx>),
StoreLocal(Arg<NameIdx>),
StoreGlobal(Arg<NameIdx>),
StoreDeref(Arg<NameIdx>),
DeleteFast(Arg<NameIdx>),
DeleteLocal(Arg<NameIdx>),
DeleteGlobal(Arg<NameIdx>),
DeleteDeref(Arg<NameIdx>),
LoadClosure(Arg<NameIdx>),
Subscript,
StoreSubscript,
DeleteSubscript,
StoreAttr {
idx: Arg<NameIdx>,
},
DeleteAttr {
idx: Arg<NameIdx>,
},
LoadConst {
idx: Arg<u32>,
},
UnaryOperation {
op: Arg<UnaryOperator>,
},
BinaryOperation {
op: Arg<BinaryOperator>,
},
BinaryOperationInplace {
op: Arg<BinaryOperator>,
},
LoadAttr {
idx: Arg<NameIdx>,
},
TestOperation {
op: Arg<TestOperator>,
},
CompareOperation {
op: Arg<ComparisonOperator>,
},
Pop,
Rotate2,
Rotate3,
Duplicate,
Duplicate2,
GetIter,
Continue {
target: Arg<Label>,
},
Break {
target: Arg<Label>,
},
Jump {
target: Arg<Label>,
},
JumpIfTrue {
target: Arg<Label>,
},
JumpIfFalse {
target: Arg<Label>,
},
JumpIfTrueOrPop {
target: Arg<Label>,
},
JumpIfFalseOrPop {
target: Arg<Label>,
},
MakeFunction(Arg<MakeFunctionFlags>),
CallFunctionPositional {
nargs: Arg<u32>,
},
CallFunctionKeyword {
nargs: Arg<u32>,
},
CallFunctionEx {
has_kwargs: Arg<bool>,
},
LoadMethod {
idx: Arg<NameIdx>,
},
CallMethodPositional {
nargs: Arg<u32>,
},
CallMethodKeyword {
nargs: Arg<u32>,
},
CallMethodEx {
has_kwargs: Arg<bool>,
},
ForIter {
target: Arg<Label>,
},
ReturnValue,
YieldValue,
YieldFrom,
SetupAnnotation,
SetupLoop,
SetupFinally {
handler: Arg<Label>,
},
EnterFinally,
EndFinally,
SetupExcept {
handler: Arg<Label>,
},
SetupWith {
end: Arg<Label>,
},
WithCleanupStart,
WithCleanupFinish,
PopBlock,
Raise {
kind: Arg<RaiseKind>,
},
BuildString {
size: Arg<u32>,
},
BuildTuple {
size: Arg<u32>,
},
BuildTupleUnpack {
size: Arg<u32>,
},
BuildList {
size: Arg<u32>,
},
BuildListUnpack {
size: Arg<u32>,
},
BuildSet {
size: Arg<u32>,
},
BuildSetUnpack {
size: Arg<u32>,
},
BuildMap {
size: Arg<u32>,
},
BuildMapForCall {
size: Arg<u32>,
},
DictUpdate,
BuildSlice {
step: Arg<bool>,
},
ListAppend {
i: Arg<u32>,
},
SetAdd {
i: Arg<u32>,
},
MapAdd {
i: Arg<u32>,
},
PrintExpr,
LoadBuildClass,
UnpackSequence {
size: Arg<u32>,
},
UnpackEx {
args: Arg<UnpackExArgs>,
},
FormatValue {
conversion: Arg<ConversionFlag>,
},
PopException,
Reverse {
amount: Arg<u32>,
},
GetAwaitable,
BeforeAsyncWith,
SetupAsyncWith {
end: Arg<Label>,
},
GetAIter,
GetANext,
EndAsyncFor,
ExtendedArg,
}
const _: () = assert!(mem::size_of::<Instruction>() == 1);
impl From<Instruction> for u8 {
#[inline]
fn from(ins: Instruction) -> u8 {
unsafe { std::mem::transmute::<Instruction, u8>(ins) }
}
}
impl TryFrom<u8> for Instruction {
type Error = crate::marshal::MarshalError;
#[inline]
fn try_from(value: u8) -> Result<Self, crate::marshal::MarshalError> {
if value <= u8::from(Instruction::ExtendedArg) {
Ok(unsafe { std::mem::transmute::<u8, Instruction>(value) })
} else {
Err(crate::marshal::MarshalError::InvalidBytecode)
}
}
}
#[derive(Copy, Clone)]
#[repr(C)]
pub struct CodeUnit {
pub op: Instruction,
pub arg: OpArgByte,
}
const _: () = assert!(mem::size_of::<CodeUnit>() == 2);
impl CodeUnit {
pub fn new(op: Instruction, arg: OpArgByte) -> Self {
Self { op, arg }
}
}
use self::Instruction::*;
bitflags! {
#[derive(Copy, Clone, Debug, PartialEq)]
pub struct MakeFunctionFlags: u8 {
const CLOSURE = 0x01;
const ANNOTATIONS = 0x02;
const KW_ONLY_DEFAULTS = 0x04;
const DEFAULTS = 0x08;
}
}
impl OpArgType for MakeFunctionFlags {
#[inline(always)]
fn from_op_arg(x: u32) -> Option<Self> {
MakeFunctionFlags::from_bits(x as u8)
}
#[inline(always)]
fn to_op_arg(self) -> u32 {
self.bits().into()
}
}
#[derive(Debug, Clone)]
pub enum ConstantData {
Tuple { elements: Vec<ConstantData> },
Integer { value: BigInt },
Float { value: f64 },
Complex { value: Complex64 },
Boolean { value: bool },
Str { value: String },
Bytes { value: Vec<u8> },
Code { code: Box<CodeObject> },
None,
Ellipsis,
}
impl PartialEq for ConstantData {
fn eq(&self, other: &Self) -> bool {
use ConstantData::*;
match (self, other) {
(Integer { value: a }, Integer { value: b }) => a == b,
(Float { value: a }, Float { value: b }) => a.to_bits() == b.to_bits(),
(Complex { value: a }, Complex { value: b }) => {
a.re.to_bits() == b.re.to_bits() && a.im.to_bits() == b.im.to_bits()
}
(Boolean { value: a }, Boolean { value: b }) => a == b,
(Str { value: a }, Str { value: b }) => a == b,
(Bytes { value: a }, Bytes { value: b }) => a == b,
(Code { code: a }, Code { code: b }) => std::ptr::eq(a.as_ref(), b.as_ref()),
(Tuple { elements: a }, Tuple { elements: b }) => a == b,
(None, None) => true,
(Ellipsis, Ellipsis) => true,
_ => false,
}
}
}
impl Eq for ConstantData {}
impl hash::Hash for ConstantData {
fn hash<H: hash::Hasher>(&self, state: &mut H) {
use ConstantData::*;
mem::discriminant(self).hash(state);
match self {
Integer { value } => value.hash(state),
Float { value } => value.to_bits().hash(state),
Complex { value } => {
value.re.to_bits().hash(state);
value.im.to_bits().hash(state);
}
Boolean { value } => value.hash(state),
Str { value } => value.hash(state),
Bytes { value } => value.hash(state),
Code { code } => std::ptr::hash(code.as_ref(), state),
Tuple { elements } => elements.hash(state),
None => {}
Ellipsis => {}
}
}
}
pub enum BorrowedConstant<'a, C: Constant> {
Integer { value: &'a BigInt },
Float { value: f64 },
Complex { value: Complex64 },
Boolean { value: bool },
Str { value: &'a str },
Bytes { value: &'a [u8] },
Code { code: &'a CodeObject<C> },
Tuple { elements: &'a [C] },
None,
Ellipsis,
}
impl<C: Constant> Copy for BorrowedConstant<'_, C> {}
impl<C: Constant> Clone for BorrowedConstant<'_, C> {
fn clone(&self) -> Self {
*self
}
}
impl<C: Constant> BorrowedConstant<'_, C> {
pub fn fmt_display(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
BorrowedConstant::Integer { value } => write!(f, "{value}"),
BorrowedConstant::Float { value } => write!(f, "{value}"),
BorrowedConstant::Complex { value } => write!(f, "{value}"),
BorrowedConstant::Boolean { value } => {
write!(f, "{}", if *value { "True" } else { "False" })
}
BorrowedConstant::Str { value } => write!(f, "{value:?}"),
BorrowedConstant::Bytes { value } => write!(f, "b\"{}\"", value.escape_ascii()),
BorrowedConstant::Code { code } => write!(f, "{code:?}"),
BorrowedConstant::Tuple { elements } => {
write!(f, "(")?;
let mut first = true;
for c in *elements {
if first {
first = false
} else {
write!(f, ", ")?;
}
c.borrow_constant().fmt_display(f)?;
}
write!(f, ")")
}
BorrowedConstant::None => write!(f, "None"),
BorrowedConstant::Ellipsis => write!(f, "..."),
}
}
pub fn to_owned(self) -> ConstantData {
use ConstantData::*;
match self {
BorrowedConstant::Integer { value } => Integer {
value: value.clone(),
},
BorrowedConstant::Float { value } => Float { value },
BorrowedConstant::Complex { value } => Complex { value },
BorrowedConstant::Boolean { value } => Boolean { value },
BorrowedConstant::Str { value } => Str {
value: value.to_owned(),
},
BorrowedConstant::Bytes { value } => Bytes {
value: value.to_owned(),
},
BorrowedConstant::Code { code } => Code {
code: Box::new(code.map_clone_bag(&BasicBag)),
},
BorrowedConstant::Tuple { elements } => Tuple {
elements: elements
.iter()
.map(|c| c.borrow_constant().to_owned())
.collect(),
},
BorrowedConstant::None => None,
BorrowedConstant::Ellipsis => Ellipsis,
}
}
}
op_arg_enum!(
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
#[repr(u8)]
pub enum ComparisonOperator {
Less = 0b001,
Greater = 0b010,
NotEqual = 0b011,
Equal = 0b100,
LessOrEqual = 0b101,
GreaterOrEqual = 0b110,
}
);
op_arg_enum!(
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
#[repr(u8)]
pub enum TestOperator {
In = 0,
NotIn = 1,
Is = 2,
IsNot = 3,
ExceptionMatch = 4,
}
);
op_arg_enum!(
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
#[repr(u8)]
pub enum BinaryOperator {
Power = 0,
Multiply = 1,
MatrixMultiply = 2,
Divide = 3,
FloorDivide = 4,
Modulo = 5,
Add = 6,
Subtract = 7,
Lshift = 8,
Rshift = 9,
And = 10,
Xor = 11,
Or = 12,
}
);
op_arg_enum!(
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
#[repr(u8)]
pub enum UnaryOperator {
Not = 0,
Invert = 1,
Minus = 2,
Plus = 3,
}
);
#[derive(Copy, Clone)]
pub struct UnpackExArgs {
pub before: u8,
pub after: u8,
}
impl OpArgType for UnpackExArgs {
#[inline(always)]
fn from_op_arg(x: u32) -> Option<Self> {
let [before, after, ..] = x.to_le_bytes();
Some(Self { before, after })
}
#[inline(always)]
fn to_op_arg(self) -> u32 {
u32::from_le_bytes([self.before, self.after, 0, 0])
}
}
impl fmt::Display for UnpackExArgs {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "before: {}, after: {}", self.before, self.after)
}
}
pub struct Arguments<'a, N: AsRef<str>> {
pub posonlyargs: &'a [N],
pub args: &'a [N],
pub vararg: Option<&'a N>,
pub kwonlyargs: &'a [N],
pub varkwarg: Option<&'a N>,
}
impl<N: AsRef<str>> fmt::Debug for Arguments<'_, N> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
macro_rules! fmt_slice {
($x:expr) => {
format_args!("[{}]", $x.iter().map(AsRef::as_ref).format(", "))
};
}
f.debug_struct("Arguments")
.field("posonlyargs", &fmt_slice!(self.posonlyargs))
.field("args", &fmt_slice!(self.posonlyargs))
.field("vararg", &self.vararg.map(N::as_ref))
.field("kwonlyargs", &fmt_slice!(self.kwonlyargs))
.field("varkwarg", &self.varkwarg.map(N::as_ref))
.finish()
}
}
impl<C: Constant> CodeObject<C> {
pub fn arg_names(&self) -> Arguments<C::Name> {
let nargs = self.arg_count as usize;
let nkwargs = self.kwonlyarg_count as usize;
let mut varargs_pos = nargs + nkwargs;
let posonlyargs = &self.varnames[..self.posonlyarg_count as usize];
let args = &self.varnames[..nargs];
let kwonlyargs = &self.varnames[nargs..varargs_pos];
let vararg = if self.flags.contains(CodeFlags::HAS_VARARGS) {
let vararg = &self.varnames[varargs_pos];
varargs_pos += 1;
Some(vararg)
} else {
None
};
let varkwarg = if self.flags.contains(CodeFlags::HAS_VARKEYWORDS) {
Some(&self.varnames[varargs_pos])
} else {
None
};
Arguments {
posonlyargs,
args,
vararg,
kwonlyargs,
varkwarg,
}
}
pub fn label_targets(&self) -> BTreeSet<Label> {
let mut label_targets = BTreeSet::new();
let mut arg_state = OpArgState::default();
for instruction in &*self.instructions {
let (instruction, arg) = arg_state.get(*instruction);
if let Some(l) = instruction.label_arg() {
label_targets.insert(l.get(arg));
}
}
label_targets
}
fn display_inner(
&self,
f: &mut fmt::Formatter,
expand_code_objects: bool,
level: usize,
) -> fmt::Result {
let label_targets = self.label_targets();
let line_digits = (3).max(self.locations.last().unwrap().row.to_string().len());
let offset_digits = (4).max(self.instructions.len().to_string().len());
let mut last_line = OneIndexed::MAX;
let mut arg_state = OpArgState::default();
for (offset, &instruction) in self.instructions.iter().enumerate() {
let (instruction, arg) = arg_state.get(instruction);
let line = self.locations[offset].row;
if line != last_line {
if last_line != OneIndexed::MAX {
writeln!(f)?;
}
last_line = line;
write!(f, "{line:line_digits$}")?;
} else {
for _ in 0..line_digits {
write!(f, " ")?;
}
}
write!(f, " ")?;
for _ in 0..level {
write!(f, " ")?;
}
let arrow = if label_targets.contains(&Label(offset as u32)) {
">>"
} else {
" "
};
write!(f, "{arrow} {offset:offset_digits$} ")?;
instruction.fmt_dis(arg, f, self, expand_code_objects, 21, level)?;
writeln!(f)?;
}
Ok(())
}
pub fn display_expand_code_objects(&self) -> impl fmt::Display + '_ {
struct Display<'a, C: Constant>(&'a CodeObject<C>);
impl<C: Constant> fmt::Display for Display<'_, C> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
self.0.display_inner(f, true, 1)
}
}
Display(self)
}
pub fn map_bag<Bag: ConstantBag>(self, bag: Bag) -> CodeObject<Bag::Constant> {
let map_names = |names: Box<[C::Name]>| {
names
.into_vec()
.into_iter()
.map(|x| bag.make_name(x.as_ref()))
.collect::<Box<[_]>>()
};
CodeObject {
constants: self
.constants
.into_vec()
.into_iter()
.map(|x| bag.make_constant(x.borrow_constant()))
.collect(),
names: map_names(self.names),
varnames: map_names(self.varnames),
cellvars: map_names(self.cellvars),
freevars: map_names(self.freevars),
source_path: bag.make_name(self.source_path.as_ref()),
obj_name: bag.make_name(self.obj_name.as_ref()),
instructions: self.instructions,
locations: self.locations,
flags: self.flags,
posonlyarg_count: self.posonlyarg_count,
arg_count: self.arg_count,
kwonlyarg_count: self.kwonlyarg_count,
first_line_number: self.first_line_number,
max_stackdepth: self.max_stackdepth,
cell2arg: self.cell2arg,
}
}
pub fn map_clone_bag<Bag: ConstantBag>(&self, bag: &Bag) -> CodeObject<Bag::Constant> {
let map_names =
|names: &[C::Name]| names.iter().map(|x| bag.make_name(x.as_ref())).collect();
CodeObject {
constants: self
.constants
.iter()
.map(|x| bag.make_constant(x.borrow_constant()))
.collect(),
names: map_names(&self.names),
varnames: map_names(&self.varnames),
cellvars: map_names(&self.cellvars),
freevars: map_names(&self.freevars),
source_path: bag.make_name(self.source_path.as_ref()),
obj_name: bag.make_name(self.obj_name.as_ref()),
instructions: self.instructions.clone(),
locations: self.locations.clone(),
flags: self.flags,
posonlyarg_count: self.posonlyarg_count,
arg_count: self.arg_count,
kwonlyarg_count: self.kwonlyarg_count,
first_line_number: self.first_line_number,
max_stackdepth: self.max_stackdepth,
cell2arg: self.cell2arg.clone(),
}
}
}
impl<C: Constant> fmt::Display for CodeObject<C> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
self.display_inner(f, false, 1)?;
for constant in &*self.constants {
if let BorrowedConstant::Code { code } = constant.borrow_constant() {
writeln!(f, "\nDisassembly of {code:?}")?;
code.fmt(f)?;
}
}
Ok(())
}
}
impl Instruction {
#[inline]
pub fn label_arg(&self) -> Option<Arg<Label>> {
match self {
Jump { target: l }
| JumpIfTrue { target: l }
| JumpIfFalse { target: l }
| JumpIfTrueOrPop { target: l }
| JumpIfFalseOrPop { target: l }
| ForIter { target: l }
| SetupFinally { handler: l }
| SetupExcept { handler: l }
| SetupWith { end: l }
| SetupAsyncWith { end: l }
| Break { target: l }
| Continue { target: l } => Some(*l),
_ => None,
}
}
pub fn unconditional_branch(&self) -> bool {
matches!(
self,
Jump { .. } | Continue { .. } | Break { .. } | ReturnValue | Raise { .. }
)
}
pub fn stack_effect(&self, arg: OpArg, jump: bool) -> i32 {
match self {
ImportName { .. } | ImportNameless => -1,
ImportStar => -1,
ImportFrom { .. } => 1,
LoadFast(_) | LoadNameAny(_) | LoadGlobal(_) | LoadDeref(_) | LoadClassDeref(_) => 1,
StoreFast(_) | StoreLocal(_) | StoreGlobal(_) | StoreDeref(_) => -1,
DeleteFast(_) | DeleteLocal(_) | DeleteGlobal(_) | DeleteDeref(_) => 0,
LoadClosure(_) => 1,
Subscript => -1,
StoreSubscript => -3,
DeleteSubscript => -2,
LoadAttr { .. } => 0,
StoreAttr { .. } => -2,
DeleteAttr { .. } => -1,
LoadConst { .. } => 1,
UnaryOperation { .. } => 0,
BinaryOperation { .. }
| BinaryOperationInplace { .. }
| TestOperation { .. }
| CompareOperation { .. } => -1,
Pop => -1,
Rotate2 | Rotate3 => 0,
Duplicate => 1,
Duplicate2 => 2,
GetIter => 0,
Continue { .. } => 0,
Break { .. } => 0,
Jump { .. } => 0,
JumpIfTrue { .. } | JumpIfFalse { .. } => -1,
JumpIfTrueOrPop { .. } | JumpIfFalseOrPop { .. } => {
if jump {
0
} else {
-1
}
}
MakeFunction(flags) => {
let flags = flags.get(arg);
-2 - flags.contains(MakeFunctionFlags::CLOSURE) as i32
- flags.contains(MakeFunctionFlags::ANNOTATIONS) as i32
- flags.contains(MakeFunctionFlags::KW_ONLY_DEFAULTS) as i32
- flags.contains(MakeFunctionFlags::DEFAULTS) as i32
+ 1
}
CallFunctionPositional { nargs } => -(nargs.get(arg) as i32) - 1 + 1,
CallMethodPositional { nargs } => -(nargs.get(arg) as i32) - 3 + 1,
CallFunctionKeyword { nargs } => -1 - (nargs.get(arg) as i32) - 1 + 1,
CallMethodKeyword { nargs } => -1 - (nargs.get(arg) as i32) - 3 + 1,
CallFunctionEx { has_kwargs } => -1 - (has_kwargs.get(arg) as i32) - 1 + 1,
CallMethodEx { has_kwargs } => -1 - (has_kwargs.get(arg) as i32) - 3 + 1,
LoadMethod { .. } => -1 + 3,
ForIter { .. } => {
if jump {
-1
} else {
1
}
}
ReturnValue => -1,
YieldValue => 0,
YieldFrom => -1,
SetupAnnotation | SetupLoop | SetupFinally { .. } | EnterFinally | EndFinally => 0,
SetupExcept { .. } => jump as i32,
SetupWith { .. } => (!jump) as i32,
WithCleanupStart => 0,
WithCleanupFinish => -1,
PopBlock => 0,
Raise { kind } => -(kind.get(arg) as u8 as i32),
BuildString { size }
| BuildTuple { size, .. }
| BuildTupleUnpack { size, .. }
| BuildList { size, .. }
| BuildListUnpack { size, .. }
| BuildSet { size, .. }
| BuildSetUnpack { size, .. } => -(size.get(arg) as i32) + 1,
BuildMap { size } => {
let nargs = size.get(arg) * 2;
-(nargs as i32) + 1
}
BuildMapForCall { size } => {
let nargs = size.get(arg);
-(nargs as i32) + 1
}
DictUpdate => -1,
BuildSlice { step } => -2 - (step.get(arg) as i32) + 1,
ListAppend { .. } | SetAdd { .. } => -1,
MapAdd { .. } => -2,
PrintExpr => -1,
LoadBuildClass => 1,
UnpackSequence { size } => -1 + size.get(arg) as i32,
UnpackEx { args } => {
let UnpackExArgs { before, after } = args.get(arg);
-1 + before as i32 + 1 + after as i32
}
FormatValue { .. } => -1,
PopException => 0,
Reverse { .. } => 0,
GetAwaitable => 0,
BeforeAsyncWith => 1,
SetupAsyncWith { .. } => {
if jump {
-1
} else {
0
}
}
GetAIter => 0,
GetANext => 1,
EndAsyncFor => -2,
ExtendedArg => 0,
}
}
pub fn display<'a>(
&'a self,
arg: OpArg,
ctx: &'a impl InstrDisplayContext,
) -> impl fmt::Display + 'a {
struct FmtFn<F>(F);
impl<F: Fn(&mut fmt::Formatter) -> fmt::Result> fmt::Display for FmtFn<F> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
(self.0)(f)
}
}
FmtFn(move |f: &mut fmt::Formatter| self.fmt_dis(arg, f, ctx, false, 0, 0))
}
#[allow(clippy::too_many_arguments)]
fn fmt_dis(
&self,
arg: OpArg,
f: &mut fmt::Formatter,
ctx: &impl InstrDisplayContext,
expand_code_objects: bool,
pad: usize,
level: usize,
) -> fmt::Result {
macro_rules! w {
($variant:ident) => {
write!(f, stringify!($variant))
};
($variant:ident, $map:ident = $arg_marker:expr) => {{
let arg = $arg_marker.get(arg);
write!(f, "{:pad$}({}, {})", stringify!($variant), arg, $map(arg))
}};
($variant:ident, $arg_marker:expr) => {
write!(f, "{:pad$}({})", stringify!($variant), $arg_marker.get(arg))
};
($variant:ident, ?$arg_marker:expr) => {
write!(
f,
"{:pad$}({:?})",
stringify!($variant),
$arg_marker.get(arg)
)
};
}
let varname = |i: u32| ctx.get_varname(i as usize);
let name = |i: u32| ctx.get_name(i as usize);
let cell_name = |i: u32| ctx.get_cell_name(i as usize);
match self {
ImportName { idx } => w!(ImportName, name = idx),
ImportNameless => w!(ImportNameless),
ImportStar => w!(ImportStar),
ImportFrom { idx } => w!(ImportFrom, name = idx),
LoadFast(idx) => w!(LoadFast, varname = idx),
LoadNameAny(idx) => w!(LoadNameAny, name = idx),
LoadGlobal(idx) => w!(LoadGlobal, name = idx),
LoadDeref(idx) => w!(LoadDeref, cell_name = idx),
LoadClassDeref(idx) => w!(LoadClassDeref, cell_name = idx),
StoreFast(idx) => w!(StoreFast, varname = idx),
StoreLocal(idx) => w!(StoreLocal, name = idx),
StoreGlobal(idx) => w!(StoreGlobal, name = idx),
StoreDeref(idx) => w!(StoreDeref, cell_name = idx),
DeleteFast(idx) => w!(DeleteFast, varname = idx),
DeleteLocal(idx) => w!(DeleteLocal, name = idx),
DeleteGlobal(idx) => w!(DeleteGlobal, name = idx),
DeleteDeref(idx) => w!(DeleteDeref, cell_name = idx),
LoadClosure(i) => w!(LoadClosure, cell_name = i),
Subscript => w!(Subscript),
StoreSubscript => w!(StoreSubscript),
DeleteSubscript => w!(DeleteSubscript),
StoreAttr { idx } => w!(StoreAttr, name = idx),
DeleteAttr { idx } => w!(DeleteAttr, name = idx),
LoadConst { idx } => {
let value = ctx.get_constant(idx.get(arg) as usize);
match value.borrow_constant() {
BorrowedConstant::Code { code } if expand_code_objects => {
write!(f, "{:pad$}({:?}):", "LoadConst", code)?;
code.display_inner(f, true, level + 1)?;
Ok(())
}
c => {
write!(f, "{:pad$}(", "LoadConst")?;
c.fmt_display(f)?;
write!(f, ")")
}
}
}
UnaryOperation { op } => w!(UnaryOperation, ?op),
BinaryOperation { op } => w!(BinaryOperation, ?op),
BinaryOperationInplace { op } => w!(BinaryOperationInplace, ?op),
LoadAttr { idx } => w!(LoadAttr, name = idx),
TestOperation { op } => w!(TestOperation, ?op),
CompareOperation { op } => w!(CompareOperation, ?op),
Pop => w!(Pop),
Rotate2 => w!(Rotate2),
Rotate3 => w!(Rotate3),
Duplicate => w!(Duplicate),
Duplicate2 => w!(Duplicate2),
GetIter => w!(GetIter),
Continue { target } => w!(Continue, target),
Break { target } => w!(Break, target),
Jump { target } => w!(Jump, target),
JumpIfTrue { target } => w!(JumpIfTrue, target),
JumpIfFalse { target } => w!(JumpIfFalse, target),
JumpIfTrueOrPop { target } => w!(JumpIfTrueOrPop, target),
JumpIfFalseOrPop { target } => w!(JumpIfFalseOrPop, target),
MakeFunction(flags) => w!(MakeFunction, ?flags),
CallFunctionPositional { nargs } => w!(CallFunctionPositional, nargs),
CallFunctionKeyword { nargs } => w!(CallFunctionKeyword, nargs),
CallFunctionEx { has_kwargs } => w!(CallFunctionEx, has_kwargs),
LoadMethod { idx } => w!(LoadMethod, name = idx),
CallMethodPositional { nargs } => w!(CallMethodPositional, nargs),
CallMethodKeyword { nargs } => w!(CallMethodKeyword, nargs),
CallMethodEx { has_kwargs } => w!(CallMethodEx, has_kwargs),
ForIter { target } => w!(ForIter, target),
ReturnValue => w!(ReturnValue),
YieldValue => w!(YieldValue),
YieldFrom => w!(YieldFrom),
SetupAnnotation => w!(SetupAnnotation),
SetupLoop => w!(SetupLoop),
SetupExcept { handler } => w!(SetupExcept, handler),
SetupFinally { handler } => w!(SetupFinally, handler),
EnterFinally => w!(EnterFinally),
EndFinally => w!(EndFinally),
SetupWith { end } => w!(SetupWith, end),
WithCleanupStart => w!(WithCleanupStart),
WithCleanupFinish => w!(WithCleanupFinish),
BeforeAsyncWith => w!(BeforeAsyncWith),
SetupAsyncWith { end } => w!(SetupAsyncWith, end),
PopBlock => w!(PopBlock),
Raise { kind } => w!(Raise, ?kind),
BuildString { size } => w!(BuildString, size),
BuildTuple { size } => w!(BuildTuple, size),
BuildTupleUnpack { size } => w!(BuildTupleUnpack, size),
BuildList { size } => w!(BuildList, size),
BuildListUnpack { size } => w!(BuildListUnpack, size),
BuildSet { size } => w!(BuildSet, size),
BuildSetUnpack { size } => w!(BuildSetUnpack, size),
BuildMap { size } => w!(BuildMap, size),
BuildMapForCall { size } => w!(BuildMap, size),
DictUpdate => w!(DictUpdate),
BuildSlice { step } => w!(BuildSlice, step),
ListAppend { i } => w!(ListAppend, i),
SetAdd { i } => w!(SetAdd, i),
MapAdd { i } => w!(MapAdd, i),
PrintExpr => w!(PrintExpr),
LoadBuildClass => w!(LoadBuildClass),
UnpackSequence { size } => w!(UnpackSequence, size),
UnpackEx { args } => w!(UnpackEx, args),
FormatValue { conversion } => w!(FormatValue, ?conversion),
PopException => w!(PopException),
Reverse { amount } => w!(Reverse, amount),
GetAwaitable => w!(GetAwaitable),
GetAIter => w!(GetAIter),
GetANext => w!(GetANext),
EndAsyncFor => w!(EndAsyncFor),
ExtendedArg => w!(ExtendedArg, Arg::<u32>::marker()),
}
}
}
pub trait InstrDisplayContext {
type Constant: Constant;
fn get_constant(&self, i: usize) -> &Self::Constant;
fn get_name(&self, i: usize) -> &str;
fn get_varname(&self, i: usize) -> &str;
fn get_cell_name(&self, i: usize) -> &str;
}
impl<C: Constant> InstrDisplayContext for CodeObject<C> {
type Constant = C;
fn get_constant(&self, i: usize) -> &C {
&self.constants[i]
}
fn get_name(&self, i: usize) -> &str {
self.names[i].as_ref()
}
fn get_varname(&self, i: usize) -> &str {
self.varnames[i].as_ref()
}
fn get_cell_name(&self, i: usize) -> &str {
self.cellvars
.get(i)
.unwrap_or_else(|| &self.freevars[i - self.cellvars.len()])
.as_ref()
}
}
impl fmt::Display for ConstantData {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
self.borrow_constant().fmt_display(f)
}
}
impl<C: Constant> fmt::Debug for CodeObject<C> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(
f,
"<code object {} at ??? file {:?}, line {}>",
self.obj_name.as_ref(),
self.source_path.as_ref(),
self.first_line_number.map_or(-1, |x| x.get() as i32)
)
}
}