Repository: bytecodealliance/wasmtime
Pull request #7525 (2023-11-10 19:21:55)

winch: Do not use `unconditional_jump` with `br_table`
This patch fixes how jumps are handled in `br_table`; prior to this change, `br_table` was implemented using
`CodeGenContext::unconditional_jump`; this function ensures, among other invariants that the value stack and stack pointer must be balanced according to the expectation of the target branch. Even though in `br_table` there's branch to a potentially known location, it's impossible be certain at compile time, which branch will be taken; in that regard, `br_table` behaves more like `br_if`. Using `unconditional_jump` resulted in the stack being manipulated multiple times and breaking the other existing invariants around stack balancing.

This commit makes it so that `br_table` doesn't rely on `unconditional_jump` anymore and instead it delegates control flow to the target branch, which will ensure that the value stack and stack pointer are correctly balanced when restoring reachability, very similar to what happens with `br_if`.

This issue was discovered while fuzzing and a file test is included with the test case.

<!--
Please make sure you include the following information:

- If this work has been discussed elsewhere, please include a link to that
  conversation. If it was discussed in an issue, just mention "issue #...".

- Explain why this change is needed. If the details are in an issue already,
  this can be brief.

Our development process is documented in the Wasmtime book:
https://docs.wasmtime.dev/contributing-development-process.html

Please ensure all communication follows the code of conduct:
https://github.com/bytecodealliance/wasmtime/blob/main/CODE_OF_CONDUCT.md
-->



Patch:
diff --git a/cranelift/codegen/src/egraph.rs b/cranelift/codegen/src/egraph.rs
index 2d5802572..0f699b68d 100644
--- a/cranelift/codegen/src/egraph.rs
+++ b/cranelift/codegen/src/egraph.rs
@@ -12,13 +12,13 @@ use crate::ir::{
     Block, DataFlowGraph, Function, Inst, InstructionData, Type, Value, ValueDef, ValueListPool,
 };
 use crate::loop_analysis::LoopAnalysis;
-use crate::opts::generated_code::ContextIter;
 use crate::opts::IsleContext;
 use crate::scoped_hash_map::{Entry as ScopedEntry, ScopedHashMap};
 use crate::trace;
 use crate::unionfind::UnionFind;
 use cranelift_entity::packed_option::ReservedValue;
 use cranelift_entity::SecondaryMap;
+use smallvec::SmallVec;
 use std::hash::Hasher;
 
 mod cost;
@@ -69,6 +69,9 @@ pub struct EgraphPass<'a> {
     eclasses: UnionFind<Value>,
 }
 
+// The maximum number of rewrites we will take from a single call into ISLE.
+const MATCHES_LIMIT: usize = 5;
+
 /// Context passed through node insertion and optimization.
 pub(crate) struct OptimizeCtx<'opt, 'analysis>
 where
@@ -87,6 +90,7 @@ where
     // Held locally during optimization of one node (recursively):
     pub(crate) rewrite_depth: usize,
     pub(crate) subsume_values: FxHashSet<Value>,
+    optimized_values: SmallVec<[Value; MATCHES_LIMIT]>,
 }
 
 /// For passing to `insert_pure_enode`. Sometimes the enode already
@@ -213,6 +217,7 @@ where
         // A pure node always has exactly one result.
         let orig_value = self.func.dfg.first_result(inst);
 
+        let mut optimized_values = std::mem::take(&mut self.optimized_values);
         let mut isle_ctx = IsleContext { ctx: self };
 
         // Limit rewrite depth. When we apply optimization rules, they
@@ -239,26 +244,26 @@ where
         // values produced as equivalents to this value.
         trace!("Calling into ISLE with original value {}", orig_value);
         isle_ctx.ctx.stats.rewrite_rule_invoked += 1;
-        let mut optimized_values =
-            crate::opts::generated_code::constructor_simplify(&mut isle_ctx, orig_value);
+        debug_assert!(optimized_values.is_empty());
+        crate::opts::generated_code::constructor_simplify(
+            &mut isle_ctx,
+            orig_value,
+            &mut optimized_values,
+        );
         trace!(
-            "  -> returned from ISLE, optimized values's size hint = {:?}",
-            optimized_values.size_hint()
+            "  -> returned from ISLE, generated {} optimized values",
+            optimized_values.len()
         );
+        if optimized_values.len() > MATCHES_LIMIT {
+            trace!("Reached maximum matches limit; too many optimized values, ignoring rest.");
+            optimized_values.truncate(MATCHES_LIMIT);
+        }
 
         // Create a union of all new values with the original (or
         // maybe just one new value marked as "subsuming" the
         // original, if present.)
-        let mut i = 0;
         let mut union_value = orig_value;
