Refactoring: extract Storage out of WasmBackend

This commit is contained in:
Brian Carroll 2021-10-12 10:17:55 +01:00
parent d2e01bd10a
commit 472943a068
2 changed files with 444 additions and 406 deletions

View file

@ -10,13 +10,10 @@ use roc_module::symbol::Symbol;
use roc_mono::ir::{CallType, Expr, JoinPointId, Literal, Proc, Stmt};
use roc_mono::layout::{Builtin, Layout};
use crate::code_builder::{CodeBuilder, VirtualMachineSymbolState};
use crate::code_builder::CodeBuilder;
use crate::layout::WasmLayout;
use crate::storage::{StackMemoryLocation, StoredValue};
use crate::{
copy_memory, pop_stack_frame, push_stack_frame, round_up_to_alignment, CopyMemoryConfig,
LocalId, ALIGN_1, ALIGN_2, ALIGN_4, ALIGN_8, PTR_SIZE, PTR_TYPE,
};
use crate::storage::{LocalKind, StackMemoryLocation, Storage, StoredValue};
use crate::{copy_memory, pop_stack_frame, push_stack_frame, CopyMemoryConfig, LocalId, PTR_TYPE};
// Don't allocate any constant data at address zero or near it. Would be valid, but bug-prone.
// Follow Emscripten's example by using 1kB (4 bytes would probably do)
@ -25,30 +22,20 @@ const UNUSED_DATA_SECTION_BYTES: u32 = 1024;
#[derive(Clone, Copy, Debug)]
struct LabelId(u32);
enum LocalKind {
Parameter,
Variable,
}
// TODO: use Bumpalo Vec once parity_wasm supports general iterators (>=0.43)
pub struct WasmBackend<'a> {
// Module: Wasm AST
// Module level: Wasm AST
pub module_builder: ModuleBuilder,
// Module: internal state & IR mappings
// Module level: internal state & IR mappings
_data_offset_map: MutMap<Literal<'a>, u32>,
_data_offset_next: u32,
proc_symbol_map: MutMap<Symbol, CodeLocation>,
// Functions: Wasm AST
// Function level
code_builder: CodeBuilder,
arg_types: std::vec::Vec<ValueType>,
local_types: std::vec::Vec<ValueType>,
storage: Storage,
// Functions: internal state & IR mappings
stack_memory: i32,
stack_frame_pointer: Option<LocalId>,
symbol_storage_map: MutMap<Symbol, StoredValue>,
/// how many blocks deep are we (used for jumps)
block_depth: u32,
joinpoint_label_map: MutMap<JoinPointId, (u32, std::vec::Vec<StoredValue>)>,
@ -67,27 +54,16 @@ impl<'a> WasmBackend<'a> {
block_depth: 0,
joinpoint_label_map: MutMap::default(),
stack_memory: 0,
// Functions: Wasm AST
// Functions
code_builder: CodeBuilder::new(),
arg_types: std::vec::Vec::with_capacity(8),
local_types: std::vec::Vec::with_capacity(32),
stack_frame_pointer: None,
symbol_storage_map: MutMap::default(),
storage: Storage::new(),
}
}
fn reset(&mut self) {
// Functions: Wasm AST
self.code_builder.clear();
self.arg_types.clear();
self.local_types.clear();
// Functions: internal state & IR mappings
self.stack_memory = 0;
self.stack_frame_pointer = None;
self.symbol_storage_map.clear();
self.storage.clear();
self.joinpoint_label_map.clear();
assert_eq!(self.block_depth, 0);
}
@ -119,7 +95,7 @@ impl<'a> WasmBackend<'a> {
let ret_layout = WasmLayout::new(&proc.ret_layout);
let signature_builder = if let WasmLayout::StackMemory { .. } = ret_layout {
self.arg_types.push(PTR_TYPE);
self.storage.arg_types.push(PTR_TYPE);
self.start_block(BlockType::NoResult); // block to ensure all paths pop stack memory (if any)
builder::signature()
} else {
@ -129,10 +105,11 @@ impl<'a> WasmBackend<'a> {
};
for (layout, symbol) in proc.args {
self.create_storage(&WasmLayout::new(layout), *symbol, LocalKind::Parameter);
self.storage
.allocate(&WasmLayout::new(layout), *symbol, LocalKind::Parameter);
}
signature_builder.with_params(self.arg_types.clone())
signature_builder.with_params(self.storage.arg_types.clone())
}
fn finalize_proc(&mut self, signature_builder: SignatureBuilder) -> FunctionDefinition {
@ -140,32 +117,32 @@ impl<'a> WasmBackend<'a> {
let mut final_instructions = Vec::with_capacity(self.code_builder.len() + 10);
if self.stack_memory > 0 {
if self.storage.stack_memory > 0 {
push_stack_frame(
&mut final_instructions,
self.stack_memory,
self.stack_frame_pointer.unwrap(),
self.storage.stack_memory,
self.storage.stack_frame_pointer.unwrap(),
);
}
self.code_builder.finalize_into(&mut final_instructions);
if self.stack_memory > 0 {
if self.storage.stack_memory > 0 {
pop_stack_frame(
&mut final_instructions,
self.stack_memory,
self.stack_frame_pointer.unwrap(),
self.storage.stack_memory,
self.storage.stack_frame_pointer.unwrap(),
);
}
final_instructions.push(End);
// Declare local variables (in batches of the same type)
let num_locals = self.local_types.len();
let num_locals = self.storage.local_types.len();
let mut locals = Vec::with_capacity(num_locals);
if num_locals > 0 {
let mut batch_type = self.local_types[0];
let mut batch_type = self.storage.local_types[0];
let mut batch_size = 0;
for t in &self.local_types {
for t in &self.storage.local_types {
if *t == batch_type {
batch_size += 1;
} else {
@ -186,340 +163,6 @@ impl<'a> WasmBackend<'a> {
.build() // function
}
/**********************************************************
SYMBOL STORAGE, LOCALS, AND COPYING
***********************************************************/
fn get_next_local_id(&self) -> LocalId {
LocalId((self.arg_types.len() + self.local_types.len()) as u32)
}
fn create_storage(
&mut self,
wasm_layout: &WasmLayout,
symbol: Symbol,
kind: LocalKind,
) -> StoredValue {
let next_local_id = self.get_next_local_id();
let storage = match wasm_layout {
WasmLayout::Primitive(value_type, size) => match kind {
LocalKind::Parameter => {
self.arg_types.push(*value_type);
StoredValue::Local {
local_id: next_local_id,
value_type: *value_type,
size: *size,
}
}
LocalKind::Variable => StoredValue::VirtualMachineStack {
vm_state: VirtualMachineSymbolState::NotYetPushed,
value_type: *value_type,
size: *size,
},
},
WasmLayout::HeapMemory => {
match kind {
LocalKind::Parameter => self.arg_types.push(PTR_TYPE),
LocalKind::Variable => self.local_types.push(PTR_TYPE),
}
StoredValue::Local {
local_id: next_local_id,
value_type: PTR_TYPE,
size: PTR_SIZE,
}
}
WasmLayout::StackMemory {
size,
alignment_bytes,
} => {
let location = match kind {
LocalKind::Parameter => {
self.arg_types.push(PTR_TYPE);
StackMemoryLocation::PointerArg(next_local_id)
}
LocalKind::Variable => {
if self.stack_frame_pointer.is_none() {
self.stack_frame_pointer = Some(next_local_id);
self.local_types.push(PTR_TYPE);
}
let offset =
round_up_to_alignment(self.stack_memory, *alignment_bytes as i32);
self.stack_memory = offset + (*size as i32);
StackMemoryLocation::FrameOffset(offset as u32)
}
};
StoredValue::StackMemory {
location,
size: *size,
alignment_bytes: *alignment_bytes,
}
}
};
self.symbol_storage_map.insert(symbol, storage.clone());
storage
}
fn get_symbol_storage(&self, sym: &Symbol) -> &StoredValue {
self.symbol_storage_map.get(sym).unwrap_or_else(|| {
panic!(
"Symbol {:?} not found in function scope:\n{:?}",
sym, self.symbol_storage_map
)
})
}
/// Load symbols to the top of the VM stack
/// (There is no method for one symbol. This is deliberate, since
/// if anyone ever called it in a loop, it would generate inefficient code)
fn load_symbols(&mut self, symbols: &[Symbol]) {
if self.code_builder.verify_stack_match(symbols) {
// The symbols were already at the top of the stack, do nothing!
// This should be quite common due to the structure of the Mono IR
return;
}
for sym in symbols.iter() {
let storage = self.get_symbol_storage(sym).to_owned();
match storage {
StoredValue::VirtualMachineStack {
vm_state,
value_type,
size,
} => {
let next_local_id = self.get_next_local_id();
let maybe_next_vm_state =
self.code_builder.load_symbol(*sym, vm_state, next_local_id);
match maybe_next_vm_state {
// The act of loading the value changed the VM state, so update it
Some(next_vm_state) => {
self.symbol_storage_map.insert(
*sym,
StoredValue::VirtualMachineStack {
vm_state: next_vm_state,
value_type,
size,
},
);
}
None => {
// Loading the value required creating a new local, because
// it was not in a convenient position in the VM stack.
self.local_types.push(value_type);
self.symbol_storage_map.insert(
*sym,
StoredValue::Local {
local_id: next_local_id,
value_type,
size,
},
);
}
}
}
StoredValue::Local { local_id, .. }
| StoredValue::StackMemory {
location: StackMemoryLocation::PointerArg(local_id),
..
} => {
self.code_builder.add_one(GetLocal(local_id.0));
self.code_builder.set_top_symbol(*sym);
}
StoredValue::StackMemory {
location: StackMemoryLocation::FrameOffset(offset),
..
} => {
self.code_builder.add_many(&[
GetLocal(self.stack_frame_pointer.unwrap().0),
I32Const(offset as i32),
I32Add,
]);
self.code_builder.set_top_symbol(*sym);
}
}
}
}
fn copy_symbol_to_memory(
&mut self,
to_ptr: LocalId,
to_offset: u32,
from_symbol: Symbol,
) -> u32 {
let from_storage = self.get_symbol_storage(&from_symbol).to_owned();
match from_storage {
StoredValue::StackMemory {
location,
size,
alignment_bytes,
} => {
let (from_ptr, from_offset) = location.local_and_offset(self.stack_frame_pointer);
copy_memory(
&mut self.code_builder,
CopyMemoryConfig {
from_ptr,
from_offset,
to_ptr,
to_offset,
size,
alignment_bytes,
},
);
size
}
StoredValue::VirtualMachineStack {
value_type, size, ..
}
| StoredValue::Local {
value_type, size, ..
} => {
let store_instruction = match (value_type, size) {
(ValueType::I64, 8) => I64Store(ALIGN_8, to_offset),
(ValueType::I32, 4) => I32Store(ALIGN_4, to_offset),
(ValueType::I32, 2) => I32Store16(ALIGN_2, to_offset),
(ValueType::I32, 1) => I32Store8(ALIGN_1, to_offset),
(ValueType::F32, 4) => F32Store(ALIGN_4, to_offset),
(ValueType::F64, 8) => F64Store(ALIGN_8, to_offset),
_ => {
panic!("Cannot store {:?} with alignment of {:?}", value_type, size);
}
};
self.code_builder.add_one(GetLocal(to_ptr.0));
self.load_symbols(&[from_symbol]);
self.code_builder.add_one(store_instruction);
size
}
}
}
/// generate code to copy a value from one SymbolStorage to another
pub fn copy_value_by_storage(
&mut self,
to: &StoredValue,
from: &StoredValue,
from_symbol: Symbol,
) {
use StoredValue::*;
match (to, from) {
(
Local {
local_id: to_local_id,
value_type: to_value_type,
size: to_size,
},
VirtualMachineStack {
value_type: from_value_type,
size: from_size,
..
},
) => {
debug_assert!(to_value_type == from_value_type);
debug_assert!(to_size == from_size);
self.load_symbols(&[from_symbol]);
self.code_builder.add_one(SetLocal(to_local_id.0));
self.symbol_storage_map.insert(from_symbol, to.clone());
}
(
Local {
local_id: to_local_id,
value_type: to_value_type,
size: to_size,
},
Local {
local_id: from_local_id,
value_type: from_value_type,
size: from_size,
},
) => {
debug_assert!(to_value_type == from_value_type);
debug_assert!(to_size == from_size);
self.code_builder
.add_many(&[GetLocal(from_local_id.0), SetLocal(to_local_id.0)]);
}
(
StackMemory {
location: to_location,
size: to_size,
alignment_bytes: to_alignment_bytes,
},
StackMemory {
location: from_location,
size: from_size,
alignment_bytes: from_alignment_bytes,
},
) => {
let (from_ptr, from_offset) =
from_location.local_and_offset(self.stack_frame_pointer);
let (to_ptr, to_offset) = to_location.local_and_offset(self.stack_frame_pointer);
debug_assert!(*to_size == *from_size);
debug_assert!(*to_alignment_bytes == *from_alignment_bytes);
copy_memory(
&mut self.code_builder,
CopyMemoryConfig {
from_ptr,
from_offset,
to_ptr,
to_offset,
size: *from_size,
alignment_bytes: *from_alignment_bytes,
},
);
}
_ => {
panic!("Cannot copy storage from {:?} to {:?}", from, to);
}
}
}
/// Ensure SymbolStorage has an associated local.
/// (Blocks can't access the VM stack of their parent scope, but they can access locals.)
fn ensure_symbol_storage_has_local(
&mut self,
symbol: Symbol,
storage: StoredValue,
) -> StoredValue {
if let StoredValue::VirtualMachineStack {
vm_state,
value_type,
size,
} = storage
{
let local_id = self.get_next_local_id();
if vm_state != VirtualMachineSymbolState::NotYetPushed {
self.code_builder.load_symbol(symbol, vm_state, local_id);
self.code_builder.add_one(SetLocal(local_id.0));
}
self.local_types.push(value_type);
let new_storage = StoredValue::Local {
local_id,
value_type,
size,
};
self.symbol_storage_map.insert(symbol, new_storage.clone());
return new_storage;
} else {
storage
}
}
/**********************************************************
STATEMENTS
@ -561,14 +204,14 @@ impl<'a> WasmBackend<'a> {
size,
alignment_bytes,
};
self.symbol_storage_map.insert(*let_sym, storage);
self.storage.symbol_storage_map.insert(*let_sym, storage);
}
self.build_expr(let_sym, expr, layout)?;
if let WasmLayout::Primitive(value_type, size) = wasm_layout {
let vm_state = self.code_builder.set_top_symbol(*let_sym);
self.symbol_storage_map.insert(
self.storage.symbol_storage_map.insert(
*let_sym,
StoredValue::VirtualMachineStack {
vm_state,
@ -585,12 +228,13 @@ impl<'a> WasmBackend<'a> {
Stmt::Let(sym, expr, layout, following) => {
let wasm_layout = WasmLayout::new(layout);
self.create_storage(&wasm_layout, *sym, LocalKind::Variable);
self.storage
.allocate(&wasm_layout, *sym, LocalKind::Variable);
self.build_expr(sym, expr, layout)?;
if let WasmLayout::Primitive(value_type, size) = wasm_layout {
let vm_state = self.code_builder.set_top_symbol(*sym);
self.symbol_storage_map.insert(
self.storage.symbol_storage_map.insert(
*sym,
StoredValue::VirtualMachineStack {
vm_state,
@ -607,7 +251,7 @@ impl<'a> WasmBackend<'a> {
Stmt::Ret(sym) => {
use crate::storage::StoredValue::*;
let storage = self.symbol_storage_map.get(sym).unwrap();
let storage = self.storage.symbol_storage_map.get(sym).unwrap();
match storage {
StackMemory {
@ -616,7 +260,7 @@ impl<'a> WasmBackend<'a> {
alignment_bytes,
} => {
let (from_ptr, from_offset) =
location.local_and_offset(self.stack_frame_pointer);
location.local_and_offset(self.storage.stack_frame_pointer);
copy_memory(
&mut self.code_builder,
CopyMemoryConfig {
@ -631,7 +275,7 @@ impl<'a> WasmBackend<'a> {
}
_ => {
self.load_symbols(&[*sym]);
self.storage.load_symbols(&mut self.code_builder, &[*sym]);
self.code_builder.add_one(Br(self.block_depth)); // jump to end of function (for stack frame pop)
}
}
@ -652,8 +296,12 @@ impl<'a> WasmBackend<'a> {
// Ensure the condition value is not stored only in the VM stack
// Otherwise we can't reach it from inside the block
let cond_storage = self.get_symbol_storage(cond_symbol).to_owned();
self.ensure_symbol_storage_has_local(*cond_symbol, cond_storage);
let cond_storage = self.storage.get(cond_symbol).to_owned();
self.storage.ensure_symbol_storage_has_local(
&mut self.code_builder,
*cond_symbol,
cond_storage,
);
// create (number_of_branches - 1) new blocks.
for _ in 0..branches.len() {
@ -663,7 +311,8 @@ impl<'a> WasmBackend<'a> {
// then, we jump whenever the value under scrutiny is equal to the value of a branch
for (i, (value, _, _)) in branches.iter().enumerate() {
// put the cond_symbol on the top of the stack
self.load_symbols(&[*cond_symbol]);
self.storage
.load_symbols(&mut self.code_builder, &[*cond_symbol]);
self.code_builder.add_one(I32Const(*value as i32));
@ -699,9 +348,13 @@ impl<'a> WasmBackend<'a> {
for parameter in parameters.iter() {
let wasm_layout = WasmLayout::new(&parameter.layout);
let mut param_storage =
self.create_storage(&wasm_layout, parameter.symbol, LocalKind::Variable);
param_storage =
self.ensure_symbol_storage_has_local(parameter.symbol, param_storage);
self.storage
.allocate(&wasm_layout, parameter.symbol, LocalKind::Variable);
param_storage = self.storage.ensure_symbol_storage_has_local(
&mut self.code_builder,
parameter.symbol,
param_storage,
);
jp_param_storages.push(param_storage);
}
@ -730,8 +383,13 @@ impl<'a> WasmBackend<'a> {
let (target, param_storages) = self.joinpoint_label_map[id].clone();
for (arg_symbol, param_storage) in arguments.iter().zip(param_storages.iter()) {
let arg_storage = self.get_symbol_storage(arg_symbol).clone();
self.copy_value_by_storage(&param_storage, &arg_storage, *arg_symbol);
let arg_storage = self.storage.get(arg_symbol).clone();
self.storage.copy_value_by_storage(
&mut self.code_builder,
&param_storage,
&arg_storage,
*arg_symbol,
);
}
// jump
@ -764,7 +422,8 @@ impl<'a> WasmBackend<'a> {
arguments,
}) => match call_type {
CallType::ByName { name: func_sym, .. } => {
self.load_symbols(*arguments);
self.storage
.load_symbols(&mut self.code_builder, *arguments);
let function_location = self.proc_symbol_map.get(func_sym).ok_or(format!(
"Cannot find function {:?} called from {:?}",
func_sym, sym
@ -831,18 +490,22 @@ impl<'a> WasmBackend<'a> {
) -> Result<(), String> {
// TODO: we just calculated storage and now we're getting it out of a map
// Not passing it as an argument because I'm trying to match Backend method signatures
let storage = self.get_symbol_storage(sym).to_owned();
let storage = self.storage.get(sym).to_owned();
if let Layout::Struct(field_layouts) = layout {
match storage {
StoredValue::StackMemory { location, size, .. } => {
if size > 0 {
let (local_id, struct_offset) =
location.local_and_offset(self.stack_frame_pointer);
location.local_and_offset(self.storage.stack_frame_pointer);
let mut field_offset = struct_offset;
for (field, _) in fields.iter().zip(field_layouts.iter()) {
field_offset +=
self.copy_symbol_to_memory(local_id, field_offset, *field);
field_offset += self.storage.copy_symbol_to_memory(
&mut self.code_builder,
local_id,
field_offset,
*field,
);
}
} else {
return Err(format!("Not supported yet: zero-size struct at {:?}", sym));
@ -857,8 +520,13 @@ impl<'a> WasmBackend<'a> {
};
} else {
// Struct expression but not Struct layout => single element. Copy it.
let field_storage = self.get_symbol_storage(&fields[0]).to_owned();
self.copy_value_by_storage(&storage, &field_storage, fields[0]);
let field_storage = self.storage.get(&fields[0]).to_owned();
self.storage.copy_value_by_storage(
&mut self.code_builder,
&storage,
&field_storage,
fields[0],
);
}
Ok(())
}
@ -869,7 +537,7 @@ impl<'a> WasmBackend<'a> {
args: &'a [Symbol],
return_layout: &Layout<'a>,
) -> Result<(), String> {
self.load_symbols(args);
self.storage.load_symbols(&mut self.code_builder, args);
let wasm_layout = WasmLayout::new(return_layout);
self.build_instructions_lowlevel(lowlevel, wasm_layout.value_type())?;
Ok(())