Skip to content

Commit 6082049

Browse files
coolreader18coolreader18
committed
Remove the need for a label_map in CodeObject
1 parent fd8d9c9 commit 6082049

File tree

3 files changed

+118
-53
lines changed

3 files changed

+118
-53
lines changed

‎bytecode/src/bytecode.rs‎

Lines changed: 44 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ use itertools::Itertools;
77
use num_bigint::BigInt;
88
use num_complex::Complex64;
99
use serde::{Deserialize, Serialize};
10-
use std::collections::{BTreeMap, HashSet};
10+
use std::collections::BTreeSet;
1111
use std::fmt;
1212

1313
/// Sourcecode location.
@@ -94,8 +94,6 @@ impl ConstantBag for BasicBag {
9494
#[derive(Clone, PartialEq, Serialize, Deserialize)]
9595
pub struct CodeObject<C: Constant = ConstantData> {
9696
pub instructions: Vec<Instruction>,
97-
/// Jump targets.
98-
pub label_map: BTreeMap<Label, usize>,
9997
pub locations: Vec<Location>,
10098
pub flags: CodeFlags,
10199
pub posonlyarg_count: usize, // Number of positional-only arguments
@@ -146,11 +144,13 @@ impl CodeFlags {
146144
}
147145

148146
#[derive(Serialize, Debug, Deserialize, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]
149-
pub struct Label(usize);
150-
151-
impl Label {
152-
pub fn new(label: usize) -> Self {
153-
Label(label)
147+
#[repr(transparent)]
148+
// XXX: if you add a new instruction that stores a Label, make sure to add it in
149+
// compile::CodeInfo::finalize_code
150+
pub struct Label(pub usize);
151+
impl fmt::Display for Label {
152+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
153+
self.0.fmt(f)
154154
}
155155
}
156156

@@ -551,7 +551,6 @@ impl<C: Constant> CodeObject<C> {
551551
) -> Self {
552552
CodeObject {
553553
instructions: Vec::new(),
554-
label_map: BTreeMap::new(),
555554
locations: Vec::new(),
556555
flags,
557556
posonlyarg_count,
@@ -602,9 +601,31 @@ impl<C: Constant> CodeObject<C> {
602601
expand_codeobjects: bool,
603602
level: usize,
604603
) -> fmt::Result {
605-
let label_targets: HashSet<&usize> = self.label_map.values().collect();
604+
let mut label_targets = BTreeSet::new();
605+
for instruction in &self.instructions {
606+
let label = match instruction {
607+
Jump { target } => target,
608+
JumpIfTrue { target } => target,
609+
JumpIfFalse { target } => target,
610+
JumpIfTrueOrPop { target } => target,
611+
JumpIfFalseOrPop { target } => target,
612+
ForIter { target } => target,
613+
SetupLoop { start, end } => {
614+
label_targets.insert(*start);
615+
label_targets.insert(*end);
616+
continue;
617+
}
618+
SetupFinally { handler } => handler,
619+
SetupExcept { handler } => handler,
620+
SetupWith { end } => end,
621+
SetupAsyncWith { end } => end,
622+
_ => continue,
623+
};
624+
label_targets.insert(*label);
625+
}
626+
606627
for (offset, instruction) in self.instructions.iter().enumerate() {
607-
let arrow = if label_targets.contains(&offset) {
628+
let arrow = if label_targets.contains(&Label(offset)) {
608629
">>"
609630
} else {
610631
" "
@@ -613,14 +634,7 @@ impl<C: Constant> CodeObject<C> {
613634
write!(f, " ")?;
614635
}
615636
write!(f, "{} {:5} ", arrow, offset)?;
616-
instruction.fmt_dis(
617-
f,
618-
&self.label_map,
619-
&self.constants,
620-
&self.names,
621-
expand_codeobjects,
622-
level,
623-
)?;
637+
instruction.fmt_dis(f, &self.constants, &self.names, expand_codeobjects, level)?;
624638
}
625639
Ok(())
626640
}
@@ -649,7 +663,6 @@ impl<C: Constant> CodeObject<C> {
649663
.collect(),
650664

651665
instructions: self.instructions,
652-
label_map: self.label_map,
653666
locations: self.locations,
654667
flags: self.flags,
655668
posonlyarg_count: self.posonlyarg_count,
@@ -675,7 +688,6 @@ impl<C: Constant> CodeObject<C> {
675688
.collect(),
676689

677690
instructions: self.instructions.clone(),
678-
label_map: self.label_map.clone(),
679691
locations: self.locations.clone(),
680692
flags: self.flags,
681693
posonlyarg_count: self.posonlyarg_count,
@@ -723,7 +735,6 @@ impl Instruction {
723735
fn fmt_dis<C: Constant>(
724736
&self,
725737
f: &mut fmt::Formatter,
726-
label_map: &BTreeMap<Label, usize>,
727738
constants: &[C],
728739
names: &[C::Name],
729740
expand_codeobjects: bool,
@@ -803,28 +814,28 @@ impl Instruction {
803814
GetIter => w!(GetIter),
804815
Continue => w!(Continue),
805816
Break => w!(Break),
806-
Jump { target } => w!(Jump, label_map[target]),
807-
JumpIfTrue { target } => w!(JumpIfTrue, label_map[target]),
808-
JumpIfFalse { target } => w!(JumpIfFalse, label_map[target]),
809-
JumpIfTrueOrPop { target } => w!(JumpIfTrueOrPop, label_map[target]),
810-
JumpIfFalseOrPop { target } => w!(JumpIfFalseOrPop, label_map[target]),
817+
Jump { target } => w!(Jump, target),
818+
JumpIfTrue { target } => w!(JumpIfTrue, target),
819+
JumpIfFalse { target } => w!(JumpIfFalse, target),
820+
JumpIfTrueOrPop { target } => w!(JumpIfTrueOrPop, target),
821+
JumpIfFalseOrPop { target } => w!(JumpIfFalseOrPop, target),
811822
MakeFunction => w!(MakeFunction),
812823
CallFunction { typ } => w!(CallFunction, format!("{:?}", typ)),
813-
ForIter { target } => w!(ForIter, label_map[target]),
824+
ForIter { target } => w!(ForIter, target),
814825
ReturnValue => w!(ReturnValue),
815826
YieldValue => w!(YieldValue),
816827
YieldFrom => w!(YieldFrom),
817828
SetupAnnotation => w!(SetupAnnotation),
818-
SetupLoop { start, end } => w!(SetupLoop, label_map[start], label_map[end]),
819-
SetupExcept { handler } => w!(SetupExcept, label_map[handler]),
820-
SetupFinally { handler } => w!(SetupFinally, label_map[handler]),
829+
SetupLoop { start, end } => w!(SetupLoop, start, end),
830+
SetupExcept { handler } => w!(SetupExcept, handler),
831+
SetupFinally { handler } => w!(SetupFinally, handler),
821832
EnterFinally => w!(EnterFinally),
822833
EndFinally => w!(EndFinally),
823-
SetupWith { end } => w!(SetupWith, label_map[end]),
834+
SetupWith { end } => w!(SetupWith, end),
824835
WithCleanupStart => w!(WithCleanupStart),
825836
WithCleanupFinish => w!(WithCleanupFinish),
826837
BeforeAsyncWith => w!(BeforeAsyncWith),
827-
SetupAsyncWith { end } => w!(SetupAsyncWith, label_map[end]),
838+
SetupAsyncWith { end } => w!(SetupAsyncWith, end),
828839
PopBlock => w!(PopBlock),
829840
Raise { argc } => w!(Raise, argc),
830841
BuildString { size } => w!(BuildString, size),

‎compiler/src/compile.rs‎

Lines changed: 72 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,52 @@ use rustpython_bytecode::bytecode::{self, CallType, CodeObject, Instruction, Lab
1818

1919
type CompileResult<T> = Result<T, CompileError>;
2020

21+
struct CodeInfo {
22+
code: CodeObject,
23+
name_cache: IndexSet<String>,
24+
label_map: Vec<Option<Label>>,
25+
}
26+
impl CodeInfo {
27+
fn finalize_code(self) -> CodeObject {
28+
let CodeInfo {
29+
mut code,
30+
name_cache,
31+
label_map,
32+
} = self;
33+
code.names.extend(name_cache);
34+
for instruction in &mut code.instructions {
35+
use Instruction::*;
36+
// this is a little bit hacky, as until now the data stored inside Labels in
37+
// Instructions is just bookkeeping, but I think it's the best way to do this
38+
// XXX: any new instruction that uses a label has to be added here
39+
let label = match instruction {
40+
Jump { target } => target,
41+
JumpIfTrue { target } => target,
42+
JumpIfFalse { target } => target,
43+
JumpIfTrueOrPop { target } => target,
44+
JumpIfFalseOrPop { target } => target,
45+
ForIter { target } => target,
46+
SetupLoop { start, end } => {
47+
*start = label_map[start.0].expect("label never set");
48+
*end = label_map[end.0].expect("label never set");
49+
continue;
50+
}
51+
SetupFinally { handler } => handler,
52+
SetupExcept { handler } => handler,
53+
SetupWith { end } => end,
54+
SetupAsyncWith { end } => end,
55+
_ => continue,
56+
};
57+
*label = label_map[label.0].expect("label never set");
58+
}
59+
code
60+
}
61+
}
62+
2163
/// Main structure holding the state of compilation.
2264
struct Compiler {
23-
output_stack: Vec<(CodeObject, IndexSet<String>)>,
65+
code_stack: Vec<CodeInfo>,
2466
symbol_table_stack: Vec<SymbolTable>,
25-
nxt_label: usize,
2667
source_path: String,
2768
current_source_location: ast::Location,
2869
current_qualified_path: Option<String>,
@@ -124,9 +165,8 @@ pub fn compile_program_single(
124165
impl Compiler {
125166
fn new(opts: CompileOpts, source_path: String) -> Self {
126167
Compiler {
127-
output_stack: Vec::new(),
168+
code_stack: Vec::new(),
128169
symbol_table_stack: Vec::new(),
129-
nxt_label: 0,
130170
source_path,
131171
current_source_location: ast::Location::default(),
132172
current_qualified_path: None,
@@ -151,7 +191,11 @@ impl Compiler {
151191
}
152192

153193
fn push_output(&mut self, code: CodeObject) {
154-
self.output_stack.push((code, IndexSet::new()));
194+
self.code_stack.push(CodeInfo {
195+
code,
196+
name_cache: IndexSet::new(),
197+
label_map: Vec::new(),
198+
});
155199
}
156200

157201
fn push_new_code_object(&mut self, obj_name: String) {
@@ -168,15 +212,17 @@ impl Compiler {
168212
}
169213

170214
fn pop_code_object(&mut self) -> CodeObject {
171-
let (mut code, names) = self.output_stack.pop().unwrap();
172-
code.names.extend(names);
173-
code
215+
self.code_stack.pop().unwrap().finalize_code()
174216
}
175217

176218
// could take impl Into<Cow<str>>, but everything is borrowed from ast structs; we never
177219
// actually have a `String` to pass
178220
fn name(&mut self, name: &str) -> bytecode::NameIdx {
179-
let cache = &mut self.output_stack.last_mut().expect("nothing on stack").1;
221+
let cache = &mut self
222+
.code_stack
223+
.last_mut()
224+
.expect("nothing on stack")
225+
.name_cache;
180226
if let Some(x) = cache.get_index_of(name) {
181227
x
182228
} else {
@@ -189,7 +235,7 @@ impl Compiler {
189235
program: &ast::Program,
190236
symbol_table: SymbolTable,
191237
) -> CompileResult<()> {
192-
let size_before = self.output_stack.len();
238+
let size_before = self.code_stack.len();
193239
self.symbol_table_stack.push(symbol_table);
194240

195241
let (statements, doc) = get_doc(&program.statements);
@@ -203,7 +249,7 @@ impl Compiler {
203249
}
204250
self.compile_statements(statements)?;
205251

206-
assert_eq!(self.output_stack.len(), size_before);
252+
assert_eq!(self.code_stack.len(), size_before);
207253

208254
// Emit None at end:
209255
self.emit_constant(bytecode::ConstantData::None);
@@ -2241,24 +2287,31 @@ impl Compiler {
22412287

22422288
fn current_code(&mut self) -> &mut CodeObject {
22432289
&mut self
2244-
.output_stack
2290+
.code_stack
22452291
.last_mut()
22462292
.expect("No OutputStream on stack")
2247-
.0
2293+
.code
22482294
}
22492295

22502296
// Generate a new label
22512297
fn new_label(&mut self) -> Label {
2252-
let l = Label::new(self.nxt_label);
2253-
self.nxt_label += 1;
2254-
l
2298+
let label_map = &mut self.code_stack.last_mut().unwrap().label_map;
2299+
let label = Label(label_map.len());
2300+
label_map.push(None);
2301+
label
22552302
}
22562303

22572304
// Assign current position the given label
22582305
fn set_label(&mut self, label: Label) {
2259-
let code = self.current_code();
2260-
let pos = code.instructions.len();
2261-
code.label_map.insert(label, pos);
2306+
let CodeInfo {
2307+
code, label_map, ..
2308+
} = self.code_stack.last_mut().unwrap();
2309+
let actual_label = Label(code.instructions.len());
2310+
let prev_val = std::mem::replace(&mut label_map[label.0], Some(actual_label));
2311+
debug_assert!(
2312+
prev_val.map_or(true, |x| x == actual_label),
2313+
"double-set a label"
2314+
);
22622315
}
22632316

22642317
fn set_source_location(&mut self, location: ast::Location) {

‎vm/src/frame.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1233,8 +1233,9 @@ impl ExecutingFrame<'_> {
12331233
}
12341234
}
12351235

1236+
#[inline]
12361237
fn jump(&mut self, label: bytecode::Label) {
1237-
let target_pc = self.code.label_map[&label];
1238+
let target_pc = label.0;
12381239
vm_trace!("jump from {:?} to {:?}", self.lasti(), target_pc);
12391240
self.lasti.store(target_pc, Ordering::Relaxed);
12401241
}

0 commit comments

Comments
 (0)