use crate::vm::{

AsObject, PyObjectRef, PyResult, VirtualMachine,

bytecode::{AnyOpcode, oparg},

};

fn try_from_i32(raw: i32) -> Result<AnyOpcode, ()> {

u16::try_from(raw)

.map_err(|_| ())?

.try_into()

.map_err(|_| ())

}

// prepare specialization

#[pyattr]

const ENABLE_SPECIALIZATION: i8 = 1;

#[pyattr]

const ENABLE_SPECIALIZATION_FT: i8 = 1;

#[derive(FromArgs)]

struct StackEffectArgs {

#[pyarg(positional)]

opcode: i32,

#[pyarg(positional, optional)]

oparg: Option<PyObjectRef>,

#[pyarg(named, optional)]

jump: Option<PyObjectRef>,

}

#[pyfunction]

#[allow(clippy::unnecessary_cast)] // `c_long` is `i32` or `i64`

fn stack_effect(args: StackEffectArgs, vm: &VirtualMachine) -> PyResult<i32> {

let oparg = match args.oparg {

// (int)PyLong_AsLong(oparg)

Some(v) if !vm.is_none(&v) => v.try_into_value::<core::ffi::c_long>(vm)? as i32,

_ => 0,

};

let jump = match args.jump {

None => None,

Some(v) if vm.is_none(&v) => None,

Some(v) if v.is(&vm.ctx.true_value) => Some(true),

Some(v) if v.is(&vm.ctx.false_value) => Some(false),

Some(_) => {

return Err(vm.new_value_error("stack_effect: jump must be False, True or None"));

}

};

let invalid = || vm.new_value_error("invalid opcode or oparg");

let instr = try_from_i32(args.opcode).map_err(|_| invalid())?;

// Specialized instructions are not supported.

if instr.real().is_some_and(|op| op.deopt().is_some()) {

return Err(invalid());

}

let oparg = u32::try_from(oparg).map_err(|_| invalid())?;

let effect = if jump == Some(false) {

instr.stack_effect(oparg)

} else {

instr.stack_effect_jump(oparg)

};

Ok(effect)

}

#[derive(FromArgs)]

struct OpcodeArg {

#[pyarg(any)]

opcode: i32,

}

#[pyfunction]

fn is_valid(OpcodeArg { opcode }: OpcodeArg) -> bool {

try_from_i32(opcode).is_ok()

}

#[pyfunction]

fn has_arg(OpcodeArg { opcode }: OpcodeArg) -> bool {

try_from_i32(opcode).is_ok_and(|op| op.has_arg())

}

#[pyfunction]

fn has_const(OpcodeArg { opcode }: OpcodeArg) -> bool {

try_from_i32(opcode).is_ok_and(|op| op.has_const())

}

#[pyfunction]

fn has_name(OpcodeArg { opcode }: OpcodeArg) -> bool {

try_from_i32(opcode).is_ok_and(|op| op.has_name())

}

#[pyfunction]

fn has_jump(OpcodeArg { opcode }: OpcodeArg) -> bool {

try_from_i32(opcode).is_ok_and(|op| op.has_jump())

}

#[pyfunction]

fn has_free(OpcodeArg { opcode }: OpcodeArg) -> bool {

try_from_i32(opcode).is_ok_and(|op| op.has_free())

}

#[pyfunction]

fn has_local(OpcodeArg { opcode }: OpcodeArg) -> bool {

try_from_i32(opcode).is_ok_and(|op| op.has_local())

}

#[pyfunction]

fn has_exc(OpcodeArg { opcode }: OpcodeArg) -> bool {

try_from_i32(opcode).is_ok_and(|op| op.is_block_push())

}

#[pyfunction]

fn get_intrinsic1_descs(vm: &VirtualMachine) -> Vec<PyObjectRef> {

oparg::IntrinsicFunction1::iter()

.map(|x| vm.ctx.new_str(x.desc()).into())

.collect()

}

#[pyfunction]

fn get_intrinsic2_descs(vm: &VirtualMachine) -> Vec<PyObjectRef> {

oparg::IntrinsicFunction2::iter()

.map(|x| vm.ctx.new_str(x.desc()).into())

.collect()

}

#[pyfunction]

fn get_nb_ops(vm: &VirtualMachine) -> Vec<PyObjectRef> {

oparg::BinaryOperator::iter()

.map(|x| {

vm.ctx

.new_tuple(vec![

vm.ctx.new_str(x.desc()).into(),

vm.ctx.new_str(x.to_string()).into(),

])

.into()

})

.collect()

}

#[pyfunction]

fn get_special_method_names(vm: &VirtualMachine) -> Vec<PyObjectRef> {

oparg::SpecialMethod::iter()

.map(|x| vm.ctx.new_str(x.to_string()).into())

.collect()

}

#[derive(FromArgs)]

