Cleanup cranelift-frontend (#4989)

* cranelift-wasm: Assume block is reachable

In handling the WebAssembly "end" operator, cranelift-wasm had logic to
skip generating a jump instruction if the block was both unreachable and
"pristine", meaning no instructions had been added.

However, `translate_operator` checks first that `state.reachable` is
true, so this logic only runs when cranelift-wasm believes that the
current block _is_ reachable. Therefore the condition should always be
true, whether the block is pristine or not.

I've left a debug_assert in case `state.reachable` ever doesn't agree
with `builder.is_unreachable()`, but the assert doesn't fail in any of
the tests. We'll see if fuzzing finds something.

Anyway, outside of cranelift-frontend, this eliminates the only use of
`is_pristine()`, and there were no uses of `is_filled()`. So I've made
both of those private. They're now only used in a nearby debug assert.

* cranelift-frontend: Clarify pristine/filled states

There was a comment here saying "A filled block cannot be pristine."
Given that the intent was for those two states to be mutually exclusive,
I've replaced the two booleans with a three-state enum.

I also replaced all reads of these two flags with method calls. In all
but one case these are only checked in debug assertions, so I don't even
care whether they get inlined. They're easier to read, and this will
make it easier to replace their implementations, which I hope to do
soon.

Finally, I replaced all assignments to either flag with an appropriate
assignment of the corresponding enum state. Keep in mind this
correspondence between the new enum and the old flags:

- Empty: pristine true, filled false
- Partial: pristine false, filled false
- Filled: pristine false, filled true

Every existing update to these flags could only move to a later state.
(For example, Partial couldn't go back to Empty.) In the old flags that
meant that pristine could only go from true to false, and filled could
only go from false to true.

`fill_current_block` was a weird case because at first glance it looks
like it could allow both pristine and filled to be true at the same
time. However, it's only called from `FuncInstBuilder::build`, which
calls `ensure_inserted_block` before doing anything else, and _that_
cleared the pristine flag.

Similarly, `handle_ssa_side_effects` looks like it could allow both
pristine and filled to be true for anything in `split_blocks_created`.
However, those blocks are created by SSABuilder, so their BlockData is
not initialized by `create_block`, and instead uses BlockData::default.
The `Default` implementation here previously set both flags false, while
`create_block` would instead set pristine to true. So these split blocks
were correctly set to the Filled state, and after this patch they are
still set correctly.

* cranelift-frontend: Separate SSA and user block params

Previously there was a `user_param_count` field in BlockData, used
purely to debug-assert that no user parameters are added to a block
after `use_var` adds SSA parameters.

Instead, this patch enforces a strict phase separation between the
period after a block is created when user parameters can be added to it,
and the period when `use_var` may be called and instructions may be
added.

I'm assuming that calls to `use_var` are _always_ followed by inserting
one or more instructions into the block. (If you don't want to insert an
instruction, why do you need to know where instructions in this block
would get variable definitions from?) This patch has no visible effect
for callers which follow that rule.

However, it was previously legal to call `use_var`, then append a block
parameter before adding instructions, so long as `use_var` didn't
actually need to add a block parameter. That could only happen if the
current block is sealed and has exactly one predecessor. So anyone who
was counting on this behavior was playing a dangerous game anyway.

* cranelift-frontend: Defer initializing block data

Every reference to the func_ctx.status SecondaryMap will automatically
create the appropriate entries on-demand, with the sole exception of
`finalize`. In that function, debug assertions use SecondaryMap::keys to
find out which blocks need to be checked.

However, those assertions always succeed for blocks which never had any
instructions added. So it's okay to skip them for blocks which aren't
touched after `create_block`.

Author: jameysharp

cranelift/frontend/src/frontend.rs

@@ -24,7 +24,7 @@ use cranelift_codegen::packed_option::PackedOption;
 #[derive(Default)]
 pub struct FunctionBuilderContext {
     ssa: SSABuilder,
-    blocks: SecondaryMap<Block, BlockData>,
+    status: SecondaryMap<Block, BlockStatus>,
     types: SecondaryMap<Variable, Type>,
 }
 
@@ -41,20 +41,15 @@ pub struct FunctionBuilder<'a> {
     position: PackedOption<Block>,
 }
 
-#[derive(Clone, Default)]
-struct BlockData {
-    /// A Block is "pristine" iff no instructions have been added since the last
-    /// call to `switch_to_block()`.
-    pristine: bool,
-
-    /// A Block is "filled" iff a terminator instruction has been inserted since
-    /// the last call to `switch_to_block()`.
-    ///
-    /// A filled block cannot be pristine.
-    filled: bool,
-
-    /// Count of parameters not supplied implicitly by the SSABuilder.
-    user_param_count: usize,
+#[derive(Clone, Default, Eq, PartialEq)]
+enum BlockStatus {
+    /// No instructions have been added.
+    #[default]
+    Empty,
+    /// Some instructions have been added, but no terminator.
+    Partial,
+    /// A terminator has been added; no further instructions may be added.
+    Filled,
 }
 
 impl FunctionBuilderContext {
@@ -66,12 +61,12 @@ impl FunctionBuilderContext {
 
     fn clear(&mut self) {
         self.ssa.clear();
-        self.blocks.clear();
+        self.status.clear();
         self.types.clear();
     }
 
     fn is_empty(&self) -> bool {
-        self.ssa.is_empty() && self.blocks.is_empty() && self.types.is_empty()
+        self.ssa.is_empty() && self.status.is_empty() && self.types.is_empty()
     }
 }
 
@@ -304,11 +299,6 @@ impl<'a> FunctionBuilder<'a> {
     pub fn create_block(&mut self) -> Block {
         let block = self.func.dfg.make_block();
         self.func_ctx.ssa.declare_block(block);
-        self.func_ctx.blocks[block] = BlockData {
-            filled: false,
-            pristine: true,
-            user_param_count: 0,
-        };
         block
     }
 
@@ -337,13 +327,13 @@ impl<'a> FunctionBuilder<'a> {
         debug_assert!(
             self.position.is_none()
                 || self.is_unreachable()
-                || self.is_pristine()
-                || self.is_filled(),
+                || self.is_pristine(self.position.unwrap())
+                || self.is_filled(self.position.unwrap()),
             "you have to fill your block before switching"
         );
         // We cannot switch to a filled block
         debug_assert!(
-            !self.func_ctx.blocks[block].filled,
+            !self.is_filled(block),
             "you cannot switch to a block which is already filled"
         );
 
@@ -393,6 +383,12 @@ impl<'a> FunctionBuilder<'a> {
     /// Returns the Cranelift IR necessary to use a previously defined user
     /// variable, returning an error if this is not possible.
     pub fn try_use_var(&mut self, var: Variable) -> Result<Value, UseVariableError> {
+        // Assert that we're about to add instructions to this block using the definition of the
+        // given variable. ssa.use_var is the only part of this crate which can add block parameters
+        // behind the caller's back. If we disallow calling append_block_param as soon as use_var is
+        // called, then we enforce a strict separation between user parameters and SSA parameters.
+        self.ensure_inserted_block();
+
         let (val, side_effects) = {
             let ty = *self
                 .func_ctx
@@ -534,14 +530,14 @@ impl<'a> FunctionBuilder<'a> {
     /// Make sure that the current block is inserted in the layout.
     pub fn ensure_inserted_block(&mut self) {
         let block = self.position.unwrap();
-        if self.func_ctx.blocks[block].pristine {
+        if self.is_pristine(block) {
             if !self.func.layout.is_block_inserted(block) {
                 self.func.layout.append_block(block);
             }
-            self.func_ctx.blocks[block].pristine = false;
+            self.func_ctx.status[block] = BlockStatus::Partial;
         } else {
             debug_assert!(
-                !self.func_ctx.blocks[block].filled,
+                !self.is_filled(block),
                 "you cannot add an instruction to a block already filled"
             );
         }
@@ -569,9 +565,12 @@ impl<'a> FunctionBuilder<'a> {
 
         // These parameters count as "user" parameters here because they aren't
         // inserted by the SSABuilder.
-        let user_param_count = &mut self.func_ctx.blocks[block].user_param_count;
+        debug_assert!(
+            self.is_pristine(block),
+            "You can't add block parameters after adding any instruction"
+        );
+
         for argtyp in &self.func.stencil.signature.params {
-            *user_param_count += 1;
             self.func
                 .stencil
                 .dfg
@@ -585,9 +584,12 @@ impl<'a> FunctionBuilder<'a> {
     pub fn append_block_params_for_function_returns(&mut self, block: Block) {
         // These parameters count as "user" parameters here because they aren't
         // inserted by the SSABuilder.
-        let user_param_count = &mut self.func_ctx.blocks[block].user_param_count;
+        debug_assert!(
+            self.is_pristine(block),
+            "You can't add block parameters after adding any instruction"
+        );
+
         for argtyp in &self.func.stencil.signature.returns {
-            *user_param_count += 1;
             self.func
                 .stencil
                 .dfg
@@ -602,25 +604,27 @@ impl<'a> FunctionBuilder<'a> {
         // Check that all the `Block`s are filled and sealed.
         #[cfg(debug_assertions)]
         {
-            for (block, block_data) in self.func_ctx.blocks.iter() {
-                assert!(
-                    block_data.pristine || self.func_ctx.ssa.is_sealed(block),
-                    "FunctionBuilder finalized, but block {} is not sealed",
-                    block,
-                );
-                assert!(
-                    block_data.pristine || block_data.filled,
-                    "FunctionBuilder finalized, but block {} is not filled",
-                    block,
-                );
+            for block in self.func_ctx.status.keys() {
+                if !self.is_pristine(block) {
+                    assert!(
+                        self.func_ctx.ssa.is_sealed(block),
+                        "FunctionBuilder finalized, but block {} is not sealed",
+                        block,
+                    );
+                    assert!(
+                        self.is_filled(block),
+                        "FunctionBuilder finalized, but block {} is not filled",
+                        block,
+                    );
+                }
             }
         }
 
         // In debug mode, check that all blocks are valid basic blocks.
         #[cfg(debug_assertions)]
         {
             // Iterate manually to provide more helpful error messages.
-            for block in self.func_ctx.blocks.keys() {
+            for block in self.func_ctx.status.keys() {
                 if let Err((inst, msg)) = self.func.is_block_basic(block) {
                     let inst_str = self.func.dfg.display_inst(inst);
                     panic!(
@@ -665,14 +669,9 @@ impl<'a> FunctionBuilder<'a> {
     /// instructions to it, otherwise this could interfere with SSA construction.
     pub fn append_block_param(&mut self, block: Block, ty: Type) -> Value {
         debug_assert!(
-            self.func_ctx.blocks[block].pristine,
+            self.is_pristine(block),
             "You can't add block parameters after adding any instruction"
         );
-        debug_assert_eq!(
-            self.func_ctx.blocks[block].user_param_count,
-            self.func.dfg.num_block_params(block)
-        );
-        self.func_ctx.blocks[block].user_param_count += 1;
         self.func.dfg.append_block_param(block, ty)
     }
 
@@ -714,14 +713,14 @@ impl<'a> FunctionBuilder<'a> {
 
     /// Returns `true` if and only if no instructions have been added since the last call to
     /// `switch_to_block`.
-    pub fn is_pristine(&self) -> bool {
-        self.func_ctx.blocks[self.position.unwrap()].pristine
+    fn is_pristine(&self, block: Block) -> bool {
+        self.func_ctx.status[block] == BlockStatus::Empty
     }
 
     /// Returns `true` if and only if a terminator instruction has been inserted since the
     /// last call to `switch_to_block`.
-    pub fn is_filled(&self) -> bool {
-        self.func_ctx.blocks[self.position.unwrap()].filled
+    fn is_filled(&self, block: Block) -> bool {
+        self.func_ctx.status[block] == BlockStatus::Filled
     }
 }
 
@@ -1078,7 +1077,7 @@ fn greatest_divisible_power_of_two(size: u64) -> u64 {
 impl<'a> FunctionBuilder<'a> {
     /// A Block is 'filled' when a terminator instruction is present.
     fn fill_current_block(&mut self) {
-        self.func_ctx.blocks[self.position.unwrap()].filled = true;
+        self.func_ctx.status[self.position.unwrap()] = BlockStatus::Filled;
     }
 
     fn declare_successor(&mut self, dest_block: Block, jump_inst: Inst) {
@@ -1089,10 +1088,12 @@ impl<'a> FunctionBuilder<'a> {
 
     fn handle_ssa_side_effects(&mut self, side_effects: SideEffects) {
         for split_block in side_effects.split_blocks_created {
-            self.func_ctx.blocks[split_block].filled = true
+            self.func_ctx.status[split_block] = BlockStatus::Filled;
         }
         for modified_block in side_effects.instructions_added_to_blocks {
-            self.func_ctx.blocks[modified_block].pristine = false
+            if self.is_pristine(modified_block) {
+                self.func_ctx.status[modified_block] = BlockStatus::Partial;
+            }
         }
     }
 }

cranelift/wasm/src/code_translator.rs

@@ -115,6 +115,9 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
         return Ok(());
     }
 
+    // Given that we believe the current block is reachable, the FunctionBuilder ought to agree.
+    debug_assert!(!builder.is_unreachable());
+
     // This big match treats all Wasm code operators.
     match op {
         /********************************** Locals ****************************************
@@ -380,18 +383,16 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
         Operator::End => {
             let frame = state.control_stack.pop().unwrap();
             let next_block = frame.following_code();
-
-            if !builder.is_unreachable() || !builder.is_pristine() {
-                let return_count = frame.num_return_values();
-                let return_args = state.peekn_mut(return_count);
-                canonicalise_then_jump(builder, frame.following_code(), return_args);
-                // You might expect that if we just finished an `if` block that
-                // didn't have a corresponding `else` block, then we would clean
-                // up our duplicate set of parameters that we pushed earlier
-                // right here. However, we don't have to explicitly do that,
-                // since we truncate the stack back to the original height
-                // below.
-            }
+            let return_count = frame.num_return_values();
+            let return_args = state.peekn_mut(return_count);
+
+            canonicalise_then_jump(builder, next_block, return_args);
+            // You might expect that if we just finished an `if` block that
+            // didn't have a corresponding `else` block, then we would clean
+            // up our duplicate set of parameters that we pushed earlier
+            // right here. However, we don't have to explicitly do that,
+            // since we truncate the stack back to the original height
+            // below.
 
             builder.switch_to_block(next_block);
             builder.seal_block(next_block);