Skip to content

Commit db8e648

Browse files
committed
Make PyFunction.code a PyCodeRef, PyGenerator.frame a FrameRef, and
other improvements to increase use of specific ref types.
1 parent dc68101 commit db8e648

File tree

9 files changed

+107
-181
lines changed

9 files changed

+107
-181
lines changed

vm/src/frame.rs

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ use crate::bytecode;
1111
use crate::function::PyFuncArgs;
1212
use crate::obj::objbool;
1313
use crate::obj::objbuiltinfunc::PyBuiltinFunction;
14-
use crate::obj::objcode;
14+
use crate::obj::objcode::PyCodeRef;
1515
use crate::obj::objdict;
1616
use crate::obj::objdict::PyDict;
1717
use crate::obj::objint::PyInt;
@@ -22,7 +22,7 @@ use crate::obj::objstr;
2222
use crate::obj::objtype;
2323
use crate::obj::objtype::PyClassRef;
2424
use crate::pyobject::{
25-
DictProtocol, IdProtocol, PyContext, PyObjectRef, PyResult, PyValue, TryFromObject,
25+
DictProtocol, IdProtocol, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
2626
TypeProtocol,
2727
};
2828
use crate::vm::VirtualMachine;
@@ -186,6 +186,8 @@ enum BlockType {
186186
},
187187
}
188188