#[allow(dead_code)]

struct ExecutorArgs {

#[pyarg(any)]

code: PyObjectRef,

#[pyarg(any)]

offset: i32,

}

#[pyfunction]

fn get_executor(_args: ExecutorArgs, vm: &VirtualMachine) -> PyObjectRef {

// TODO

vm.ctx.none()

}

#[pyfunction]

fn get_specialization_stats(vm: &VirtualMachine) -> PyObjectRef {

vm.ctx.none()

}

use crate::vm::{self, compiler::Mode};

macro_rules! assert_dis_snapshot {

($value:expr) => {

insta::assert_snapshot!(dis($value, false))

};

(positions: $value:expr) => {

insta::assert_snapshot!(dis($value, true))

};

}

/// Returns the [`dis.dis`](https://docs.python.org/3/library/dis.html#dis.dis) output,

/// with the source positions of the instructions if `show_positions` is true.

///

/// # Notes

///

/// Memory addresses in the output are replaced with `0xdeadbeef` for consistency.

fn dis(source: &str, show_positions: bool) -> String {

const FNAME: &str = "<?>";

let builder = vm::Interpreter::builder(Default::default());

let stdlib_defs = crate::stdlib_module_defs(&builder.ctx);

let interp = builder

.add_native_modules(&stdlib_defs)

.add_frozen_modules(rustpython_pylib::FROZEN_STDLIB)

.build();

interp.enter(|vm| {

let scope = vm.new_scope_with_builtins();

let code_obj = vm

.compile(source.trim(), Mode::Exec, FNAME)

.map_err(|err| err.into_pyexception(vm, Some(source)))

.unwrap();

scope.globals.set_item("code", code_obj.into(), vm).unwrap();

scope

.globals

.set_item("show_positions", vm.new_pyobj(show_positions), vm)

.unwrap();

let py_source = r#"

let py_code_obj = vm

.compile(py_source, Mode::Exec, FNAME)

.map_err(|err| err.into_pyexception(vm, Some(py_source)))

.unwrap();

vm.run_code_obj(py_code_obj, scope.clone())

.unwrap_or_else(|err| {

let mut traceback = String::new();

let _ = vm.write_exception(&mut traceback, &err);

panic!("Python disassembly failed: {err:?}\n{traceback}");

});

let py_output = scope.globals.get_item("output", vm).unwrap();

py_output.str(vm).unwrap().to_string()

})

}

#[test]

fn if_ors() {

assert_dis_snapshot!(

r#"

)

}

#[test]

fn if_ands() {

assert_dis_snapshot!(

r#"

)

}

#[test]

fn if_mixed() {

assert_dis_snapshot!(

r#"

)

}

#[test]

fn nested_bool_op() {

assert_dis_snapshot!(

r#"

)

}

#[test]

fn const_no_op() {

assert_dis_snapshot!(

r#"

)

}

#[test]

fn constant_true_if_pass_keeps_line_anchor_nop() {

assert_dis_snapshot!(

r#"

)

}

#[test]

fn nested_double_async_with() {

assert_dis_snapshot!(

r#"

)

}

#[test]

fn bare_function_annotations_check_attribute_and_subscript_expressions() {

assert_dis_snapshot!(

r#"

)

}

#[test]

fn empty_module_implicit_return_location() {

// The implicit return is emitted with NO_LOCATION; propagate_line_numbers()

// gives it the module RESUME location, whose line is 0.

assert_dis_snapshot!(positions: "")

}

#[test]

