[rust] enter_ssa recurses into function expressions
Updates `enter_ssa()` to recurse into function expressions. In the TS compiler we use a single Builder instance and copy (references to) the function expression's blocks into the builder. That type of sharing just does not play nicely with Rust. But... we don't need to do that! We already know the context variables of the function expression, so we can lookup each of them to find their re-mapped identifier, and set that as the starting state for the entry block of the function expression. That lets us use normal recursion and otherwise not share any information between the outer and inner builders. Of course to make this work we actually have to populate Function.context, but the algorithm _should_ work.
Joe Savona committed
Jul 13, 2023 at 12:24 UTC
4762208b0f5f885ec68bf5afc989688d8e19048b
3 files changed
+52
-14
compiler/forget/crates/build-hir/src/build.rs
+2
@@ -84,6 +84,8 @@ pub fn build<'a>(
84
.map(|id| String::from_str_in(&id.name, &env.allocator)),
85
body,
86
params,
87
+ // TODO: populate context!
88
+ context: env.vec_new(),
89
is_async: fun.is_async,
90
is_generator: fun.is_generator,
91
}))
compiler/forget/crates/hir-ssa/src/enter.rs
+49
-14
@@ -3,7 +3,7 @@ use std::{cell::RefCell, rc::Rc};
3
use bumpalo::collections::{CollectIn, Vec};
4
use hir::{
5
BasicBlock, BlockId, Blocks, Environment, Function, Identifier, IdentifierData, IdentifierId,
6
- IdentifierOperand, Instruction, LValue, MutableRange, Phi,
6
+ IdentifierOperand, Instruction, InstructionValue, LValue, MutableRange, Phi,
7
};
8
use indexmap::{IndexMap, IndexSet};
9
use thiserror::Error;
@@ -13,13 +13,25 @@ use thiserror::Error;
13
pub struct SSAError;
14
15
pub fn enter_ssa<'a>(env: &Environment<'a>, fun: &mut Function<'a>) -> Result<(), SSAError> {
16
+ // assert_eq!(fun.context.is_empty())
17
+ enter_ssa_impl(env, fun, None)
18
+}
19
+
20
+pub fn enter_ssa_impl<'a>(
21
+ env: &Environment<'a>,
22
+ fun: &mut Function<'a>,
23
+ context_defs: Option<IndexMap<IdentifierId, Identifier<'a>>>,
24
+) -> Result<(), SSAError> {
25
let blocks = &fun.body.blocks;
26
let instructions = &mut fun.body.instructions;
27
let mut builder = Builder::new(env, fun.body.entry, blocks);
28
+ if let Some(context_defs) = context_defs {
29
+ builder.initialize_context(context_defs);
30
+ }
31
for param in &mut fun.params {
32
builder.visit_param(param);
33
}
22
- enter_ssa_impl(&mut builder, instructions)?;
34
+ visit_instructions(env, &mut builder, instructions)?;
35
36
let mut states = builder.complete();
37
@@ -31,7 +43,8 @@ pub fn enter_ssa<'a>(env: &Environment<'a>, fun: &mut Function<'a>) -> Result<()
43
Ok(())
44
}
45
34
-fn enter_ssa_impl<'a, 'e, 'f>(
46
+fn visit_instructions<'a, 'e, 'f>(
47
+ env: &Environment<'a>,
48
builder: &mut Builder<'a, 'e, 'f>,
49
instructions: &mut Vec<'a, Instruction<'a>>,
50
) -> Result<(), SSAError> {
@@ -40,9 +53,27 @@ fn enter_ssa_impl<'a, 'e, 'f>(
53
let instr = &mut instructions[usize::from(*instr_ix)];
54
instr.each_identifier_store(|store| builder.visit_store(store));
55
instr.each_identifier_load(|load| builder.visit_load(load));
56
+
57
+ if let InstructionValue::Function(fun) = &mut instr.value {
58
+ // Lookup each of the context variables referenced in the function
59
+ // against the current block. Note that variables which are reassigned somewhere
60
+ // and referenced in a function expression are always promoted to LoadContext
61
+ // so the context variables are all guaranteed to be const. We're just remapping
62
+ // the id along w the original declaration and other usages.
63
+ let context_defs: IndexMap<IdentifierId, Identifier> = fun
64
+ .lowered_function
65
+ .context
66
+ .iter()
67
+ .map(|id| {
68
+ let identifier = builder.get_id_at(block.id, &id.identifier);
69
+ (id.identifier.id, identifier)
70
+ })
71
+ .collect();
72
+ enter_ssa_impl(env, &mut fun.lowered_function, Some(context_defs))?;
73
+ }
74
}
44
- });
45
- Ok(())
75
+ Ok(())
76
+ })
77
}
78
79
#[derive(Debug)]
@@ -97,6 +128,11 @@ impl<'a, 'e, 'f> Builder<'a, 'e, 'f> {
128
}
129
}
130
131
+ fn initialize_context(&mut self, defs: IndexMap<IdentifierId, Identifier<'a>>) {
132
+ let state = self.states.get_mut(&self.current).unwrap();
133
+ state.defs = defs;
134
+ }
135
+
136
fn complete(self) -> IndexMap<BlockId, BlockState<'a>> {
137
self.states
138
}
@@ -217,9 +253,9 @@ impl<'a, 'e, 'f> Builder<'a, 'e, 'f> {
253
}
254
}
255
220
- fn each_block<F>(&mut self, mut f: F) -> ()
256
+ fn each_block<F>(&mut self, mut f: F) -> Result<(), SSAError>
257
where
222
- F: FnMut(&BasicBlock<'a>, &mut Self) -> (),
258
+ F: FnMut(&BasicBlock<'a>, &mut Self) -> Result<(), SSAError>,
259
{
260
let mut visited = IndexSet::new();
261
let block_ids: Vec<_> = self.blocks.keys().cloned().collect_in(self.env.allocator);
@@ -227,21 +263,20 @@ impl<'a, 'e, 'f> Builder<'a, 'e, 'f> {
263
visited.insert(block_id);
264
self.current = block_id;
265
let block = self.blocks.get(&block_id).unwrap();
230
- f(block, self);
266
+ f(block, self)?;
267
let successors = block.terminal.value.successors();
268
for successor in successors {
269
let block = self.blocks.get(&successor).unwrap();
270
let count = self
271
.unsealed_predecessors
236
- .get(&successor)
237
- .cloned()
238
- .unwrap_or(block.predecessors.len())
239
- - 1;
240
- self.unsealed_predecessors.insert(successor, count);
241
- if count == 0 && visited.contains(&successor) {
272
+ .entry(successor)
273
+ .or_insert(block.predecessors.len());
274
+ *count -= 1;
275
+ if *count == 0 && visited.contains(&successor) {
276
self.fix_incomplete_phis(successor)
277
}
278
}
279
}
280
+ Ok(())
281
}
282
}
compiler/forget/crates/hir/src/function.rs
+1
@@ -12,6 +12,7 @@ pub struct Function<'a> {
12
pub id: Option<String<'a>>,
13
pub body: HIR<'a>,
14
pub params: Vec<'a, IdentifierOperand<'a>>,
15
+ pub context: Vec<'a, IdentifierOperand<'a>>,
16
pub is_async: bool,
17
pub is_generator: bool,
18
}