Expose a FunctionBuilder through LocalFunction

This should allow to continue to build expressions for a function even
after it exists

crates/macro/src/lib.rs

@@ -752,7 +752,7 @@ fn create_matchers(variants: &[WalrusVariant]) -> impl quote::ToTokens {
                                 true #(
                                     && #self_args.is_match(
                                         local_func,
-                                        &local_func.exprs[*#args]
+                                        local_func.get(*#args)
                                     )
                                 )*
                             }

src/function_builder.rs

@@ -7,7 +7,7 @@ use std::ops::{Deref, DerefMut, Drop};
 /// A helpful struct used for building instances of `LocalFunction`
 #[derive(Default, Debug)]
 pub struct FunctionBuilder {
-    arena: Arena<Expr>,
+    pub(crate) arena: Arena<Expr>,
 }
 
 impl FunctionBuilder {
@@ -127,7 +127,7 @@ impl FunctionBuilder {
             results: ty.results().to_vec().into_boxed_slice(),
             exprs,
         });
-        let func = LocalFunction::new(ty_id, args, self.arena, entry);
+        let func = LocalFunction::new(ty_id, args, self, entry);
         module.funcs.add_local(func)
     }
 }

src/ir/mod.rs

@@ -990,7 +990,7 @@ impl<'expr> Visit<'expr> for ExprId {
     where
         V: Visitor<'expr>,
     {
-        visitor.visit_expr(&visitor.local_function().exprs[*self])
+        visitor.visit_expr(&visitor.local_function().get(*self))
     }
 }
 
@@ -1004,11 +1004,11 @@ impl VisitMut for ExprId {
         // use graph traversal instead of a simple tree walk like we have today
         // when that comes about.
         let mut expr = mem::replace(
-            &mut visitor.local_function_mut().exprs[*self],
+            visitor.local_function_mut().get_mut(*self),
             Expr::Unreachable(Unreachable {}),
         );
         visitor.visit_expr_mut(&mut expr);
-        visitor.local_function_mut().exprs[*self] = expr;
+        *visitor.local_function_mut().get_mut(*self) = expr;
     }
 }
 

src/lib.rs

@@ -18,7 +18,7 @@ pub mod passes;
 mod ty;
 
 pub use crate::error::{ErrorKind, Result};
-pub use crate::function_builder::FunctionBuilder;
+pub use crate::function_builder::{FunctionBuilder, BlockBuilder};
 pub use crate::init_expr::InitExpr;
 pub use crate::ir::{Local, LocalId};
 pub use crate::module::*;

src/module/functions/local_function/context.rs

@@ -170,7 +170,7 @@ impl<'a> FunctionContext<'a> {
                 }
             }
 
-            let block = self.func.exprs[frame.block.into()].unwrap_block_mut();
+            let block = self.func.block_mut(frame.block);
             block.exprs.extend(extra_exprs);
             block.exprs.push(expr);
         }
