From 5f2f2a4bb4523821a8de21e7c4788daf849e0474 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Schr=C3=B6der?= <info@swedishembedded.com>
Date: Mon, 22 Jun 2026 20:41:19 +0000
Subject: [PATCH] Cranelift: use a BTreeMap for JITModule code ranges

`JITModule` resolves a PC back to its defining function (for exception
unwinding) via a `Vec<(start, end, FuncId)>` that was re-sorted in full on
every `finalize_definitions` call. Code that defines and finalizes functions
one at a time therefore paid an O(n) sort per finalize, i.e. O(n^2) overall.

Replace the vector with a `BTreeMap` keyed on the start address. Inserts stay
sorted incrementally (O(log n)) and the lookup becomes an O(log n) range
query, so finalizing is linear in the number of functions with no sort.
---
 cranelift/jit/src/backend.rs | 35 ++++++++++-------------------------
 1 file changed, 10 insertions(+), 25 deletions(-)

diff --git a/cranelift/jit/src/backend.rs b/cranelift/jit/src/backend.rs
index cbed2c3217bf..2f04f67db524 100644
--- a/cranelift/jit/src/backend.rs
+++ b/cranelift/jit/src/backend.rs
@@ -16,6 +16,7 @@ use cranelift_module::{
 };
 use log::info;
 use std::cell::RefCell;
+use std::collections::BTreeMap;
 use std::collections::HashMap;
 use std::ffi::CString;
 use std::io::Write;
@@ -175,7 +176,9 @@ pub struct JITModule {
     declarations: ModuleDeclarations,
     compiled_functions: SecondaryMap<FuncId, Option<CompiledBlob>>,
     compiled_data_objects: SecondaryMap<DataId, Option<CompiledBlob>>,
-    code_ranges: Vec<(usize, usize, FuncId)>,
+    /// Map from a function's start address to its (end address, FuncId), used
+    /// to resolve a PC back to its function for exception unwinding.
+    code_ranges: BTreeMap<usize, (usize, FuncId)>,
     functions_to_finalize: Vec<FuncId>,
     data_objects_to_finalize: Vec<DataId>,
 }
@@ -334,9 +337,6 @@ impl JITModule {
             data.perform_relocations(|name| self.get_address(name));
         }
 
-        self.code_ranges
-            .sort_unstable_by_key(|(start, _end, _)| *start);
-
         // Now that we're done patching, prepare the memory for execution!
         let branch_protection = if cfg!(target_arch = "aarch64") && use_bti(&self.isa.isa_flags()) {
             BranchProtection::BTI
@@ -367,7 +367,7 @@ impl JITModule {
             declarations: ModuleDeclarations::default(),
             compiled_functions: SecondaryMap::new(),
             compiled_data_objects: SecondaryMap::new(),
-            code_ranges: Vec::new(),
+            code_ranges: BTreeMap::new(),
             functions_to_finalize: Vec::new(),
             data_objects_to_finalize: Vec::new(),
         }
@@ -381,24 +381,10 @@ impl JITModule {
         &'a self,
         pc: usize,
     ) -> Option<(usize, wasmtime_unwinder::ExceptionTable<'a>)> {
-        // Search the sorted code-ranges for the PC.
-        let idx = match self
-            .code_ranges
-            .binary_search_by_key(&pc, |(start, _end, _func)| *start)
-        {
-            Ok(exact_start_match) => Some(exact_start_match),
-            Err(least_upper_bound) if least_upper_bound > 0 => {
-                let last_range_before_pc = &self.code_ranges[least_upper_bound - 1];
-                if last_range_before_pc.0 <= pc && pc < last_range_before_pc.1 {
-                    Some(least_upper_bound - 1)
-                } else {
-                    None
-                }
-            }
-            _ => None,
-        }?;
-
-        let (start, _, func) = self.code_ranges[idx];
+        let (&start, &(end, func)) = self.code_ranges.range(..=pc).next_back()?;
+        if pc >= end {
+            return None;
+        }
 
         // Get the ExceptionTable. The "parse" here simply reads two
         // u32s for lengths and constructs borrowed slices, so it's
@@ -521,8 +507,7 @@ impl Module for JITModule {
 
         let range_start = ptr.addr();
         let range_end = range_start + size;
-        // These will be sorted when we finalize.
-        self.code_ranges.push((range_start, range_end, id));
+        self.code_ranges.insert(range_start, (range_end, id));
 
         self.functions_to_finalize.push(id);