Skip to content

Commit 67f17b6

Browse files
youknowoneyouknowone
committed
refactor ast
1 parent 889382b commit 67f17b6

10 files changed

Lines changed: 308 additions & 353 deletions

File tree

‎crates/vm/src/stdlib/ast.rs‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,18 @@ mod string;
5656
mod type_ignore;
5757
mod type_parameters;
5858

59+
/// Return the cached singleton instance for an operator/context node type,
60+
/// or create a new instance if none exists.
61+
fn singleton_node_to_object(vm: &VirtualMachine, node_type: &'static Py<PyType>) -> PyObjectRef {
62+
if let Some(instance) = node_type.get_attr(vm.ctx.intern_str("_instance")) {
63+
return instance;
64+
}
65+
NodeAst
66+
.into_ref_with_type(vm, node_type.to_owned())
67+
.unwrap()
68+
.into()
69+
}
70+
5971
fn get_node_field(vm: &VirtualMachine, obj: &PyObject, field: &'static str, typ: &str) -> PyResult {
6072
vm.get_attribute_opt(obj.to_owned(), field)?
6173
.ok_or_else(|| vm.new_type_error(format!(r#"required field "{field}" missing from {typ}"#)))
@@ -772,7 +784,7 @@ pub(crate) fn validate_ast_object(vm: &VirtualMachine, object: PyObjectRef) -> P
772784
}
773785

774786
// Used by builtins::compile()
775-
pub const PY_COMPILE_FLAG_AST_ONLY: i32 = 0x0400;
787+
pub const PY_CF_ONLY_AST: i32 = 0x0400;
776788

777789
// The following flags match the values from Include/cpython/compile.h
778790
// Caveat emptor: These flags are undocumented on purpose and depending
@@ -781,7 +793,7 @@ pub const PY_CF_SOURCE_IS_UTF8: i32 = 0x0100;
781793
pub const PY_CF_DONT_IMPLY_DEDENT: i32 = 0x200;
782794
pub const PY_CF_IGNORE_COOKIE: i32 = 0x0800;
783795
pub const PY_CF_ALLOW_INCOMPLETE_INPUT: i32 = 0x4000;
784-
pub const PY_CF_OPTIMIZED_AST: i32 = 0x8000 | PY_COMPILE_FLAG_AST_ONLY;
796+
pub const PY_CF_OPTIMIZED_AST: i32 = 0x8000 | PY_CF_ONLY_AST;
785797
pub const PY_CF_TYPE_COMMENTS: i32 = 0x1000;
786798
pub const PY_CF_ALLOW_TOP_LEVEL_AWAIT: i32 = 0x2000;
787799

@@ -802,7 +814,7 @@ const CO_FUTURE_GENERATOR_STOP: i32 = 0x800000;
802814
const CO_FUTURE_ANNOTATIONS: i32 = 0x1000000;
803815

804816
// Used by builtins::compile() - the summary of all flags
805-
pub const PY_COMPILE_FLAGS_MASK: i32 = PY_COMPILE_FLAG_AST_ONLY
817+
pub const PY_COMPILE_FLAGS_MASK: i32 = PY_CF_ONLY_AST
806818
| PY_CF_SOURCE_IS_UTF8
807819
| PY_CF_DONT_IMPLY_DEDENT
808820
| PY_CF_IGNORE_COOKIE

‎crates/vm/src/stdlib/ast/exception.rs‎

Lines changed: 24 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -9,22 +9,22 @@ impl Node for ast::ExceptHandler {
99
}
1010
}
1111
fn ast_from_object(
12-
_vm: &VirtualMachine,
12+
vm: &VirtualMachine,
1313
source_file: &SourceFile,
14-
_object: PyObjectRef,
14+
object: PyObjectRef,
1515
) -> PyResult<Self> {
16-
let _cls = _object.class();
16+
let cls = object.class();
1717
Ok(
18-
if _cls.is(pyast::NodeExceptHandlerExceptHandler::static_type()) {
18+
if cls.is(pyast::NodeExceptHandlerExceptHandler::static_type()) {
1919
Self::ExceptHandler(ast::ExceptHandlerExceptHandler::ast_from_object(
20-
_vm,
20+
vm,
2121
source_file,
22-
_object,
22+
object,
2323
)?)
2424
} else {
25-
return Err(_vm.new_type_error(format!(
25+
return Err(vm.new_type_error(format!(
2626
"expected some sort of excepthandler, but got {}",
27-
_object.repr(_vm)?
27+
object.repr(vm)?
2828
)));
2929
},
3030
)
@@ -33,50 +33,50 @@ impl Node for ast::ExceptHandler {
3333

3434
// constructor
3535
impl Node for ast::ExceptHandlerExceptHandler {
36-
fn ast_to_object(self, _vm: &VirtualMachine, source_file: &SourceFile) -> PyObjectRef {
36+
fn ast_to_object(self, vm: &VirtualMachine, source_file: &SourceFile) -> PyObjectRef {
3737
let Self {
3838
node_index: _,
3939
type_,
4040
name,
4141
body,
42-
range: _range,
42+
range,
4343
} = self;
4444
let node = NodeAst
4545
.into_ref_with_type(
46-
_vm,
46+
vm,
4747
pyast::NodeExceptHandlerExceptHandler::static_type().to_owned(),
4848
)
4949
.unwrap();
5050
let dict = node.as_object().dict().unwrap();
51-
dict.set_item("type", type_.ast_to_object(_vm, source_file), _vm)
51+
dict.set_item("type", type_.ast_to_object(vm, source_file), vm)
5252
.unwrap();
53-
dict.set_item("name", name.ast_to_object(_vm, source_file), _vm)
53+
dict.set_item("name", name.ast_to_object(vm, source_file), vm)
5454
.unwrap();
55-
dict.set_item("body", body.ast_to_object(_vm, source_file), _vm)
55+
dict.set_item("body", body.ast_to_object(vm, source_file), vm)
5656
.unwrap();
57-
node_add_location(&dict, _range, _vm, source_file);
57+
node_add_location(&dict, range, vm, source_file);
5858
node.into()
5959
}
6060

6161
fn ast_from_object(
62-
_vm: &VirtualMachine,
62+
vm: &VirtualMachine,
6363
source_file: &SourceFile,
64-
_object: PyObjectRef,
64+
object: PyObjectRef,
6565
) -> PyResult<Self> {
6666
Ok(Self {
6767
node_index: Default::default(),
68-
type_: get_node_field_opt(_vm, &_object, "type")?
69-
.map(|obj| Node::ast_from_object(_vm, source_file, obj))
68+
type_: get_node_field_opt(vm, &object, "type")?
69+
.map(|obj| Node::ast_from_object(vm, source_file, obj))
7070
.transpose()?,
71-
name: get_node_field_opt(_vm, &_object, "name")?
72-
.map(|obj| Node::ast_from_object(_vm, source_file, obj))
71+
name: get_node_field_opt(vm, &object, "name")?
72+
.map(|obj| Node::ast_from_object(vm, source_file, obj))
7373
.transpose()?,
7474
body: Node::ast_from_object(
75-
_vm,
75+
vm,
7676
source_file,
77-
get_node_field(_vm, &_object, "body", "ExceptHandler")?,
77+
get_node_field(vm, &object, "body", "ExceptHandler")?,
7878
)?,
79-
range: range_from_object(_vm, source_file, _object, "ExceptHandler")?,
79+
range: range_from_object(vm, source_file, object, "ExceptHandler")?,
8080
})
8181
}
8282
}

‎crates/vm/src/stdlib/ast/expression.rs‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1256,26 +1256,20 @@ impl Node for ast::ExprContext {
12561256
unimplemented!("Invalid expression context is not allowed in Python AST")
12571257
}
12581258
};
1259-
if let Some(instance) = node_type.get_attr(vm.ctx.intern_str("_instance")) {
1260-
return instance;
1261-
}
1262-
NodeAst
1263-
.into_ref_with_type(vm, node_type.to_owned())
1264-
.unwrap()
1265-
.into()
1259+
singleton_node_to_object(vm, node_type)
12661260
}
12671261

12681262
fn ast_from_object(
12691263
vm: &VirtualMachine,
12701264
_source_file: &SourceFile,
12711265
object: PyObjectRef,
12721266
) -> PyResult<Self> {
1273-
let _cls = object.class();
1274-
Ok(if _cls.is(pyast::NodeExprContextLoad::static_type()) {
1267+
let cls = object.class();
1268+
Ok(if cls.is(pyast::NodeExprContextLoad::static_type()) {
12751269
Self::Load
1276-
} else if _cls.is(pyast::NodeExprContextStore::static_type()) {
1270+
} else if cls.is(pyast::NodeExprContextStore::static_type()) {
12771271
Self::Store
1278-
} else if _cls.is(pyast::NodeExprContextDel::static_type()) {
1272+
} else if cls.is(pyast::NodeExprContextDel::static_type()) {
12791273
Self::Del
12801274
} else {
12811275
return Err(vm.new_type_error(format!(

‎crates/vm/src/stdlib/ast/operator.rs‎

Lines changed: 53 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -8,29 +8,23 @@ impl Node for ast::BoolOp {
88
Self::And => pyast::NodeBoolOpAnd::static_type(),
99
Self::Or => pyast::NodeBoolOpOr::static_type(),
1010
};
11-
if let Some(instance) = node_type.get_attr(vm.ctx.intern_str("_instance")) {
12-
return instance;
13-
}
14-
NodeAst
15-
.into_ref_with_type(vm, node_type.to_owned())
16-
.unwrap()
17-
.into()
11+
singleton_node_to_object(vm, node_type)
1812
}
1913

2014
fn ast_from_object(
21-
_vm: &VirtualMachine,
15+
vm: &VirtualMachine,
2216
_source_file: &SourceFile,
23-
_object: PyObjectRef,
17+
object: PyObjectRef,
2418
) -> PyResult<Self> {
25-
let _cls = _object.class();
26-
Ok(if _cls.is(pyast::NodeBoolOpAnd::static_type()) {
19+
let cls = object.class();
20+
Ok(if cls.is(pyast::NodeBoolOpAnd::static_type()) {
2721
Self::And
28-
} else if _cls.is(pyast::NodeBoolOpOr::static_type()) {
22+
} else if cls.is(pyast::NodeBoolOpOr::static_type()) {
2923
Self::Or
3024
} else {
31-
return Err(_vm.new_type_error(format!(
25+
return Err(vm.new_type_error(format!(
3226
"expected some sort of boolop, but got {}",
33-
_object.repr(_vm)?
27+
object.repr(vm)?
3428
)));
3529
})
3630
}
@@ -54,51 +48,45 @@ impl Node for ast::Operator {
5448
Self::BitAnd => pyast::NodeOperatorBitAnd::static_type(),
5549
Self::FloorDiv => pyast::NodeOperatorFloorDiv::static_type(),
5650
};
57-
if let Some(instance) = node_type.get_attr(vm.ctx.intern_str("_instance")) {
58-
return instance;
59-
}
60-
NodeAst
61-
.into_ref_with_type(vm, node_type.to_owned())
62-
.unwrap()
63-
.into()
51+
singleton_node_to_object(vm, node_type)
6452
}
6553

6654
fn ast_from_object(
67-
_vm: &VirtualMachine,
55+
vm: &VirtualMachine,
6856
_source_file: &SourceFile,
69-
_object: PyObjectRef,
57+
object: PyObjectRef,
7058
) -> PyResult<Self> {
71-
let _cls = _object.class();
72-
Ok(if _cls.is(pyast::NodeOperatorAdd::static_type()) {
59+
let cls = object.class();
60+
Ok(if cls.is(pyast::NodeOperatorAdd::static_type()) {
7361
Self::Add
74-
} else if _cls.is(pyast::NodeOperatorSub::static_type()) {
62+
} else if cls.is(pyast::NodeOperatorSub::static_type()) {
7563
Self::Sub
76-
} else if _cls.is(pyast::NodeOperatorMult::static_type()) {
64+
} else if cls.is(pyast::NodeOperatorMult::static_type()) {
7765
Self::Mult
78-
} else if _cls.is(pyast::NodeOperatorMatMult::static_type()) {
66+
} else if cls.is(pyast::NodeOperatorMatMult::static_type()) {
7967
Self::MatMult
80-
} else if _cls.is(pyast::NodeOperatorDiv::static_type()) {
68+
} else if cls.is(pyast::NodeOperatorDiv::static_type()) {
8169
Self::Div
82-
} else if _cls.is(pyast::NodeOperatorMod::static_type()) {
70+
} else if cls.is(pyast::NodeOperatorMod::static_type()) {
8371
Self::Mod
84-
} else if _cls.is(pyast::NodeOperatorPow::static_type()) {
72+
} else if cls.is(pyast::NodeOperatorPow::static_type()) {
8573
Self::Pow
86-
} else if _cls.is(pyast::NodeOperatorLShift::static_type()) {
74+
} else if cls.is(pyast::NodeOperatorLShift::static_type()) {
8775
Self::LShift
88-
} else if _cls.is(pyast::NodeOperatorRShift::static_type()) {
76+
} else if cls.is(pyast::NodeOperatorRShift::static_type()) {
8977
Self::RShift
90-
} else if _cls.is(pyast::NodeOperatorBitOr::static_type()) {
78+
} else if cls.is(pyast::NodeOperatorBitOr::static_type()) {
9179
Self::BitOr
92-
} else if _cls.is(pyast::NodeOperatorBitXor::static_type()) {
80+
} else if cls.is(pyast::NodeOperatorBitXor::static_type()) {
9381
Self::BitXor
94-
} else if _cls.is(pyast::NodeOperatorBitAnd::static_type()) {
82+
} else if cls.is(pyast::NodeOperatorBitAnd::static_type()) {
9583
Self::BitAnd
96-
} else if _cls.is(pyast::NodeOperatorFloorDiv::static_type()) {
84+
} else if cls.is(pyast::NodeOperatorFloorDiv::static_type()) {
9785
Self::FloorDiv
9886
} else {
99-
return Err(_vm.new_type_error(format!(
87+
return Err(vm.new_type_error(format!(
10088
"expected some sort of operator, but got {}",
101-
_object.repr(_vm)?
89+
object.repr(vm)?
10290
)));
10391
})
10492
}
@@ -113,33 +101,27 @@ impl Node for ast::UnaryOp {
113101
Self::UAdd => pyast::NodeUnaryOpUAdd::static_type(),
114102
Self::USub => pyast::NodeUnaryOpUSub::static_type(),
115103
};
116-
if let Some(instance) = node_type.get_attr(vm.ctx.intern_str("_instance")) {
117-
return instance;
118-
}
119-
NodeAst
120-
.into_ref_with_type(vm, node_type.to_owned())
121-
.unwrap()
122-
.into()
104+
singleton_node_to_object(vm, node_type)
123105
}
124106

125107
fn ast_from_object(
126-
_vm: &VirtualMachine,
108+
vm: &VirtualMachine,
127109
_source_file: &SourceFile,
128-
_object: PyObjectRef,
110+
object: PyObjectRef,
129111
) -> PyResult<Self> {
130-
let _cls = _object.class();
131-
Ok(if _cls.is(pyast::NodeUnaryOpInvert::static_type()) {
112+
let cls = object.class();
113+
Ok(if cls.is(pyast::NodeUnaryOpInvert::static_type()) {
132114
Self::Invert
133-
} else if _cls.is(pyast::NodeUnaryOpNot::static_type()) {
115+
} else if cls.is(pyast::NodeUnaryOpNot::static_type()) {
134116
Self::Not
135-
} else if _cls.is(pyast::NodeUnaryOpUAdd::static_type()) {
117+
} else if cls.is(pyast::NodeUnaryOpUAdd::static_type()) {
136118
Self::UAdd
137-
} else if _cls.is(pyast::NodeUnaryOpUSub::static_type()) {
119+
} else if cls.is(pyast::NodeUnaryOpUSub::static_type()) {
138120
Self::USub
139121
} else {
140-
return Err(_vm.new_type_error(format!(
122+
return Err(vm.new_type_error(format!(
141123
"expected some sort of unaryop, but got {}",
142-
_object.repr(_vm)?
124+
object.repr(vm)?
143125
)));
144126
})
145127
}
@@ -160,45 +142,39 @@ impl Node for ast::CmpOp {
160142
Self::In => pyast::NodeCmpOpIn::static_type(),
161143
Self::NotIn => pyast::NodeCmpOpNotIn::static_type(),
162144
};
163-
if let Some(instance) = node_type.get_attr(vm.ctx.intern_str("_instance")) {
164-
return instance;
165-
}
166-
NodeAst
167-
.into_ref_with_type(vm, node_type.to_owned())
168-
.unwrap()
169-
.into()
145+
singleton_node_to_object(vm, node_type)
170146
}
171147

172148
fn ast_from_object(
173-
_vm: &VirtualMachine,
149+
vm: &VirtualMachine,
174150
_source_file: &SourceFile,
175-
_object: PyObjectRef,
151+
object: PyObjectRef,
176152
) -> PyResult<Self> {
177-
let _cls = _object.class();
178-
Ok(if _cls.is(pyast::NodeCmpOpEq::static_type()) {
153+
let cls = object.class();
154+
Ok(if cls.is(pyast::NodeCmpOpEq::static_type()) {
179155
Self::Eq
180-
} else if _cls.is(pyast::NodeCmpOpNotEq::static_type()) {
156+
} else if cls.is(pyast::NodeCmpOpNotEq::static_type()) {
181157
Self::NotEq
182-
} else if _cls.is(pyast::NodeCmpOpLt::static_type()) {
158+
} else if cls.is(pyast::NodeCmpOpLt::static_type()) {
183159
Self::Lt
184-
} else if _cls.is(pyast::NodeCmpOpLtE::static_type()) {
160+
} else if cls.is(pyast::NodeCmpOpLtE::static_type()) {
185161
Self::LtE
186-
} else if _cls.is(pyast::NodeCmpOpGt::static_type()) {
162+
} else if cls.is(pyast::NodeCmpOpGt::static_type()) {
187163
Self::Gt
188-
} else if _cls.is(pyast::NodeCmpOpGtE::static_type()) {
164+
} else if cls.is(pyast::NodeCmpOpGtE::static_type()) {
189165
Self::GtE
190-
} else if _cls.is(pyast::NodeCmpOpIs::static_type()) {
166+
} else if cls.is(pyast::NodeCmpOpIs::static_type()) {
191167
Self::Is
192-
} else if _cls.is(pyast::NodeCmpOpIsNot::static_type()) {
168+
} else if cls.is(pyast::NodeCmpOpIsNot::static_type()) {
193169
Self::IsNot
194-
} else if _cls.is(pyast::NodeCmpOpIn::static_type()) {
170+
} else if cls.is(pyast::NodeCmpOpIn::static_type()) {
195171
Self::In
196-
} else if _cls.is(pyast::NodeCmpOpNotIn::static_type()) {
172+
} else if cls.is(pyast::NodeCmpOpNotIn::static_type()) {
197173
Self::NotIn
198174
} else {
199-
return Err(_vm.new_type_error(format!(
175+
return Err(vm.new_type_error(format!(
200176
"expected some sort of cmpop, but got {}",
201-
_object.repr(_vm)?
177+
object.repr(vm)?
202178
)));
203179
})
204180
}

0 commit comments

Comments
 (0)