-        while let Some(optimized_value) = optimized_values.next(&mut isle_ctx) {
-            i += 1;
-            const MATCHES_LIMIT: u32 = 5;
-            if i > MATCHES_LIMIT {
-                trace!("Reached maximum matches limit; too many optimized values, ignoring rest.");
-                break;
-            }
-
+        for optimized_value in optimized_values.drain(..) {
             trace!(
                 "Returned from ISLE for {}, got {:?}",
                 orig_value,
@@ -305,6 +310,9 @@ where
 
         isle_ctx.ctx.rewrite_depth -= 1;
 
+        debug_assert!(isle_ctx.ctx.optimized_values.is_empty());
+        isle_ctx.ctx.optimized_values = optimized_values;
+
         union_value
     }
 
@@ -550,6 +558,7 @@ impl<'a> EgraphPass<'a> {
                             stats: &mut self.stats,
                             alias_analysis: self.alias_analysis,
                             alias_analysis_state: &mut alias_analysis_state,
+                            optimized_values: Default::default(),
                         };
 
                         if is_pure_for_egraph(ctx.func, inst) {
diff --git a/cranelift/codegen/src/opts.rs b/cranelift/codegen/src/opts.rs
index 00f350c6e..f7f03a0be 100644
--- a/cranelift/codegen/src/opts.rs
+++ b/cranelift/codegen/src/opts.rs
@@ -26,7 +26,7 @@ pub type ValueArray3 = [Value; 3];
 pub type ConstructorVec<T> = SmallVec<[T; 8]>;
 
 pub(crate) mod generated_code;
-use generated_code::ContextIter;
+use generated_code::{ContextIter, IntoContextIter};
 
 pub(crate) struct IsleContext<'a, 'b, 'c> {
     pub(crate) ctx: &'a mut OptimizeCtx<'b, 'c>,
@@ -38,6 +38,18 @@ pub(crate) struct InstDataEtorIter<'a, 'b, 'c> {
     _phantom2: PhantomData<&'b ()>,
     _phantom3: PhantomData<&'c ()>,
 }
+
+impl Default for InstDataEtorIter<'_, '_, '_> {
+    fn default() -> Self {
+        InstDataEtorIter {
+            stack: SmallVec::default(),
+            _phantom1: PhantomData,
+            _phantom2: PhantomData,
+            _phantom3: PhantomData,
+        }
+    }
+}
+
 impl<'a, 'b, 'c> InstDataEtorIter<'a, 'b, 'c> {
     fn new(root: Value) -> Self {
         debug_assert_ne!(root, Value::reserved_value());
@@ -84,13 +96,27 @@ where
     }
 }
 
+impl<'a, 'b, 'c> IntoContextIter for InstDataEtorIter<'a, 'b, 'c>
+where
+    'b: 'a,
+    'c: 'b,
+{
+    type Context = IsleContext<'a, 'b, 'c>;
+    type Output = (Type, InstructionData);
+    type IntoIter = Self;
+
+    fn into_context_iter(self) -> Self {
+        self
+    }
+}
+
 impl<'a, 'b, 'c> generated_code::Context for IsleContext<'a, 'b, 'c> {
     isle_common_prelude_methods!();
 
-    type inst_data_etor_iter = InstDataEtorIter<'a, 'b, 'c>;
+    type inst_data_etor_returns = InstDataEtorIter<'a, 'b, 'c>;
 
-    fn inst_data_etor(&mut self, eclass: Value) -> InstDataEtorIter<'a, 'b, 'c> {
-        InstDataEtorIter::new(eclass)
+    fn inst_data_etor(&mut self, eclass: Value, returns: &mut InstDataEtorIter<'a, 'b, 'c>) {
+        *returns = InstDataEtorIter::new(eclass);
     }
 
     fn make_inst_ctor(&mut self, ty: Type, op: &InstructionData) -> Value {
diff --git a/cranelift/isle/isle/isle_examples/link/multi_constructor_main.rs b/cranelift/isle/isle/isle_examples/link/multi_constructor_main.rs
index 16725e36c..11396a93b 100644
--- a/cranelift/isle/isle/isle_examples/link/multi_constructor_main.rs
+++ b/cranelift/isle/isle/isle_examples/link/multi_constructor_main.rs
@@ -1,15 +1,15 @@
 mod multi_constructor;
-
-pub(crate) type ConstructorVec<T> = Vec<T>;
+use multi_constructor::{ContextIter, IntoContextIter};
 
 struct Context;
 
+#[derive(Default)]
 struct It {
     i: u32,
     limit: u32,
 }
 
-impl multi_constructor::ContextIter for It {
+impl ContextIter for It {
     type Context = Context;
     type Output = u32;
     fn next(&mut self, _ctx: &mut Self::Context) -> Option<u32> {
@@ -23,15 +23,25 @@ impl multi_constructor::ContextIter for It {
     }
 }
 
+impl IntoContextIter for It {
+    type Context = Context;
+    type Output = u32;
+    type IntoIter = It;
+    fn into_context_iter(self) -> It {
+        self
+    }
+}
+
 impl multi_constructor::Context for Context {
-    type etor_C_iter = It;
-    fn etor_C(&mut self, value: u32) -> It {
-        It { i: 0, limit: value }
+    type etor_C_returns = It;
+    fn etor_C(&mut self, value: u32, returns: &mut It) {
+        returns.i = 0;
+        returns.limit = value;
     }
 
-    type ctor_B_iter = multi_constructor::ContextIterWrapper<u32, std::vec::IntoIter<u32>, Context>;
-    fn ctor_B(&mut self, value: u32) -> Self::ctor_B_iter {
-        (0..value).rev().collect::<Vec<_>>().into_iter().into()
+    type ctor_B_returns = multi_constructor::ContextIterWrapper<Vec<u32>, Context>;
+    fn ctor_B(&mut self, value: u32, returns: &mut Self::ctor_B_returns) {
+        returns.extend((0..value).rev());
     }
 }
 
@@ -55,16 +65,21 @@ impl<'a, Item, I: multi_constructor::ContextIter<Output = Item, Context = Contex
 
 fn main() {
     let mut ctx = Context;
-    let l1 = multi_constructor::constructor_A(&mut ctx, 10);
-    let l2 = multi_constructor::constructor_D(&mut ctx, 5);
+
+    let mut l1 = multi_constructor::ContextIterWrapper::<Vec<_>, _>::default();
+    multi_constructor::constructor_A(&mut ctx, 10, &mut l1);
+
+    let mut l2 = multi_constructor::ContextIterWrapper::<Vec<_>, _>::default();
+    multi_constructor::constructor_D(&mut ctx, 5, &mut l2);
+
     let l1 = IterWithContext {
         ctx: &mut ctx,
-        it: l1,
+        it: l1.into_context_iter(),
     }
     .collect::<Vec<_>>();
     let l2 = IterWithContext {
         ctx: &mut ctx,
-        it: l2,
+        it: l2.into_context_iter(),
     }
     .collect::<Vec<_>>();
     println!("l1 = {:?} l2 = {:?}", l1, l2);
diff --git a/cranelift/isle/isle/isle_examples/link/multi_extractor_main.rs b/cranelift/isle/isle/isle_examples/link/multi_extractor_main.rs
index 90a730b7c..c7c7dd406 100644
--- a/cranelift/isle/isle/isle_examples/link/multi_extractor_main.rs
+++ b/cranelift/isle/isle/isle_examples/link/multi_extractor_main.rs
@@ -1,8 +1,6 @@
 mod multi_extractor;
 
-use multi_extractor::ContextIter;
-
-pub(crate) type ConstructorVec<T> = Vec<T>;
+use multi_extractor::{ContextIter, IntoContextIter};
 
 #[derive(Clone)]
 pub enum A {
@@ -10,12 +8,13 @@ pub enum A {
     C,
 }
 
+#[derive(Default)]
 struct It {
     i: u32,
     arg: u32,
 }
 
-impl multi_extractor::ContextIter for It {
+impl ContextIter for It {
     type Context = Context;
     type Output = (A, u32);
     fn next(&mut self, _ctx: &mut Self::Context) -> Option<Self::Output> {
@@ -34,17 +33,29 @@ impl multi_extractor::ContextIter for It {
     }
 }
 
+impl IntoContextIter for It {
+    type Context = Context;
+    type IntoIter = It;
+    type Output = (A, u32);
+    fn into_context_iter(self) -> It {
+        self
+    }
+}
+
 struct Context;
 impl multi_extractor::Context for Context {
-    type e1_etor_iter = It;
-    fn e1_etor(&mut self, arg0: u32) -> It {
-        It { i: 0, arg: arg0 }
+    type e1_etor_returns = It;
+    fn e1_etor(&mut self, arg0: u32, returns: &mut It) {
+        returns.i = 0;
+        returns.arg = arg0;
     }
 }
 
 fn main() {
     let mut ctx = Context;
-    let x = multi_extractor::constructor_Rule(&mut ctx, 0xf0).next(&mut ctx);
-    let y = multi_extractor::constructor_Rule(&mut ctx, 0).next(&mut ctx);
+    let mut x = vec![];
+    multi_extractor::constructor_Rule(&mut ctx, 0xf0, &mut x);
+    let mut y = vec![];
+    multi_extractor::constructor_Rule(&mut ctx, 0, &mut y);
     println!("x = {:?} y = {:?}", x, y);
 }
diff --git a/cranelift/isle/isle/src/codegen.rs b/cranelift/isle/isle/src/codegen.rs
index d540afefb..1c3dac155 100644
--- a/cranelift/isle/isle/src/codegen.rs
+++ b/cranelift/isle/isle/src/codegen.rs
@@ -154,7 +154,7 @@ impl<'a> Codegen<'a> {
         if sig.ret_kind == ReturnKind::Iterator {
             writeln!(
                 code,
-                "{indent}type {name}_iter: ContextIter<Context = Self, Output = {output}>;",
+                "{indent}type {name}_returns: Default + IntoContextIter<Context = Self, Output = {output}>;",
                 indent = indent,
                 name = sig.func_name,
                 output = ret_tuple,
@@ -165,7 +165,7 @@ impl<'a> Codegen<'a> {
         let ret_ty = match sig.ret_kind {
             ReturnKind::Plain => ret_tuple,
             ReturnKind::Option => format!("Option<{}>", ret_tuple),
-            ReturnKind::Iterator => format!("Self::{}_iter", sig.func_name),
+            ReturnKind::Iterator => format!("()"),
         };
 
         writeln!(
@@ -178,6 +178,11 @@ impl<'a> Codegen<'a> {
                 .iter()
                 .enumerate()
                 .map(|(i, &ty)| format!("arg{}: {}", i, self.type_name(ty, /* by_ref = */ true)))
+                .chain(if sig.ret_kind == ReturnKind::Iterator {
+                    Some(format!("returns: &mut Self::{}_returns", sig.func_name))
+                } else {
+                    None
+                })
                 .collect::<Vec<_>>()
                 .join(", "),
             ret_ty = ret_ty,
@@ -217,35 +222,77 @@ impl<'a> Codegen<'a> {
         writeln!(
             code,
             r#"
-           pub trait ContextIter {{
-               type Context;
-               type Output;
-               fn next(&mut self, ctx: &mut Self::Context) -> Option<Self::Output>;
-               fn size_hint(&self) -> (usize, Option<usize>) {{ (0, None) }}
-           }}
-
-           pub struct ContextIterWrapper<Item, I: Iterator < Item = Item>, C: Context> {{
-               iter: I,
-               _ctx: PhantomData<C>,
-           }}
-           impl<Item, I: Iterator<Item = Item>, C: Context> From<I> for ContextIterWrapper<Item, I, C> {{
-               fn from(iter: I) -> Self {{
-                   Self {{ iter, _ctx: PhantomData }}
-               }}
-           }}
-           impl<Item, I: Iterator<Item = Item>, C: Context> ContextIter for ContextIterWrapper<Item, I, C> {{
-               type Context = C;
-               type Output = Item;
-               fn next(&mut self, _ctx: &mut Self::Context) -> Option<Self::Output> {{
-                   self.iter.next()
-               }}
-               fn size_hint(&self) -> (usize, Option<usize>) {{
-                   self.iter.size_hint()
-               }}
-           }}
+pub trait ContextIter {{
+    type Context;
+    type Output;
+    fn next(&mut self, ctx: &mut Self::Context) -> Option<Self::Output>;
+    fn size_hint(&self) -> (usize, Option<usize>) {{ (0, None) }}
+}}
+
+pub trait IntoContextIter {{
+    type Context;
+    type Output;
+    type IntoIter: ContextIter<Context = Self::Context, Output = Self::Output>;
+    fn into_context_iter(self) -> Self::IntoIter;
+}}
+
+pub struct ContextIterWrapper<I, C> {{
+    iter: I,
+    _ctx: std::marker::PhantomData<C>,
+}}
+impl<I: Default, C> Default for ContextIterWrapper<I, C> {{
+    fn default() -> Self {{
+        ContextIterWrapper {{
+            iter: I::default(),
+            _ctx: std::marker::PhantomData
+        }}
+    }}
+}}
+impl<I, C> std::ops::Deref for ContextIterWrapper<I, C> {{
+    type Target = I;
+    fn deref(&self) -> &I {{
+        &self.iter
+    }}
+}}
+impl<I, C> std::ops::DerefMut for ContextIterWrapper<I, C> {{
+    fn deref_mut(&mut self) -> &mut I {{
+        &mut self.iter
+    }}
+}}
+impl<I: Iterator, C: Context> From<I> for ContextIterWrapper<I, C> {{
+    fn from(iter: I) -> Self {{
+        Self {{ iter, _ctx: std::marker::PhantomData }}
+    }}
+}}
+impl<I: Iterator, C: Context> ContextIter for ContextIterWrapper<I, C> {{
+    type Context = C;
+    type Output = I::Item;
+    fn next(&mut self, _ctx: &mut Self::Context) -> Option<Self::Output> {{
+        self.iter.next()
+    }}
+    fn size_hint(&self) -> (usize, Option<usize>) {{
+        self.iter.size_hint()
+    }}
+}}
+impl<I: IntoIterator, C: Context> IntoContextIter for ContextIterWrapper<I, C> {{
+    type Context = C;
+    type Output = I::Item;
+    type IntoIter = ContextIterWrapper<I::IntoIter, C>;
+    fn into_context_iter(self) -> Self::IntoIter {{
+        ContextIterWrapper {{
+            iter: self.iter.into_iter(),
+            _ctx: std::marker::PhantomData
+        }}
+    }}
+}}
+impl<T, E: Extend<T>, C> Extend<T> for ContextIterWrapper<E, C> {{
+    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {{
+        self.iter.extend(iter);
+    }}
+}}
            "#,
         )
-            .unwrap();
+        .unwrap();
     }
 
     fn generate_internal_types(&self, code: &mut String) {
@@ -353,13 +400,20 @@ impl<'a> Codegen<'a> {
                 writeln!(ctx.out, ",")?;
             }
 
-            write!(ctx.out, "{}) -> ", &ctx.indent)?;
             let (_, ret) = self.ty(sig.ret_tys[0]);
             let ret = &self.typeenv.syms[ret.index()];
+
+            if let ReturnKind::Iterator = sig.ret_kind {
+                writeln!(
+                    ctx.out,
+                    "{}    returns: &mut impl Extend<{}>,",
+                    &ctx.indent, ret
+                )?;
+            }
+
+            write!(ctx.out, "{}) -> ", &ctx.indent)?;
             match sig.ret_kind {
-                ReturnKind::Iterator => {
-                    write!(ctx.out, "impl ContextIter<Context = C, Output = {}>", ret)?
-                }
+                ReturnKind::Iterator => write!(ctx.out, "()")?,
                 ReturnKind::Option => write!(ctx.out, "Option<{}>", ret)?,
                 ReturnKind::Plain => write!(ctx.out, "{}", ret)?,
             };
@@ -367,21 +421,13 @@ impl<'a> Codegen<'a> {
             let scope = ctx.enter_scope();
             ctx.begin_block()?;
 
-            if sig.ret_kind == ReturnKind::Iterator {
-                writeln!(
-                    ctx.out,
-                    "{}let mut returns = ConstructorVec::new();",
-                    &ctx.indent
-                )?;
-            }
-
             self.emit_block(&mut ctx, &root, sig.ret_kind)?;
 
             match (sig.ret_kind, root.steps.last()) {
                     (ReturnKind::Iterator, _) => {
                         writeln!(
                             ctx.out,
-                            "{}return ContextIterWrapper::from(returns.into_iter());",
+                            "{}return;",
                             &ctx.indent
                         )?;
                     }
@@ -440,7 +486,47 @@ impl<'a> Codegen<'a> {
 
         for case in block.steps.iter() {
             for &expr in case.bind_order.iter() {
-                write!(ctx.out, "{}let v{} = ", &ctx.indent, expr.index())?;
+                let iter_return = match &ctx.ruleset.bindings[expr.index()] {
+                    Binding::Extractor { term, .. } => {
+                        let termdata = &self.termenv.terms[term.index()];
+                        let sig = termdata.extractor_sig(self.typeenv).unwrap();
+                        if sig.ret_kind == ReturnKind::Iterator {
+                            if termdata.has_external_extractor() {
+                                Some(format!("C::{}_returns", sig.func_name))
+                            } else {
+                                Some(format!("ContextIterWrapper::<ConstructorVec<_>, _>"))
+                            }
+                        } else {
+                            None
+                        }
+                    }
+                    Binding::Constructor { term, .. } => {
+                        let termdata = &self.termenv.terms[term.index()];
+                        let sig = termdata.constructor_sig(self.typeenv).unwrap();
+                        if sig.ret_kind == ReturnKind::Iterator {
+                            if termdata.has_external_constructor() {
+                                Some(format!("C::{}_returns", sig.func_name))
+                            } else {
+                                Some(format!("ContextIterWrapper::<ConstructorVec<_>, _>"))
+                            }
+                        } else {
+                            None
+                        }
+                    }
+                    _ => None,
+                };
+                if let Some(ty) = iter_return {
+                    writeln!(
+                        ctx.out,
+                        "{}let mut v{} = {}::default();",
+                        &ctx.indent,
+                        expr.index(),
+                        ty
+                    )?;
+                    write!(ctx.out, "{}", &ctx.indent)?;
+                } else {
+                    write!(ctx.out, "{}let v{} = ", &ctx.indent, expr.index())?;
+                }
                 self.emit_expr(ctx, expr)?;
                 writeln!(ctx.out, ";")?;
                 ctx.is_bound.insert(expr);
@@ -507,9 +593,15 @@ impl<'a> Codegen<'a> {
                         _ => unreachable!("Loop from a non-Iterator"),
                     };
                     let scope = ctx.enter_scope();
-                    write!(ctx.out, "{}let mut v{} = ", &ctx.indent, source.index())?;
-                    self.emit_expr(ctx, *source)?;
-                    writeln!(ctx.out, ";")?;
+
+                    writeln!(
+                        ctx.out,
+                        "{}let mut v{} = v{}.into_context_iter();",
+                        &ctx.indent,
+                        source.index(),
+                        source.index(),
+                    )?;
+
                     write!(
                         ctx.out,
                         "{}while let Some(v{}) = v{}.next(ctx)",
@@ -534,7 +626,7 @@ impl<'a> Codegen<'a> {
                     match ret_kind {
                         ReturnKind::Plain => write!(ctx.out, "return ")?,
                         ReturnKind::Option => write!(ctx.out, "return Some(")?,
-                        ReturnKind::Iterator => write!(ctx.out, "returns.push(")?,
+                        ReturnKind::Iterator => write!(ctx.out, "returns.extend(Some(")?,
                     }
                     self.emit_expr(ctx, result)?;
                     if ctx.is_ref.contains(&result) {
@@ -542,7 +634,8 @@ impl<'a> Codegen<'a> {
                     }
                     match ret_kind {
                         ReturnKind::Plain => writeln!(ctx.out, ";")?,
-                        ReturnKind::Option | ReturnKind::Iterator => writeln!(ctx.out, ");")?,
+                        ReturnKind::Option => writeln!(ctx.out, ");")?,
+                        ReturnKind::Iterator => writeln!(ctx.out, "));")?,
                     }
                 }
             }
@@ -560,6 +653,7 @@ impl<'a> Codegen<'a> {
         let mut call =
             |term: TermId,
              parameters: &[BindingId],
+
              get_sig: fn(&Term, &TypeEnv) -> Option<ExternalSig>| {
                 let termdata = &self.termenv.terms[term.index()];
                 let sig = get_sig(termdata, self.typeenv).unwrap();
@@ -584,6 +678,9 @@ impl<'a> Codegen<'a> {
                     self.emit_expr(ctx, parameter)?;
                     write!(ctx.out, "{}", after)?;
                 }
+                if let ReturnKind::Iterator = sig.ret_kind {
+                    write!(ctx.out, ", &mut v{}", result.index())?;
+                }
                 write!(ctx.out, ")")
             };
 
diff --git a/crates/wasi-http/src/types_impl.rs b/crates/wasi-http/src/types_impl.rs
index 89455e3aa..00400a07c 100644
--- a/crates/wasi-http/src/types_impl.rs
+++ b/crates/wasi-http/src/types_impl.rs
@@ -32,10 +32,10 @@ fn move_fields(table: &mut Table, id: Resource<HostFields>) -> wasmtime::Result<
     }
 }
 
-fn get_fields_mut<'a>(
+fn get_fields<'a>(
     table: &'a mut Table,
     id: &Resource<HostFields>,
-) -> wasmtime::Result<&'a mut FieldMap> {
+) -> wasmtime::Result<&'a FieldMap> {
     let fields = table.get(&id)?;
     if let HostFields::Ref { parent, get_fields } = *fields {
         let entry = table.get_any_mut(parent)?;
@@ -51,6 +51,16 @@ fn get_fields_mut<'a>(
     }
 }
 
+fn get_fields_mut<'a>(
+    table: &'a mut Table,
+    id: &Resource<HostFields>,
+) -> wasmtime::Result<Result<&'a mut FieldMap, types::HeaderError>> {
+    match table.get_mut(&id)? {
+        HostFields::Owned { fields } => Ok(Ok(fields)),
+        HostFields::Ref { .. } => Ok(Err(types::HeaderError::Immutable)),
+    }
+}
+
 fn is_forbidden_header<T: WasiHttpView>(view: &mut T, name: &HeaderName) -> bool {
     static FORBIDDEN_HEADERS: [HeaderName; 9] = [
         hyper::header::CONNECTION,
@@ -83,7 +93,7 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostFields for T {
         &mut self,
         entries: Vec<(String, Vec<u8>)>,
     ) -> wasmtime::Result<Result<Resource<HostFields>, types::HeaderError>> {
-        let mut map = hyper::HeaderMap::new();
+        let mut fields = hyper::HeaderMap::new();
 
         for (header, value) in entries {
             let header = match hyper::header::HeaderName::from_bytes(header.as_bytes()) {
@@ -100,12 +110,12 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostFields for T {
                 Err(_) => return Ok(Err(types::HeaderError::InvalidSyntax)),
             };
 
-            map.append(header, value);
+            fields.append(header, value);
         }
 
         let id = self
             .table()
-            .push(HostFields::Owned { fields: map })
+            .push(HostFields::Owned { fields })
             .context("[new_fields] pushing fields")?;
 
         Ok(Ok(id))
@@ -128,7 +138,7 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostFields for T {
             Err(_) => return Ok(vec![]),
         };
 
-        let res = get_fields_mut(self.table(), &fields)
+        let res = get_fields(self.table(), &fields)
             .context("[fields_get] getting fields")?
             .get_all(header)
             .into_iter()
@@ -160,14 +170,14 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostFields for T {
             }
         }
 
-        let m =
-            get_fields_mut(self.table(), &fields).context("[fields_set] getting mutable fields")?;
-        m.remove(&header);
-        for value in values {
-            m.append(&header, value);
-        }
-
-        Ok(Ok(()))
+        Ok(get_fields_mut(self.table(), &fields)
+            .context("[fields_set] getting mutable fields")?
+            .map(|fields| {
+                fields.remove(&header);
+                for value in values {
+                    fields.append(&header, value);
+                }
+            }))
     }
 
     fn delete(
@@ -184,9 +194,9 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostFields for T {
             return Ok(Err(types::HeaderError::Forbidden));
         }
 
-        let m = get_fields_mut(self.table(), &fields)?;
-        m.remove(header);
-        Ok(Ok(()))
+        Ok(get_fields_mut(self.table(), &fields)?.map(|fields| {
+            fields.remove(header);
+        }))
     }
 
     fn append(
@@ -209,27 +219,25 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostFields for T {
             Err(_) => return Ok(Err(types::HeaderError::InvalidSyntax)),
         };
 
-        let m = get_fields_mut(self.table(), &fields)
-            .context("[fields_append] getting mutable fields")?;
-
-        m.append(header, value);
-        Ok(Ok(()))
+        Ok(get_fields_mut(self.table(), &fields)
+            .context("[fields_append] getting mutable fields")?
+            .map(|fields| {
+                fields.append(header, value);
+            }))
     }
 
     fn entries(
         &mut self,
         fields: Resource<HostFields>,
     ) -> wasmtime::Result<Vec<(String, Vec<u8>)>> {
-        let fields = get_fields_mut(self.table(), &fields)?;
-        let result = fields
+        Ok(get_fields(self.table(), &fields)?
             .iter()
             .map(|(name, value)| (name.as_str().to_owned(), value.as_bytes().to_owned()))
-            .collect();
-        Ok(result)
+            .collect())
     }
 
     fn clone(&mut self, fields: Resource<HostFields>) -> wasmtime::Result<Resource<HostFields>> {
-        let fields = get_fields_mut(self.table(), &fields)
+        let fields = get_fields(self.table(), &fields)
             .context("[fields_clone] getting fields")?
             .clone();
 
@@ -837,7 +845,7 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostOutgoingBody for T {
             .expect("outgoing-body trailer_sender consumed by a non-owning function");
 
         let message = if let Some(ts) = ts {
-            FinishMessage::Trailers(get_fields_mut(self.table(), &ts)?.clone().into())
+            FinishMessage::Trailers(move_fields(self.table(), ts)?)
         } else {
             FinishMessage::Finished
         };
diff --git a/crates/wasi-http/wit/deps/http/types.wit b/crates/wasi-http/wit/deps/http/types.wit
index fac47f61e..4662f3fd6 100644
--- a/crates/wasi-http/wit/deps/http/types.wit
+++ b/crates/wasi-http/wit/deps/http/types.wit
@@ -42,6 +42,7 @@ interface types {
   variant header-error {
     invalid-syntax,
     forbidden,
+    immutable
   }
 
   /// Field keys are always strings.
diff --git a/crates/wasi/wit/deps/http/types.wit b/crates/wasi/wit/deps/http/types.wit
index fac47f61e..4662f3fd6 100644
--- a/crates/wasi/wit/deps/http/types.wit
+++ b/crates/wasi/wit/deps/http/types.wit
@@ -42,6 +42,7 @@ interface types {
   variant header-error {
     invalid-syntax,
     forbidden,
+    immutable
   }
 
   /// Field keys are always strings.
diff --git a/winch/codegen/src/codegen/context.rs b/winch/codegen/src/codegen/context.rs
index 1b7460277..6e3672e30 100644
--- a/winch/codegen/src/codegen/context.rs
+++ b/winch/codegen/src/codegen/context.rs
@@ -402,6 +402,13 @@ impl<'a, 'builtins> CodeGenContext<'a, 'builtins> {
         self.reachable = false;
     }
 
+    /// A combination of [Self::pop_abi_results] and [Self::push_abi_results]
+    /// to be used on conditional branches: br_if and br_table.
+    pub fn top_abi_results<M: MacroAssembler>(&mut self, result: &ABIResult, masm: &mut M) {
+        self.pop_abi_results(result, masm);
+        self.push_abi_results(result, masm);
+    }
+
     /// Handles the emission of the ABI result. This function is used at the end
     /// of a block or function to pop the results from the value stack into the
     /// corresponding ABI result representation.
diff --git a/winch/codegen/src/visitor.rs b/winch/codegen/src/visitor.rs
index a69992ac1..bbbf55ad1 100644
--- a/winch/codegen/src/visitor.rs
+++ b/winch/codegen/src/visitor.rs
@@ -1169,8 +1169,7 @@ where
             self.masm,
             |ctx, masm| ctx.pop_to_reg(masm, None),
         );
-        self.context.pop_abi_results(&result, self.masm);
-        self.context.push_abi_results(&result, self.masm);
+        self.context.top_abi_results(&result, self.masm);
         self.masm.branch(
             IntCmpKind::Ne,
             top.reg.into(),
@@ -1190,15 +1189,17 @@ where
         let labels: SmallVec<[_; 5]> = (0..len).map(|_| self.masm.get_label()).collect();
 
         let default_index = control_index(targets.default(), self.control_frames.len());
-        let default_result = self.control_frames[default_index].as_target_result();
+        let default_frame = &self.control_frames[default_index];
+        let default_result = default_frame.as_target_result();
         let (index, tmp) = self.context.maybe_without1::<(TypedReg, _), M, _>(
             default_result.result_reg(),
             self.masm,
             |cx, masm| (cx.pop_to_reg(masm, None), cx.any_gpr(masm)),
         );
 
-        self.context.pop_abi_results(&default_result, self.masm);
-        self.context.push_abi_results(&default_result, self.masm);
+        // Materialize any constants or locals into their result representation,
+        // so that when reachability is restored, they are correctly located.
+        self.context.top_abi_results(&default_result, self.masm);
         self.masm.jmp_table(&labels, index.into(), tmp);
 
         for (t, l) in targets
@@ -1208,18 +1209,16 @@ where
             .zip(labels.iter())
         {
             let control_index = control_index(t.unwrap(), self.control_frames.len());
+            let frame = &mut self.control_frames[control_index];
 
+            // NB: We don't perform any result handling as it was
+            // already taken care of above before jumping to the
+            // jump table above.
             self.masm.bind(*l);
-            self.context.unconditional_jump(
-                &mut self.control_frames[control_index],
-                self.masm,
-                // NB: We don't perform any result handling as it was
-                // already taken care of above before jumping to the
-                // jump table above. The call to `unconditional_jump`,
-                // will stil take care of the proper stack alignment.
-                |_, _, _| {},
-            )
+            self.masm.jmp(*frame.label());
+            frame.set_as_target();
         }
+        self.context.reachable = false;
         self.context.free_reg(index.reg);
         self.context.free_reg(tmp);
     }


Test patch:
diff --git a/crates/test-programs/src/bin/http_outbound_request_invalid_header.rs b/crates/test-programs/src/bin/http_outbound_request_invalid_header.rs
index 09f173298..65931dd61 100644
--- a/crates/test-programs/src/bin/http_outbound_request_invalid_header.rs
+++ b/crates/test-programs/src/bin/http_outbound_request_invalid_header.rs
@@ -1,4 +1,4 @@
-use test_programs::wasi::http::types::{HeaderError, Headers};
+use test_programs::wasi::http::types::{HeaderError, Headers, OutgoingRequest};
 
 fn main() {
     let hdrs = Headers::new();
@@ -57,4 +57,22 @@ fn main() {
         Headers::from_list(&[("ok-header-name".to_owned(), b"bad\nvalue".to_vec())]),
         Err(HeaderError::InvalidSyntax)
     ));
+
+    let req = OutgoingRequest::new(hdrs);
+    let hdrs = req.headers();
+
+    assert!(matches!(
+        hdrs.set(&"Content-Length".to_owned(), &[b"10".to_vec()]),
+        Err(HeaderError::Immutable),
+    ));
+
+    assert!(matches!(
+        hdrs.append(&"Content-Length".to_owned(), &b"10".to_vec()),
+        Err(HeaderError::Immutable),
+    ));
+
+    assert!(matches!(
+        hdrs.delete(&"Content-Length".to_owned()),
+        Err(HeaderError::Immutable),
+    ));
 }
diff --git a/winch/filetests/filetests/x64/br_table/stack_handling.wat b/winch/filetests/filetests/x64/br_table/stack_handling.wat
new file mode 100644
index 000000000..76ccf78c0
--- /dev/null
+++ b/winch/filetests/filetests/x64/br_table/stack_handling.wat
@@ -0,0 +1,37 @@
+;;! target = "x86_64"
+(module
+  (func (;0;) (param i32)
+    local.get 0
+    block ;; label = @1
+      i32.const 808727609
+      br_table 0 (;@1;) 1 (;@0;) 0 (;@1;)
+    end
+    drop
+  )
+  (export "main" (func 0))
+)
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec10             	sub	rsp, 0x10
+;;    8:	 897c240c             	mov	dword ptr [rsp + 0xc], edi
+;;    c:	 4c89742404           	mov	qword ptr [rsp + 4], r14
+;;   11:	 448b5c240c           	mov	r11d, dword ptr [rsp + 0xc]
+;;   16:	 4153                 	push	r11
+;;   18:	 b839343430           	mov	eax, 0x30343439
+;;   1d:	 b902000000           	mov	ecx, 2
+;;   22:	 39c1                 	cmp	ecx, eax
+;;   24:	 0f42c1               	cmovb	eax, ecx
+;;   27:	 4c8d1d0a000000       	lea	r11, [rip + 0xa]
+;;   2e:	 49630c83             	movsxd	rcx, dword ptr [r11 + rax*4]
+;;   32:	 4901cb               	add	r11, rcx
+;;   35:	 41ffe3               	jmp	r11
+;;   38:	 0c00                 	or	al, 0
+;;   3a:	 0000                 	add	byte ptr [rax], al
+;;   3c:	 1000                 	adc	byte ptr [rax], al
+;;   3e:	 0000                 	add	byte ptr [rax], al
+;;   40:	 0c00                 	or	al, 0
+;;   42:	 0000                 	add	byte ptr [rax], al
+;;   44:	 4883c408             	add	rsp, 8
+;;   48:	 4883c410             	add	rsp, 0x10
+;;   4c:	 5d                   	pop	rbp
+;;   4d:	 c3                   	ret