update gen_wasm to use bumpalo::collections::Vec where possible

compiler/gen_wasm/README.md

@@ -16,7 +16,7 @@
     - [x] Distinguish which variables go in locals, own stack frame, caller stack frame, etc.
     - [x] Ensure early Return statements don't skip stack cleanup
   - [x] Model the stack machine as a storage mechanism, to make generated code "less bad"
-  - [ ] Vendor-in parity_wasm library so that we can use `bumpalo::Vec`
+  - [x] Switch vectors to `bumpalo::Vec` where possible
   - [ ] Implement relocations
     - Requires knowing the _byte_ offset of each call site. This is awkward as the backend builds a `Vec<Instruction>` rather than a `Vec<u8>`. It may be worth serialising each instruction as it is inserted.
 

compiler/gen_wasm/src/backend.rs

@@ -1,3 +1,4 @@
+use bumpalo::collections::Vec;
 use parity_wasm::builder;
 use parity_wasm::builder::{CodeLocation, FunctionDefinition, ModuleBuilder, SignatureBuilder};
 use parity_wasm::elements::{
@@ -12,8 +13,10 @@ use roc_mono::layout::{Builtin, Layout};
 
 use crate::code_builder::CodeBuilder;
 use crate::layout::WasmLayout;
-use crate::storage::{Storage, StoredValue, StoredValueKind};
-use crate::{copy_memory, pop_stack_frame, push_stack_frame, CopyMemoryConfig, LocalId, PTR_TYPE};
+use crate::storage::{StackMemoryLocation, Storage, StoredValue, StoredValueKind};
+use crate::{
+    copy_memory, pop_stack_frame, push_stack_frame, CopyMemoryConfig, Env, 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)
@@ -26,6 +29,7 @@ struct LabelId(u32);
 pub struct WasmBackend<'a> {
     // Module level: Wasm AST
     pub module_builder: ModuleBuilder,
+    env: &'a Env<'a>,
 
     // Module level: internal state & IR mappings
     _data_offset_map: MutMap<Literal<'a>, u32>,
@@ -33,19 +37,20 @@ pub struct WasmBackend<'a> {
     proc_symbol_map: MutMap<Symbol, CodeLocation>,
 
     // Function level
-    code_builder: CodeBuilder,
-    storage: Storage,
+    code_builder: CodeBuilder<'a>,
+    storage: Storage<'a>,
 
     /// how many blocks deep are we (used for jumps)
     block_depth: u32,
-    joinpoint_label_map: MutMap<JoinPointId, (u32, std::vec::Vec<StoredValue>)>,
+    joinpoint_label_map: MutMap<JoinPointId, (u32, Vec<'a, StoredValue>)>,
 }
 
 impl<'a> WasmBackend<'a> {
-    pub fn new() -> Self {
+    pub fn new(env: &'a Env<'a>) -> Self {
         WasmBackend {
             // Module: Wasm AST
             module_builder: builder::module(),
+            env,
 
             // Module: internal state & IR mappings
             _data_offset_map: MutMap::default(),
@@ -56,8 +61,8 @@ impl<'a> WasmBackend<'a> {
             joinpoint_label_map: MutMap::default(),
 
             // Functions
-            code_builder: CodeBuilder::new(),
-            storage: Storage::new(),
+            code_builder: CodeBuilder::new(env.arena),
+            storage: Storage::new(env.arena),
         }
     }
 
@@ -120,7 +125,7 @@ impl<'a> WasmBackend<'a> {
 
         const STACK_FRAME_INSTRUCTIONS_LEN: usize = 10;
         let mut final_instructions =
-            Vec::with_capacity(self.code_builder.len() + STACK_FRAME_INSTRUCTIONS_LEN);
+            std::vec::Vec::with_capacity(self.code_builder.len() + STACK_FRAME_INSTRUCTIONS_LEN);
 
         if self.storage.stack_frame_size > 0 {
             push_stack_frame(
@@ -143,7 +148,7 @@ impl<'a> WasmBackend<'a> {
 
         // Declare local variables (in batches of the same type)
         let num_locals = self.storage.local_types.len();
-        let mut locals = Vec::with_capacity(num_locals);
+        let mut locals = Vec::with_capacity_in(num_locals, self.env.arena);
         if num_locals > 0 {
             let mut batch_type = self.storage.local_types[0];
             let mut batch_size = 0;
@@ -318,7 +323,7 @@ impl<'a> WasmBackend<'a> {
                 remainder,
             } => {
                 // make locals for join pointer parameters
-                let mut jp_param_storages = std::vec::Vec::with_capacity(parameters.len());
+                let mut jp_param_storages = Vec::with_capacity_in(parameters.len(), self.env.arena);
                 for parameter in parameters.iter() {
                     let wasm_layout = WasmLayout::new(&parameter.layout);
                     let mut param_storage = self.storage.allocate(
@@ -406,7 +411,8 @@ impl<'a> WasmBackend<'a> {
                     let mut wasm_args_tmp: Vec<Symbol>;
                     let (wasm_args, has_return_val) = match wasm_layout {
                         WasmLayout::StackMemory { .. } => {
-                            wasm_args_tmp = Vec::with_capacity(arguments.len() + 1); // TODO: bumpalo
+                            wasm_args_tmp =
+                                Vec::with_capacity_in(arguments.len() + 1, self.env.arena);
                             wasm_args_tmp.push(*sym);
                             wasm_args_tmp.extend_from_slice(*arguments);
                             (wasm_args_tmp.as_slice(), false)

compiler/gen_wasm/src/code_builder.rs

@@ -1,3 +1,5 @@
+use bumpalo::collections::Vec;
+use bumpalo::Bump;
 use core::panic;
 use std::collections::BTreeMap;
 use std::fmt::Debug;
@@ -25,9 +27,9 @@ pub enum VirtualMachineSymbolState {
 }
 
 #[derive(Debug)]
-pub struct CodeBuilder {
+pub struct CodeBuilder<'a> {
     /// The main container for the instructions
-    code: Vec<Instruction>,
+    code: Vec<'a, Instruction>,
 
     /// Extra instructions to insert at specific positions during finalisation
     /// (Go back and set locals when we realise we need them)
@@ -39,16 +41,16 @@ pub struct CodeBuilder {
 
     /// Our simulation model of the Wasm stack machine
     /// Keeps track of where Symbol values are in the VM stack
-    vm_stack: Vec<Symbol>,
+    vm_stack: Vec<'a, Symbol>,
 }
 
 #[allow(clippy::new_without_default)]
-impl CodeBuilder {
-    pub fn new() -> Self {
+impl<'a> CodeBuilder<'a> {
+    pub fn new(arena: &'a Bump) -> Self {
         CodeBuilder {
-            vm_stack: Vec::with_capacity(32),
+            vm_stack: Vec::with_capacity_in(32, arena),
             insertions: BTreeMap::default(),
-            code: Vec::with_capacity(1024),
+            code: Vec::with_capacity_in(1024, arena),
         }
     }
 
@@ -101,7 +103,8 @@ impl CodeBuilder {
     pub fn push_call(&mut self, function_index: u32, pops: usize, push: bool) {
         let stack_depth = self.vm_stack.len();
         if pops > stack_depth {
-            let mut final_code = Vec::with_capacity(self.code.len() + self.insertions.len());
+            let mut final_code =
+                std::vec::Vec::with_capacity(self.code.len() + self.insertions.len());
             self.finalize_into(&mut final_code);
             panic!(
                 "Trying to call to call function {:?} with {:?} values but only {:?} on the VM stack\nfinal_code={:?}\nvm_stack={:?}",
@@ -120,7 +123,7 @@ impl CodeBuilder {
     }
 
     /// Finalize a function body by copying all instructions into a vector
-    pub fn finalize_into(&mut self, final_code: &mut Vec<Instruction>) {
+    pub fn finalize_into(&mut self, final_code: &mut std::vec::Vec<Instruction>) {
         let mut insertions_iter = self.insertions.iter();
         let mut next_insertion = insertions_iter.next();
 

compiler/gen_wasm/src/lib.rs

@@ -4,6 +4,7 @@ pub mod from_wasm32_memory;
 mod layout;
 mod storage;
 
+use bumpalo::collections::Vec;
 use bumpalo::Bump;
 use parity_wasm::builder;
 use parity_wasm::elements::{Instruction, Instruction::*, Internal, ValueType};
@@ -32,7 +33,7 @@ pub const STACK_ALIGNMENT_BYTES: i32 = 16;
 pub struct LocalId(pub u32);
 
 pub struct Env<'a> {
-    pub arena: &'a Bump, // not really using this much, parity_wasm works with std::vec a lot
+    pub arena: &'a Bump,
     pub interns: Interns,
     pub exposed_to_host: MutSet<Symbol>,
 }
@@ -40,7 +41,7 @@ pub struct Env<'a> {
 pub fn build_module<'a>(
     env: &'a Env,
     procedures: MutMap<(Symbol, ProcLayout<'a>), Proc<'a>>,
-) -> Result<Vec<u8>, String> {
+) -> Result<std::vec::Vec<u8>, String> {
     let (builder, _) = build_module_help(env, procedures)?;
     let module = builder.build();
     module
@@ -52,7 +53,7 @@ pub fn build_module_help<'a>(
     env: &'a Env,
     procedures: MutMap<(Symbol, ProcLayout<'a>), Proc<'a>>,
 ) -> Result<(builder::ModuleBuilder, u32), String> {
-    let mut backend = WasmBackend::new();
+    let mut backend = WasmBackend::new(env);
     let mut layout_ids = LayoutIds::default();
 
     // Sort procedures by occurrence order
@@ -65,7 +66,7 @@ pub fn build_module_help<'a>(
     //
     // This means that for now other functions in the file have to be ordered "in reverse": if A
     // uses B, then the name of A must first occur after the first occurrence of the name of B
-    let mut procedures: std::vec::Vec<_> = procedures.into_iter().collect();
+    let mut procedures = Vec::from_iter_in(procedures.into_iter(), env.arena);
     procedures.sort_by(|a, b| b.0 .0.cmp(&a.0 .0));
 
     let mut function_index: u32 = 0;
@@ -178,7 +179,7 @@ pub fn round_up_to_alignment(unaligned: i32, alignment_bytes: i32) -> i32 {
 }
 
 pub fn push_stack_frame(
-    instructions: &mut Vec<Instruction>,
+    instructions: &mut std::vec::Vec<Instruction>,
     size: i32,
     local_frame_pointer: LocalId,
 ) {
@@ -193,7 +194,7 @@ pub fn push_stack_frame(
 }
 
 pub fn pop_stack_frame(
-    instructions: &mut Vec<Instruction>,
+    instructions: &mut std::vec::Vec<Instruction>,
     size: i32,
     local_frame_pointer: LocalId,
 ) {

compiler/gen_wasm/src/storage.rs

@@ -1,3 +1,5 @@
+use bumpalo::collections::Vec;
+use bumpalo::Bump;
 use parity_wasm::elements::{Instruction::*, ValueType};
 
 use roc_collections::all::MutMap;
@@ -58,19 +60,19 @@ pub enum StoredValue {
 
 /// Helper structure for WasmBackend, to keep track of how values are stored,
 /// including the VM stack, local variables, and linear memory
-pub struct Storage {
-    pub arg_types: std::vec::Vec<ValueType>,
-    pub local_types: std::vec::Vec<ValueType>,
+pub struct Storage<'a> {
+    pub arg_types: Vec<'a, ValueType>,
+    pub local_types: Vec<'a, ValueType>,
     pub symbol_storage_map: MutMap<Symbol, StoredValue>,
     pub stack_frame_pointer: Option<LocalId>,
     pub stack_frame_size: i32,
 }
 
-impl Storage {
-    pub fn new() -> Self {
+impl<'a> Storage<'a> {
+    pub fn new(arena: &'a Bump) -> Self {
         Storage {
-            arg_types: std::vec::Vec::with_capacity(8),
-            local_types: std::vec::Vec::with_capacity(32),
+            arg_types: Vec::with_capacity_in(8, arena),
+            local_types: Vec::with_capacity_in(32, arena),
             symbol_storage_map: MutMap::default(),
             stack_frame_pointer: None,
             stack_frame_size: 0,

compiler/gen_wasm/tests/helpers/eval.rs

@@ -106,7 +106,7 @@ pub fn helper_wasm<'a, T: Wasm32TestResult>(
         roc_gen_wasm::build_module_help(&env, procedures).unwrap();
     T::insert_test_wrapper(&mut builder, TEST_WRAPPER_NAME, main_function_index);
 
-    let module_bytes = builder.build().to_bytes().unwrap();
+    let module_bytes = builder.build().into_bytes().unwrap();
 
     // for debugging (e.g. with wasm2wat)
     if false {
@@ -190,7 +190,6 @@ macro_rules! assert_wasm_evals_to {
         match $crate::helpers::eval::assert_wasm_evals_to_help::<$ty>($src, $expected) {
             Err(msg) => println!("{:?}", msg),
             Ok(actual) => {
-                #[allow(clippy::bool_assert_comparison)]
                 assert_eq!($transform(actual), $expected)
             }
         }