cranelift: per-table mutability tracking + call_indirect elisions on immutable funcref tables

crates/cranelift/src/func_environ.rs

@@ -1652,8 +1652,11 @@ impl FuncEnvironment<'_> {
         .unwrap();
 
         // A fixed-size table can't be resized, so its base address won't change
-        // and its base load is `readonly` and `can_move`.
-        let is_static = Some(table.limits.min) == table.limits.max;
+        // and its base load is `readonly` and `can_move`. That holds when the
+        // table is declared with `min == max`, or when translation proved it
+        // is never grown or mutated.
+        let is_static =
+            Some(table.limits.min) == table.limits.max || !self.translation.tables_mutated[index];
         let mut base_flags = ir::MemFlagsData::trusted();
         if is_static {
             base_flags = base_flags.with_readonly().with_can_move();
@@ -2147,6 +2150,14 @@ impl<'a, 'func, 'module_env> Call<'a, 'func, 'module_env> {
         callee: ir::Value,
         call_args: &[ir::Value],
     ) -> WasmResult<Option<CallRets>> {
+        // Fast path: if we can statically resolve this indirect call to a
+        // single defined function (immutable funcref table + constant
+        // callee index + matching signature), emit a direct call instead.
+        // See `try_static_resolve_indirect_call`.
+        if let Some(target) = self.try_static_resolve_indirect_call(table_index, ty_index, callee) {
+            return self.direct_call(target, sig_ref, call_args).map(Some);
+        }
+
         let (code_ptr, callee_vmctx) = match self.check_and_load_code_and_callee_vmctx(
             table_index,
             ty_index,
@@ -2161,6 +2172,198 @@ impl<'a, 'func, 'module_env> Call<'a, 'func, 'module_env> {
             .map(Some)
     }
 
+    /// Try to statically resolve a `call_indirect` site to a single defined
+    /// function so the call can be lowered as a direct call.
+    ///
+    /// All four of these must hold for the resolution to succeed:
+    ///
+    /// 1. The target table must be provably immutable for the lifetime of
+    ///    any instance of this module: defined (not imported) and never the
+    ///    target of `table.set` / `table.fill` / `table.copy` (as the dst)
+    ///    / `table.grow` / `table.init`. This is the `tables_mutated` bit
+    ///    populated in `ModuleEnvironment::translate`.
+    ///
+    /// 2. The callee index value (the operand to `call_indirect`) must be a
+    ///    compile-time constant — i.e., the wasm did `i32.const N;
+    ///    call_indirect (table $t) (type $sig)`. This is what hand-lowered
+    ///    C++/Rust vtable calls and AOT-compiled JS-to-wasm dispatch tables
+    ///    look like in practice.
+    ///
+    /// 3. The slot at index `N` in the table must be precomputable from
+    ///    static `elem` segments: `module.table_initialization
+    ///    .initial_values[defined_index]` must be `TableInitialValue::Null
+    ///    { precomputed }` (i.e., not a fully-dynamic `Expr`-style init),
+    ///    and the index `N` must be in range and resolved to a concrete
+    ///    `FuncIndex` (not the reserved-value sentinel).
+    ///
+    /// 4. The function's signature in the module's interned type table
+    ///    must equal the `ty_index` declared by the `call_indirect` site.
+    ///    Otherwise the original semantics are "trap on signature
+    ///    mismatch", which we don't want to replace with a static direct
+    ///    call.
+    ///
+    /// Returns the resolved function on success, `None` otherwise (in
+    /// which case the caller falls back to a normal indirect call).
+    fn try_static_resolve_indirect_call(
+        &self,
+        table_index: TableIndex,
+        ty_index: TypeIndex,
+        callee: ir::Value,
+    ) -> Option<FuncIndex> {
+        let translation = self.env.translation;
+        let module = &translation.module;
+
+        // (1) Table must be provably immutable. Imported tables are
+        //     pre-marked as mutated in `ModuleEnvironment::translate`, so
+        //     this check also rules them out (along with the explicit
+        //     `defined_table_index` check below for clarity).
+        if translation.tables_mutated[table_index] {
+            return None;
+        }
+        let defined_table = module.defined_table_index(table_index)?;
+
+        // (2) Callee must be a constant `iconst`. Pattern adapted from
+        //     `bounds_checks::statically_known_in_bounds`.
+        let dfg = &self.builder.func.dfg;
+        let inst = dfg.value_def(callee).inst()?;
+        let imm = match dfg.insts[inst] {
+            ir::InstructionData::UnaryImm {
+                opcode: ir::Opcode::Iconst,
+                imm,
+            } => imm,
+            _ => return None,
+        };
+        let callee_ty = dfg.value_type(callee);
+        let callee_idx_u64 = imm
+            .zero_extend_from_width(callee_ty.bits())
+            .bits()
+            .cast_unsigned();
+
+        // (3) Slot must be precomputable from the static funcref image.
+        let precomputed = module.table_initialization.get(defined_table)?;
+        let slot = usize::try_from(callee_idx_u64).ok()?;
+        if slot >= precomputed.len() {
+            return None;
+        }
+        let target = precomputed[slot];
+        // `FuncIndex::reserved_value()` marks a null (uncovered) slot.
+        if target.is_reserved_value() {
+            return None;
+        }
+
+        // (4) Signature match. The site's declared `ty_index` and the
+        //     target function's declared signature must intern to the same
+        //     module type index.
+        let expected_ty = module.types[ty_index].unwrap_module_type_index();
+        let target_ty = module.functions[target]
+            .signature
+            .unwrap_module_type_index();
+        if expected_ty != target_ty {
+            return None;
+        }
+
+        Some(target)
+    }
+
+    /// Try to prove that the runtime signature check at a `call_indirect`
+    /// site through an untyped `funcref` table is redundant.
+    ///
+    /// True when:
+    ///
+    /// 1. The table is provably immutable (`tables_mutated[table_index] ==
+    ///    false`). Defined-not-imported is implied since imported tables
+    ///    are pre-marked as mutated.
+    ///
+    /// 2. The table is precomputable from static `elem` segments
+    ///    (`TableInitialValue::Null { precomputed }`).
+    ///
+    /// 3. Every non-null entry in `precomputed` has the same module-
+    ///    interned signature as the `ty_index` declared at the call site.
+    ///    Null slots are fine — they trap on the funcref-NULL load that
+    ///    happens after sig-check elision.
+    ///
+    /// When this returns true, the caller short-circuits to
+    /// `CheckIndirectCallTypeSignature::StaticMatch`, which removes the
+    /// sig load + compare from the hot path. Bounds-check on the table
+    /// index and the funcref-NULL check are still emitted by the
+    /// surrounding code, so the call still traps correctly on OOB or
+    /// null index — only the sig check is elided.
+    ///
+    /// This is the static analog of an inline-cache: instead of caching
+    /// the resolved target per call site, we observe at module-load that
+    /// the table contents make the sig check uninformative for the
+    /// lifetime of any instance.
+    /// True iff every slot in the precomputed `elem`-segment contents for
+    /// `table_index` is a concrete `FuncIndex` (no
+    /// `FuncIndex::reserved_value()` "no-entry" sentinel).
+    ///
+    /// Caller has already proven the table is immutable, so the contents
+    /// observed here are stable for the lifetime of any instance —
+    /// `false` here implies "no slot is ever null at runtime."
+    ///
+    /// When this is true, the runtime funcref-NULL check on the loaded
+    /// funcref pointer is provably redundant: any in-bounds index leads
+    /// to a non-null funcref. The bounds check still runs (so an
+    /// out-of-bounds index traps as before with `TRAP_TABLE_OUT_OF_BOUNDS`).
+    fn precomputed_table_has_no_null_slots(&self, table_index: TableIndex) -> bool {
+        let module = &self.env.translation.module;
+        let Some(defined_table) = module.defined_table_index(table_index) else {
+            return false;
+        };
+        let Some(precomputed) = module.table_initialization.get(defined_table) else {
+            return false;
+        };
+        if precomputed.is_empty() {
+            return false;
+        }
+        // Slots beyond `precomputed.len()` are null at runtime; coverage
+        // up to `limits.min` is required (caller proved immutable, so the
+        // table can't grow beyond min).
+        let table_min = module.tables[table_index].limits.min;
+        if (precomputed.len() as u64) < table_min {
+            return false;
+        }
+        precomputed.iter().all(|f| !f.is_reserved_value())
+    }
+
+    fn try_elide_sig_check_for_immutable_table(
+        &self,
+        table_index: TableIndex,
+        ty_index: TypeIndex,
+    ) -> bool {
+        let translation = self.env.translation;
+        let module = &translation.module;
+
+        if translation.tables_mutated[table_index] {
+            return false;
+        }
+        let defined_table = match module.defined_table_index(table_index) {
+            Some(d) => d,
+            None => return false,
+        };
+
+        let precomputed = match module.table_initialization.get(defined_table) {
+            Some(p) if !p.is_empty() => p,
+            _ => return false,
+        };
+
+        let expected_ty = module.types[ty_index].unwrap_module_type_index();
+        for &func_idx in precomputed.iter() {
+            // Null slots will trap on the funcref-NULL load anyway.
+            if func_idx.is_reserved_value() {
+                continue;
+            }
+            let actual_ty = module.functions[func_idx]
+                .signature
+                .unwrap_module_type_index();
+            if actual_ty != expected_ty {
+                return false;
+            }
+        }
+
+        true
+    }
+
     fn check_and_load_code_and_callee_vmctx(
         &mut self,
         table_index: TableIndex,
@@ -2218,6 +2421,34 @@ impl<'a, 'func, 'module_env> Call<'a, 'func, 'module_env> {
         // table of typed functions and that type matches `ty_index`, then
         // there's no need to perform a typecheck.
         match table.ref_type.heap_type {
+            // Untyped `funcref` tables ordinarily need a runtime sig check.
+            // But if (a) the table is provably immutable (`tables_mutated`
+            // bit clear) and (b) every non-null entry in the precomputed
+            // static `elem` segments has the same `VMSharedTypeIndex` as
+            // the call site, then the runtime check is provably redundant
+            // and we can elide it the same way we do for typed-funcref
+            // tables.
+            //
+            // This is the AOT-IC-seeding analog: instead of caching the
+            // resolved target at the call site, we cache the *signature*
+            // at module-load time and skip the hot-path sig load+compare.
+            // Helps the megamorphic case (computed `call_indirect` index)
+            // that the static-monomorphization fast path above can't
+            // handle.
+            WasmHeapType::Func
+                if self.try_elide_sig_check_for_immutable_table(table_index, ty_index) =>
+            {
+                // If we additionally know every entry in the precomputed
+                // table is non-null, lower `may_be_null` to false so the
+                // downstream funcref-NULL check is also elided. This is
+                // only sound if the table can't be grown or have its
+                // entries cleared after init (i.e., immutable, which we
+                // already proved above).
+                let may_be_null = table.ref_type.nullable
+                    && !self.precomputed_table_has_no_null_slots(table_index);
+                return CheckIndirectCallTypeSignature::StaticMatch { may_be_null };
+            }
+
             // Functions do not have a statically known type in the table, a
             // typecheck is required. Fall through to below to perform the
             // actual typecheck.

crates/environ/src/compile/module_environ.rs

@@ -110,6 +110,26 @@ pub struct ModuleTranslation<'data> {
     /// trampolines for each of these signatures are required.
     pub exported_signatures: Vec<ModuleInternedTypeIndex>,
 
+    /// Per-table flag indicating whether the table is ever mutated by any
+    /// function defined in this module via `table.set` / `table.fill` /
+    /// `table.copy` (as the destination) / `table.grow` / `table.init`.
+    ///
+    /// `false` (the default) means the table's contents are determined
+    /// entirely by its `elem` segments and any active initializer, and never
+    /// change at runtime — provably immutable for the lifetime of any
+    /// instance of this module.
+    ///
+    /// `true` means the contents can change at runtime (or the table is
+    /// imported, in which case we conservatively assume the importer
+    /// mutates it).
+    ///
+    /// This is groundwork for later passes that turn `call_indirect`
+    /// through provably-immutable function tables into direct calls when
+    /// the dispatched-to slot is statically known. Set during module
+    /// translation (see `analyze_table_mutability`); read by Cranelift
+    /// lowering and by Pulley AOT IC seeding.
+    pub tables_mutated: SecondaryMap<TableIndex, bool>,
+
     /// DWARF debug information, if enabled, parsed from the module.
     pub debuginfo: DebugInfoData<'data>,
 
@@ -227,6 +247,7 @@ impl<'data> ModuleTranslation<'data> {
             function_body_inputs: PrimaryMap::default(),
             known_imported_functions: SecondaryMap::default(),
             exported_signatures: Vec::default(),
+            tables_mutated: SecondaryMap::default(),
             debuginfo: DebugInfoData::default(),
             has_unparsed_debuginfo: false,
             data_align: None,
@@ -349,6 +370,8 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
             self.translate_payload(payload?)?;
         }
 
+        analyze_table_mutability(&mut self.result)?;
+
         Ok(self.result)
     }
 
@@ -1582,3 +1605,83 @@ impl ModuleTranslation<'_> {
         self.module.startup = ModuleStartup::IfMemoriesNeedInit(ty);
     }
 }
+
+/// Walk every defined function body, recording in
+/// `translation.tables_mutated` each table that is the destination of any
+/// runtime mutation opcode (`table.set`, `table.fill`, `table.copy` as the
+/// destination, `table.grow`, `table.init`).
+///
+/// Imported tables are conservatively pre-marked as mutated since the
+/// importer can mutate them in ways we can't see. Active `elem` segments
+/// applied at instantiation time are NOT counted as mutations — they are
+/// part of the table's *initial* state, not a runtime change.
+///
+/// `elem.drop` drops a passive element segment but does not write to any
+/// table directly, so it is intentionally not counted here. Conservatively,
+/// any `table.init` from a passive segment marks the destination table as
+/// mutated.
+fn analyze_table_mutability<'data>(translation: &mut ModuleTranslation<'data>) -> Result<()> {
+    // Resize the table-mutability map to cover every table in the module
+    // (imports + defined). `SecondaryMap` defaults to `false` for all
+    // unset entries, which is the correct "definitely-not-mutated" default
+    // for defined tables we haven't observed any mutations on yet.
+    let num_tables = translation.module.tables.len();
+    if num_tables == 0 {
+        return Ok(());
+    }
+
+    // Mark all imported tables as mutated up front. The importer can
+    // mutate them in ways this module can't see, so the conservative
+    // assumption is that they are not stable across calls.
+    let num_imported = translation.module.num_imported_tables;
+    for i in 0..num_imported {
+        translation.tables_mutated[TableIndex::from_u32(i as u32)] = true;
+    }
+
+    // Mark all *exported* tables as mutated as well. A host (or another
+    // instance importing the export) can call `Table::set` /
+    // `Table::grow` via the public wasmtime API on any exported table,
+    // and those mutations are not visible in this module's bytecode.
+    // The `call_indirect` optimizations that read this bit must
+    // therefore treat exported tables as conservatively non-stable.
+    for (_, entity_index) in &translation.module.exports {
+        if let EntityIndex::Table(table_index) = entity_index {
+            translation.tables_mutated[*table_index] = true;
+        }
+    }
+
+    // Walk every defined function body and look for table-mutation opcodes.
+    // The cost is O(total opcodes), one extra pass on top of the validator;
+    // typical large modules (sqlite3 ~50K opcodes) take well under a
+    // millisecond.
+    for (_, body_data) in &translation.function_body_inputs {
+        let mut reader = body_data.body.get_operators_reader()?;
+        while !reader.eof() {
+            use wasmparser::Operator;
+            match reader.read()? {
+                Operator::TableSet { table }
+                | Operator::TableFill { table }
+                | Operator::TableGrow { table } => {
+                    translation.tables_mutated[TableIndex::from_u32(table)] = true;
+                }
+                Operator::TableCopy {
+                    dst_table,
+                    src_table: _,
+                } => {
+                    // `src_table` is read-only in `table.copy`; only the
+                    // destination is mutated.
+                    translation.tables_mutated[TableIndex::from_u32(dst_table)] = true;
+                }
+                Operator::TableInit {
+                    table,
+                    elem_index: _,
+                } => {
+                    translation.tables_mutated[TableIndex::from_u32(table)] = true;
+                }
+                _ => {}
+            }
+        }
+    }
+
+    Ok(())
+}