Switch interpreter dispatch to musttail (become) threading

.github/workflows/rust.yml

@@ -143,3 +143,12 @@ jobs:
       run: rustup run nightly-2024-07-10 cargo install --locked --version 0.12.0 cargo-fuzz
     - name: Fuzz
       run: ./ci/jobs/fuzz.sh
+  nightly-musttail:
+    runs-on: ubuntu-22.04
+    continue-on-error: true
+    steps:
+    - uses: actions/checkout@v4
+    - name: Install latest nightly
+      run: rustup toolchain install nightly --profile minimal
+    - name: Build and test with interpreter-musttail-dispatch
+      run: ./ci/jobs/build-and-test-nightly-musttail.sh

ci/jobs/build-and-test-nightly-musttail.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+set -euo pipefail
+cd -- "$(dirname -- "${BASH_SOURCE[0]}")"
+cd ../..
+
+echo ">> cargo +nightly test (polkavm, interpreter-musttail-dispatch, interpreter tests only)"
+cargo +nightly test -p polkavm --features interpreter-musttail-dispatch -- tests::interpreter_
+
+echo ">> cargo +nightly test --release (polkavm, interpreter-musttail-dispatch, interpreter tests only)"
+cargo +nightly test --release -p polkavm --features interpreter-musttail-dispatch -- tests::interpreter_

crates/polkavm/Cargo.toml

@@ -50,3 +50,6 @@ generic-sandbox = []
 
 # Internal feature for testing. DO NOT USE.
 export-internals-for-testing = []
+
+# Musttail (`become`) direct-threaded interpreter dispatch. Requires nightly.
+interpreter-musttail-dispatch = []

crates/polkavm/src/interpreter.rs

@@ -25,8 +25,20 @@ use polkavm_common::program::{
 use polkavm_common::utils::{align_to_next_page_usize, slice_assume_init_mut, ArcBytes, GasVisitorT};
 
 type Target = u32;
+
+#[cfg(feature = "interpreter-musttail-dispatch")]
+type HandlerResult = InterruptKind;
+#[cfg(not(feature = "interpreter-musttail-dispatch"))]
 type HandlerResult = Target;
 
+// `become` is parse-gated; hiding it in a macro keeps the parser happy when the feature is off.
+#[cfg(feature = "interpreter-musttail-dispatch")]
+macro_rules! tail_call {
+    ($e:expr) => {
+        become $e
+    };
+}
+
 #[derive(Copy, Clone)]
 pub enum RegImm {
     Reg(Reg),
@@ -1711,16 +1723,33 @@ impl InterpretedInstance {
             log::trace!("Implicitly resuming at: [{}]", self.next_compiled_offset);
         }
 
-        let mut offset = self.next_compiled_offset;
-        loop {
+        #[cfg(feature = "interpreter-musttail-dispatch")]
+        {
+            // Normal call, not `become`: `run_impl`'s signature doesn't match a handler's.
+            // The musttail chain runs inside the handlers and unwinds back here.
+            let off = self.next_compiled_offset;
             if DEBUG {
                 self.cycle_counter += 1;
             }
-
-            if let Some(handler) = self.compiled_handlers.get(cast(offset).to_usize()) {
-                offset = handler(self, offset);
+            if let Some(&handler) = self.compiled_handlers.get(cast(off).to_usize()) {
+                handler(self, off)
             } else {
-                return self.interrupt.clone();
+                self.interrupt.clone()
+            }
+        }
+        #[cfg(not(feature = "interpreter-musttail-dispatch"))]
+        {
+            let mut offset = self.next_compiled_offset;
+            loop {
+                if DEBUG {
+                    self.cycle_counter += 1;
+                }
+
+                if let Some(handler) = self.compiled_handlers.get(cast(offset).to_usize()) {
+                    offset = handler(self, offset);
+                } else {
+                    return self.interrupt.clone();
+                }
             }
         }
     }
@@ -2949,7 +2978,22 @@ macro_rules! define_interpreter {
                 #[allow(clippy::needless_lifetimes)]
                 pub fn $handler_name<'a, $(M: $M_ty,)? $(const $const: $const_ty),+>($self: &'a mut InterpretedInstance, compiled_offset: Target) -> HandlerResult {
                     let $compiled_offset = compiled_offset;
-                    define_interpreter!(@define $handler_name $body $self $compiled_offset $($arg)*)
+                    let next_off: Target = define_interpreter!(@define $handler_name $body $self $compiled_offset $($arg)*);
+                    #[cfg(feature = "interpreter-musttail-dispatch")]
+                    {
+                        if DEBUG {
+                            $self.cycle_counter += 1;
+                        }
+                        if let Some(&handler) = $self.compiled_handlers.get(cast(next_off).to_usize()) {
+                            tail_call!(handler($self, next_off))
+                        } else {
+                            $self.interrupt.clone()
+                        }
+                    }
+                    #[cfg(not(feature = "interpreter-musttail-dispatch"))]
+                    {
+                        next_off
+                    }
                 }
             )+
         }

crates/polkavm/src/lib.rs

@@ -1,4 +1,6 @@
 #![cfg_attr(not(feature = "std"), no_std)]
+#![cfg_attr(feature = "interpreter-musttail-dispatch", feature(explicit_tail_calls))]
+#![cfg_attr(feature = "interpreter-musttail-dispatch", allow(incomplete_features))]
 #![forbid(unused_must_use)]
 #![forbid(clippy::missing_safety_doc)]
 #![deny(clippy::undocumented_unsafe_blocks)]