189+
pub type FrameRef = PyRef<Frame>;
190+
189191
pub struct Frame {
190192
pub code: bytecode::CodeObject,
191193
// We need 1 stack per frame
@@ -211,7 +213,7 @@ pub enum ExecutionResult {
211213
pub type FrameResult = Result<Option<ExecutionResult>, PyObjectRef>;
212214

213215
impl Frame {
214-
pub fn new(code: PyObjectRef, scope: Scope) -> Frame {
216+
pub fn new(code: PyCodeRef, scope: Scope) -> Frame {
215217
//populate the globals and locals
216218
//TODO: This is wrong, check https://github.com/nedbat/byterun/blob/31e6c4a8212c35b5157919abff43a7daa0f377c6/byterun/pyvm2.py#L95
217219
/*
@@ -224,7 +226,7 @@ impl Frame {
224226
// locals.extend(callargs);
225227

226228
Frame {
227-
code: objcode::get_value(&code),
229+
code: code.code.clone(),
228230
stack: RefCell::new(vec![]),
229231
blocks: RefCell::new(vec![]),
230232
// save the callargs as locals
@@ -562,7 +564,7 @@ impl Frame {
562564
}
563565
bytecode::Instruction::MakeFunction { flags } => {
564566
let _qualified_name = self.pop_value();
565-
let code_obj = self.pop_value();
567+
let code_obj = PyCodeRef::try_from_object(vm, self.pop_value())?;
566568

567569
let annotations = if flags.contains(bytecode::FunctionOpArg::HAS_ANNOTATIONS) {
568570
self.pop_value()

vm/src/obj/objcode.rs

Lines changed: 27 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ use crate::vm::VirtualMachine;
1313
pub type PyCodeRef = PyRef<PyCode>;
1414

1515
pub struct PyCode {
16-
code: bytecode::CodeObject,
16+
pub code: bytecode::CodeObject,
1717
}
1818

1919
impl PyCode {
@@ -38,87 +38,58 @@ pub fn init(context: &PyContext) {
3838
let code_type = context.code_type.as_object();
3939
extend_class!(context, code_type, {
4040
"__new__" => context.new_rustfunc(code_new),
41-
"__repr__" => context.new_rustfunc(code_repr)
41+
"__repr__" => context.new_rustfunc(code_repr),
42+
43+
"co_argcount" => context.new_property(code_co_argcount),
44+
"co_consts" => context.new_property(code_co_consts),
45+
"co_filename" => context.new_property(code_co_filename),
46+
"co_firstlineno" => context.new_property(code_co_firstlineno),
47+
"co_kwonlyargcount" => context.new_property(code_co_kwonlyargcount),
48+
"co_name" => context.new_property(code_co_name),
4249
});
43-
44-
for (name, f) in &[
45-
(
46-
"co_argcount",
47-
code_co_argcount as fn(&VirtualMachine, PyFuncArgs) -> PyResult,
48-
),
49-
("co_consts", code_co_consts),
50-
("co_filename", code_co_filename),
51-
("co_firstlineno", code_co_firstlineno),
52-
("co_kwonlyargcount", code_co_kwonlyargcount),
53-
("co_name", code_co_name),
54-
] {
55-
context.set_attr(code_type, name, context.new_property(f))
56-
}
57-
}
58-
59-
pub fn get_value(obj: &PyObjectRef) -> bytecode::CodeObject {
60-
if let Some(code) = obj.payload::<PyCode>() {
61-
code.code.clone()
62-
} else {
63-
panic!("Inner error getting code {:?}", obj)
64-
}
6550
}
6651

6752
fn code_new(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
6853
arg_check!(vm, args, required = [(_cls, None)]);
6954
Err(vm.new_type_error("Cannot directly create code object".to_string()))
7055
}
7156

72-
fn code_repr(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
73-
arg_check!(vm, args, required = [(o, Some(vm.ctx.code_type()))]);
74-
75-
let code = get_value(o);
76-
let repr = format!(
57+
fn code_repr(o: PyCodeRef, _vm: &VirtualMachine) -> String {
58+
let code = &o.code;
59+
format!(
7760
"<code object {} at 0x{:x} file {:?}, line {}>",
7861
code.obj_name,
7962
o.get_id(),
8063
code.source_path,
8164
code.first_line_number
82-
);
83-
Ok(vm.new_str(repr))
84-
}
85-
86-
fn member_code_obj(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult<bytecode::CodeObject> {
87-
arg_check!(vm, args, required = [(zelf, Some(vm.ctx.code_type()))]);
88-
Ok(get_value(zelf))
65+
)
8966
}
9067

91-
fn code_co_argcount(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
92-
let code_obj = member_code_obj(vm, args)?;
93-
Ok(vm.ctx.new_int(code_obj.arg_names.len()))
68+
fn code_co_argcount(code: PyCodeRef, _vm: &VirtualMachine) -> usize {
69+
code.code.arg_names.len()
9470
}
9571

96-
fn code_co_filename(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
97-
let code_obj = member_code_obj(vm, args)?;
98-
let source_path = code_obj.source_path;
99-
Ok(vm.new_str(source_path))
72+
fn code_co_filename(code: PyCodeRef, _vm: &VirtualMachine) -> String {
73+
code.code.source_path.clone()
10074
}
10175

102-
fn code_co_firstlineno(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
103-
let code_obj = member_code_obj(vm, args)?;
104-
Ok(vm.ctx.new_int(code_obj.first_line_number))
76+
fn code_co_firstlineno(code: PyCodeRef, _vm: &VirtualMachine) -> usize {
77+
code.code.first_line_number
10578
}
10679

107-
fn code_co_kwonlyargcount(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
108-
let code_obj = member_code_obj(vm, args)?;
109-
Ok(vm.ctx.new_int(code_obj.kwonlyarg_names.len()))
80+
fn code_co_kwonlyargcount(code: PyCodeRef, _vm: &VirtualMachine) -> usize {
81+
code.code.kwonlyarg_names.len()
11082
}
11183

112-
fn code_co_consts(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
113-
let code_obj = member_code_obj(vm, args)?;
114-
let consts = code_obj
84+
fn code_co_consts(code: PyCodeRef, vm: &VirtualMachine) -> PyObjectRef {
85+
let consts = code
86+
.code
11587
.get_constants()
11688
.map(|x| vm.ctx.unwrap_constant(x))
11789
.collect();
118-
Ok(vm.ctx.new_tuple(consts))
90+
vm.ctx.new_tuple(consts)
11991
}
12092

121-
fn code_co_name(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
122-
let code_obj = member_code_obj(vm, args)?;
123-
Ok(vm.new_str(code_obj.obj_name))
93+
fn code_co_name(code: PyCodeRef, _vm: &VirtualMachine) -> String {
94+
code.code.obj_name.clone()
12495
}

vm/src/obj/objfunction.rs

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,22 @@
11
use crate::frame::Scope;
22
use crate::function::PyFuncArgs;
3+
use crate::obj::objcode::PyCodeRef;
34
use crate::obj::objtype::PyClassRef;
4-
use crate::pyobject::{IdProtocol, PyContext, PyObjectRef, PyResult, PyValue, TypeProtocol};
5+
use crate::pyobject::{IdProtocol, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol};
56
use crate::vm::VirtualMachine;
67

8+
pub type PyFunctionRef = PyRef<PyFunction>;
9+
710
#[derive(Debug)]
811
pub struct PyFunction {
912
// TODO: these shouldn't be public
10-
pub code: PyObjectRef,
13+
pub code: PyCodeRef,
1114
pub scope: Scope,
1215
pub defaults: PyObjectRef,
1316
}
1417

1518
impl PyFunction {
16-
pub fn new(code: PyObjectRef, scope: Scope, defaults: PyObjectRef) -> Self {
19+
pub fn new(code: PyCodeRef, scope: Scope, defaults: PyObjectRef) -> Self {
1720
PyFunction {
1821
code,
1922
scope,
@@ -28,6 +31,12 @@ impl PyValue for PyFunction {
2831
}
2932
}
3033

34+
impl PyFunctionRef {
35+
fn code(self, _vm: &VirtualMachine) -> PyCodeRef {
36+
self.code.clone()
37+
}
38+
}
39+
3140
#[derive(Debug)]
3241
pub struct PyMethod {
3342
// TODO: these shouldn't be public
@@ -51,7 +60,7 @@ pub fn init(context: &PyContext) {
5160
let function_type = &context.function_type;
5261
extend_class!(context, function_type, {
5362
"__get__" => context.new_rustfunc(bind_method),
54-
"__code__" => context.new_property(function_code)
63+
"__code__" => context.new_property(PyFunctionRef::code)
5564
});
5665

5766
let builtin_function_or_method_type = &context.builtin_function_or_method_type;
@@ -73,10 +82,3 @@ fn bind_method(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
7382
Ok(vm.ctx.new_bound_method(function.clone(), obj.clone()))
7483
}
7584
}
76-
77-
fn function_code(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
78-
match args.args[0].payload() {
79-
Some(PyFunction { ref code, .. }) => Ok(code.clone()),
80-
None => Err(vm.new_type_error("no code".to_string())),
81-
}
82-
}

vm/src/obj/objgenerator.rs

Lines changed: 27 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -2,74 +2,56 @@
22
* The mythical generator.
33
*/
44

5-
use crate::frame::{ExecutionResult, Frame};
6-
use crate::function::PyFuncArgs;
5+
use crate::frame::{ExecutionResult, FrameRef};
76
use crate::obj::objtype::PyClassRef;
8-
use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol};
7+
use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue};
98
use crate::vm::VirtualMachine;
109

10+
pub type PyGeneratorRef = PyRef<PyGenerator>;
11+
1112
#[derive(Debug)]
1213
pub struct PyGenerator {
13-
frame: PyObjectRef,
14+
frame: FrameRef,
1415
}
15-
type PyGeneratorRef = PyRef<PyGenerator>;
1616

1717
impl PyValue for PyGenerator {
1818
fn class(vm: &VirtualMachine) -> PyClassRef {
1919
vm.ctx.generator_type()
2020
}
2121
}
2222

23-
pub fn init(context: &PyContext) {
24-
let generator_type = &context.generator_type;
25-
extend_class!(context, generator_type, {
26-
"__iter__" => context.new_rustfunc(generator_iter),
27-
"__next__" => context.new_rustfunc(generator_next),
28-
"send" => context.new_rustfunc(generator_send)
29-
});
30-
}
31-
32-
pub fn new_generator(frame: PyObjectRef, vm: &VirtualMachine) -> PyGeneratorRef {
33-
PyGenerator { frame }.into_ref(vm)
34-
}
35-
36-
fn generator_iter(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
37-
arg_check!(vm, args, required = [(o, Some(vm.ctx.generator_type()))]);
38-
Ok(o.clone())
39-
}
23+
impl PyGeneratorRef {
24+
pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyGeneratorRef {
25+
PyGenerator { frame }.into_ref(vm)
26+
}
4027

41-
fn generator_next(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
42-
arg_check!(vm, args, required = [(o, Some(vm.ctx.generator_type()))]);
43-
let value = vm.get_none();
44-
send(vm, o, &value)
45-
}
28+
fn iter(self, _vm: &VirtualMachine) -> PyGeneratorRef {
29+
self
30+
}
4631

47-
fn generator_send(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
48-
arg_check!(
49-
vm,
50-
args,
51-
required = [(o, Some(vm.ctx.generator_type())), (value, None)]
52-
);
53-
send(vm, o, value)
54-
}
32+
fn next(self, vm: &VirtualMachine) -> PyResult {
33+
self.send(vm.get_none(), vm)
34+
}
5535

56-
fn send(vm: &VirtualMachine, gen: &PyObjectRef, value: &PyObjectRef) -> PyResult {
57-
if let Some(PyGenerator { ref frame }) = gen.payload() {
58-
if let Some(frame) = frame.payload::<Frame>() {
59-
frame.push_value(value.clone());
60-
} else {
61-
panic!("Generator frame isn't a frame.");
62-
}
36+
fn send(self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
37+
self.frame.push_value(value.clone());
6338

64-
match vm.run_frame(frame.clone())? {
39+
match vm.run_frame(self.frame.clone())? {
6540
ExecutionResult::Yield(value) => Ok(value),
6641
ExecutionResult::Return(_value) => {
6742
// Stop iteration!
6843
let stop_iteration = vm.ctx.exceptions.stop_iteration.clone();
6944
Err(vm.new_exception(stop_iteration, "End of generator".to_string()))
7045
}
7146
}
72-
} else {
73-
panic!("Cannot extract frame from non-generator");
7447
}
7548
}
49+
50+
pub fn init(context: &PyContext) {
51+
let generator_type = &context.generator_type;
52+
extend_class!(context, generator_type, {
53+
"__iter__" => context.new_rustfunc(PyGeneratorRef::iter),
54+
"__next__" => context.new_rustfunc(PyGeneratorRef::next),
55+
"send" => context.new_rustfunc(PyGeneratorRef::send)
56+
});
57+
}

vm/src/pyobject.rs

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,15 @@ use num_traits::{One, Zero};
1414

1515
use crate::bytecode;
1616
use crate::exceptions;
17-
use crate::frame::{Frame, Scope};
17+
use crate::frame::Scope;
1818
use crate::function::{IntoPyNativeFunc, PyFuncArgs};
1919
use crate::obj::objbool;
2020
use crate::obj::objbuiltinfunc::PyBuiltinFunction;
2121
use crate::obj::objbytearray;
2222
use crate::obj::objbytes;
2323
use crate::obj::objclassmethod;
2424
use crate::obj::objcode;
25+
use crate::obj::objcode::PyCodeRef;
2526
use crate::obj::objcomplex::{self, PyComplex};
2627
use crate::obj::objdict::{self, PyDict};
2728
use crate::obj::objellipsis;
@@ -591,10 +592,6 @@ impl PyContext {
591592
)
592593
}
593594

594-
pub fn new_frame(&self, code: PyObjectRef, scope: Scope) -> PyObjectRef {
595-
PyObject::new(Frame::new(code, scope), self.frame_type())
596-
}
597-
598595
pub fn new_property<F, I, V>(&self, f: F) -> PyObjectRef
599596
where
600597
F: IntoPyNativeFunc<I, V>,
@@ -608,7 +605,7 @@ impl PyContext {
608605

609606
pub fn new_function(
610607
&self,
611-
code_obj: PyObjectRef,
608+
code_obj: PyCodeRef,
612609
scope: Scope,
613610
defaults: PyObjectRef,
614611
) -> PyObjectRef {

0 commit comments

Comments
 (0)