fn module_docstring_location() {

// codegen_body() loads the docstring at its expression location and stores

// __doc__ with NO_LOCATION.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn redundant_nop_location() {

// basicblock_remove_redundant_nops() copies the full NOP location into a

// following no-location jump.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn annotated_multiline_function_firstlineno() {

// codegen_function() computes firstlineno from the FunctionDef before

// compiling annotations.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn annotation_scope_function_location() {

// codegen_function_annotations() receives LOC(function) and uses it for the

// annotation closure's BUILD_MAP/RETURN_VALUE and the MAKE_FUNCTION annotate

// sequence.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn starred_arg_annotation_unpack_location() {

// codegen_argannotation() visits `Ts` at the annotation location, then emits

// UNPACK_SEQUENCE at LOC(function).

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn module_deferred_annotations_location() {

// codegen_process_deferred_annotations() uses the first module statement

// location from start_location().

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn not_compare_location() {

// A Compare inside UnaryOp(Not) is parsed with the UnaryOp start location.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn not_chained_compare_location() {

// A chained comparison keeps its inner range for the comparison scaffolding

// and uses the UnaryOp range only for TO_BOOL and UNARY_NOT.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn type_param_scope_locations() {

// codegen_type_params() emits type parameter ops at LOC(typeparam),

// bound/default evaluators at LOC(e), and type alias plumbing at LOC(s).

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn generic_typealias_wrapper_return_location() {

// The generic parameters wrapper of a type alias returns at LOC(type alias).

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn generic_function_annotation_scope_location() {

// codegen_function_annotations() gets LOC(function) inside the generic

// parameters scope too.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn generic_class_type_params_store_location() {

// codegen_class() stores the .type_params cell at LOC(class).

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn class_deferred_annotations_location() {

// Deferred annotation setup and the class tail inherit the class body

// location LOCATION(firstlineno, firstlineno, 0, 0).

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn future_annotation_string_location() {

// The stringized annotation is loaded at LOC(annotation) and stored at

// LOC(AnnAssign).

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn lambda_dict_literal_location() {

// codegen_dict() and codegen_subdict() use LOC(dict) for BUILD_MAP, MAP_ADD

// and DICT_UPDATE.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn match_or_alternative_locations() {

// codegen_pattern_or() emits the COPY of each alternative at LOC(alt).

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn match_success_jump_location() {

// codegen_match_inner() emits the success jump with NO_LOCATION.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn match_mapping_keys_location() {

// codegen_pattern_mapping() returns to LOC(p) for the BUILD_TUPLE/MATCH_KEYS

// scaffolding.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn match_class_scaffolding_location() {

// codegen_pattern_class() returns to LOC(p) after visiting the class.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn taken_constant_boolop_location() {

// The surviving load of a short-circuited BoolOp keeps the literal range.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn assert_false_message_location() {

// codegen_assert() emits LOAD_COMMON_CONSTANT and CALL at LOC(assert) and

// RAISE_VARARGS at LOC(test).

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn static_swap_implicit_return_location() {

// Line numbers are resolved before apply_static_swaps() reorders the stores,

// so the return epilogue keeps the location of the last store before the swap.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn genexpr_closure_and_call_location() {

// codegen_comprehension() uses LOC(e) for the closure, the outer CALL and the

// implicit .0 load.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn implicit_call_genexpr_range() {

// A sole unparenthesized generator expression argument spans the call

// parentheses.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn implicit_call_genexpr_parenthesized_element_range() {

// The call parentheses are part of the generator expression range even when

// its element starts with a parenthesis.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn genexpr_filter_cleanup_location() {

// jump_thread() copies the if_cleanup backedge location, elt_loc, to the

// threaded cleanup path.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn try_finally_exception_scaffolding_location() {

// codegen_try_finally() emits the exception path scaffolding with NO_LOCATION.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn adjacent_no_location_entries() {

// assemble_location_info() merges adjacent NO_LOCATION instructions into one

// linetable entry.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn fstring_format_location() {

// codegen_formatted_value() emits CONVERT_VALUE and FORMAT_* at

// LOC(FormattedValue).

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn joined_string_literals_location() {

// Adjacent f-string literal fragments are one Constant spanning all of them.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn decorated_definition_locations() {

// Decorators are evaluated first, the closure and the store use LOC(s), and

// each decorator CALL uses the decorator location.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn builtin_any_genexpr_call() {

// any(genexpr) is inlined behind a builtin identity check that falls back to

// the ordinary call.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn builtin_tuple_list_set_genexpr_calls() {

// tuple(), list() and set() of a generator expression are inlined behind a

// builtin identity check that falls back to the ordinary call.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn augassign_constant_slice_location() {

// codegen_augassign() emits the COPY/COPY/BINARY_OP NB_SUBSCR at LOC(target).

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn comprehension_list_iterable_location() {

// codegen_comprehension_iter() emits GET_ITER at LOC(comp->iter).

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn inlined_comprehension_backedge_location() {

// codegen_sync_comprehension_generator() emits the backedges at elt_loc.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn inlined_async_comprehension_end_async_for_location() {

// codegen_async_comprehension_generator() emits END_ASYNC_FOR at the

// comprehension location.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn async_for_anext_location() {

// codegen_async_for() emits the GET_ANEXT/yield-from scaffolding at LOC(s).

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn namedexpr_copy_location() {

// The COPY of a named expression is at its location, between the value and

// the target.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn with_cleanup_excludes_not_taken() {

// The NOT_TAKEN after the with body is outside the exception table entries.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn async_for_end_async_for() {

// END_ASYNC_FOR jumps back to the END_SEND of its awaitable loop.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn genexpr_prefix() {

// The outermost iterator is acquired before RETURN_GENERATOR, outside the

// StopIteration handler.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn yield_from_and_await() {

// GET_ITER and GET_AWAITABLE leave the iterator pair that SEND uses.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn comprehension_collection_depth() {

// The collection depth counts the slots of every nested iterator.

assert_dis_snapshot!(

positions: r#"

)

}

#[test]

fn import_flags() {

// The IMPORT_NAME oparg flags depend on the scope and the enclosing handlers.

assert_dis_snapshot!(

positions: r#"

)