@@ -192,8 +192,7 @@ impl<'a> FunctionContext<'a> {
     where
         E: Into<ExprId>,
     {
-        let block = self.func.exprs[block.into()].unwrap_block_mut();
-        block.exprs.push(expr.into());
+        self.func.block_mut(block).exprs.push(expr.into());
     }
 
     pub fn add_to_frame_block<E>(&mut self, control_frame: usize, expr: E)
@@ -228,8 +227,8 @@ impl<'a> FunctionContext<'a> {
 
     pub fn add_side_effect(&mut self, value: ExprId, side_effect: ExprId) -> ExprId {
         // If we can add it to an existing `WithSideEffects` expr for this value, then do that.
-        if let Some(Expr::WithSideEffects(WithSideEffects { side_effects, .. })) =
-            self.func.exprs.get_mut(value)
+        if let Expr::WithSideEffects(WithSideEffects { side_effects, .. }) =
+            self.func.get_mut(value)
         {
             side_effects.push(side_effect);
             return value;

src/module/functions/local_function/emit.rs

@@ -56,7 +56,7 @@ impl Emit<'_, '_> {
         let old = self.id;
         self.id = id;
 
-        match &self.func.exprs[id] {
+        match &self.func.get(id) {
             Const(e) => e.value.emit(self.encoder),
             Block(e) => self.visit_block(e),
             BrTable(e) => self.visit_br_table(e),

src/module/functions/local_function/mod.rs

@@ -13,9 +13,9 @@ use crate::ir::*;
 use crate::map::{IdHashMap, IdHashSet};
 use crate::parse::IndicesToIds;
 use crate::passes::Used;
-use crate::{FunctionId, Module, Result, TypeId, ValType, TableKind};
+use crate::{FunctionId, Module, Result, TypeId, ValType, FunctionBuilder, TableKind};
 use failure::{bail, ResultExt};
-use id_arena::{Arena, Id};
+use id_arena::Id;
 use std::collections::BTreeMap;
 use std::fmt;
 use std::mem;
@@ -27,8 +27,8 @@ pub struct LocalFunction {
     /// This function's type.
     pub ty: TypeId,
 
-    /// The arena that contains this function's expressions.
-    pub(crate) exprs: Arena<Expr>,
+    /// All of this function's expressions, contained in the arena.
+    exprs: FunctionBuilder,
 
     /// Arguments to this function, and the locals that they're assigned to
     pub args: Vec<LocalId>,
@@ -46,7 +46,7 @@ impl LocalFunction {
     pub(crate) fn new(
         ty: TypeId,
         args: Vec<LocalId>,
-        exprs: Arena<Expr>,
+        exprs: FunctionBuilder,
         entry: BlockId,
     ) -> LocalFunction {
         LocalFunction {
@@ -71,7 +71,7 @@ impl LocalFunction {
     ) -> Result<LocalFunction> {
         let mut func = LocalFunction {
             ty,
-            exprs: Arena::new(),
+            exprs: FunctionBuilder::new(),
             args,
             entry: None,
         };
@@ -101,12 +101,11 @@ impl LocalFunction {
         Ok(func)
     }
 
-    fn alloc<T>(&mut self, val: T) -> T::Id
+    pub(crate) fn alloc<T>(&mut self, val: T) -> T::Id
     where
         T: Ast,
     {
-        let id = self.exprs.alloc(val.into());
-        T::new_id(id)
+        self.exprs.alloc(val)
     }
 
     /// Get the id of this function's entry block.
@@ -116,22 +115,34 @@ impl LocalFunction {
 
     /// Get the block associated with the given id.
     pub fn block(&self, block: BlockId) -> &Block {
-        self.exprs[block.into()].unwrap_block()
+        self.get(block.into()).unwrap_block()
     }
 
     /// Get the block associated with the given id.
     pub fn block_mut(&mut self, block: BlockId) -> &mut Block {
-        self.exprs[block.into()].unwrap_block_mut()
+        self.get_mut(block.into()).unwrap_block_mut()
     }
 
     /// Get the expression associated with the given id
     pub fn get(&self, id: ExprId) -> &Expr {
-        &self.exprs[id]
+        &self.exprs.arena[id]
     }
 
     /// Get the expression associated with the given id
     pub fn get_mut(&mut self, id: ExprId) -> &mut Expr {
-        &mut self.exprs[id]
+        &mut self.exprs.arena[id]
+    }
+
+    /// Get access to a `FunctionBuilder` to continue adding expressions to
+    /// this function.
+    pub fn builder(&self) -> &FunctionBuilder {
+        &self.exprs
+    }
+
+    /// Get access to a `FunctionBuilder` to continue adding expressions to
+    /// this function.
+    pub fn builder_mut(&mut self) -> &mut FunctionBuilder {
+        &mut self.exprs
     }
 
     /// Get the size of this function, in number of expressions.
@@ -163,15 +174,12 @@ impl LocalFunction {
     /// Is this function's body a [constant
     /// expression](https://webassembly.github.io/spec/core/valid/instructions.html#constant-expressions)?
     pub fn is_const(&self) -> bool {
-        let entry = match &self.exprs[self.entry_block().into()] {
-            Expr::Block(b) => b,
-            _ => unreachable!(),
-        };
+        let entry = self.block(self.entry_block());
         let matcher = ConstMatcher::new();
         entry
             .exprs
             .iter()
-            .all(|e| matcher.is_match(self, &self.exprs[*e]))
+            .all(|e| matcher.is_match(self, &self.get(*e)))
     }
 
     /// Emit this function's compact locals declarations.