From 32e050714bbf523aedbcdf705c1bcdc81b035d48 Mon Sep 17 00:00:00 2001 From: hellovai Date: Wed, 15 Jul 2026 18:42:32 -0700 Subject: [PATCH] perf(cli): mimalloc global allocator + build ariadne SourceCache once per diagnostic batch Two independent wins for cold / wall time, re-derived from the compiler2 cold-compile audit (#4016): - baml_cli now sets mimalloc as its global allocator. The compiler workload is dominated by small short-lived allocations (Ty trees, Vecs, SmolStrs); system malloc measured ~35% of remaining single-threaded CPU in the original audit. - render_diagnostics builds the ariadne SourceCache (a line index over every project file) once per batch instead of once per diagnostic, threading it as &mut through render_ariadne / render_report_to_string. Rendered bytes are unchanged, so cached diagnostics from the #3924 disk cache replay identically (BAML_CACHE_VERIFY verified). Co-authored-by: Cursor --- baml_language/Cargo.lock | 19 ++++++++++ baml_language/Cargo.toml | 1 + baml_language/crates/baml_cli/Cargo.toml | 1 + baml_language/crates/baml_cli/src/main.rs | 7 ++++ .../baml_compiler_diagnostics/src/render.rs | 36 ++++++++++++------- 5 files changed, 51 insertions(+), 13 deletions(-) diff --git a/baml_language/Cargo.lock b/baml_language/Cargo.lock index 2bd874d7f6a..69150ea558a 100644 --- a/baml_language/Cargo.lock +++ b/baml_language/Cargo.lock @@ -747,6 +747,7 @@ dependencies = [ "indicatif", "insta", "libsui", + "mimalloc", "object 0.36.3", "regex", "reqwest", @@ -4263,6 +4264,15 @@ version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" +[[package]] +name = "libmimalloc-sys" +version = "0.1.49" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9" +dependencies = [ + "cc", +] + [[package]] name = "libredox" version = "0.1.17" @@ -4515,6 +4525,15 @@ dependencies = [ "paste", ] +[[package]] +name = "mimalloc" +version = "0.1.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d4139bb28d14ad1facf21d5eb8825051b326e172d216b39f6d31df53cc97862" +dependencies = [ + "libmimalloc-sys", +] + [[package]] name = "mime" version = "0.3.17" diff --git a/baml_language/Cargo.toml b/baml_language/Cargo.toml index 2224f93129b..73d59c06abf 100644 --- a/baml_language/Cargo.toml +++ b/baml_language/Cargo.toml @@ -203,6 +203,7 @@ bytecount = "0.6" reqwest = { version = "0.13.1", default-features = false, features = [ "stream", ] } +mimalloc = { version = "0.1.52" } rowan = { version = "0.16" } rustls = { version = "0.23.37", default-features = false } rustls-pemfile = "2" diff --git a/baml_language/crates/baml_cli/Cargo.toml b/baml_language/crates/baml_cli/Cargo.toml index abd41b8707e..a35e94f2dea 100644 --- a/baml_language/crates/baml_cli/Cargo.toml +++ b/baml_language/crates/baml_cli/Cargo.toml @@ -44,6 +44,7 @@ hex = { workspace = true } indexmap = { workspace = true } indicatif = { workspace = true } libsui = { workspace = true } +mimalloc = { workspace = true } object = { workspace = true } regex = { workspace = true } reqwest = { workspace = true, features = [ "blocking", "json", "rustls" ] } diff --git a/baml_language/crates/baml_cli/src/main.rs b/baml_language/crates/baml_cli/src/main.rs index 38047ce2492..2ef6fff5b4d 100644 --- a/baml_language/crates/baml_cli/src/main.rs +++ b/baml_language/crates/baml_cli/src/main.rs @@ -5,6 +5,13 @@ use std::io::Write as _; +/// The compiler workload is dominated by small short-lived allocations +/// (`Ty` trees, `Vec`s, `SmolStr`s): system malloc measured ~35% of +/// remaining single-threaded CPU in the cold-compile audit. mimalloc +/// substantially cuts cold `baml check` / `baml build` wall time. +#[global_allocator] +static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; + fn main() { // TODO: baml_log is disabled for now // baml_log::init()?; diff --git a/baml_language/crates/baml_compiler_diagnostics/src/render.rs b/baml_language/crates/baml_compiler_diagnostics/src/render.rs index bb974955b56..e41af4a63ee 100644 --- a/baml_language/crates/baml_compiler_diagnostics/src/render.rs +++ b/baml_language/crates/baml_compiler_diagnostics/src/render.rs @@ -208,7 +208,10 @@ pub fn render_diagnostic( config: &RenderConfig, ) -> String { match config.format { - DiagnosticFormat::Ariadne => render_ariadne(diagnostic, sources, file_paths, config.color), + DiagnosticFormat::Ariadne => { + let mut cache = SourceCache::new(sources.clone(), file_paths.clone()); + render_ariadne(diagnostic, sources, &mut cache, config.color) + } DiagnosticFormat::Concise => render_concise(diagnostic, sources, file_paths), } } @@ -223,9 +226,20 @@ pub fn render_diagnostics( file_paths: &HashMap, config: &RenderConfig, ) -> String { + // Build the ariadne source cache ONCE for the whole batch. Constructing + // it per diagnostic (each `Source::from` computes a line index for every + // file in the project) is quadratic in practice and measurably slows + // `baml check` on warning-heavy projects. + let mut ariadne_cache = matches!(config.format, DiagnosticFormat::Ariadne) + .then(|| SourceCache::new(sources.clone(), file_paths.clone())); diagnostics .iter() - .map(|d| render_diagnostic(d, sources, file_paths, config)) + .map(|d| match (&config.format, &mut ariadne_cache) { + (DiagnosticFormat::Ariadne, Some(cache)) => { + render_ariadne(d, sources, cache, config.color) + } + _ => render_concise(d, sources, file_paths), + }) .collect::>() .join("\n") } @@ -259,10 +273,13 @@ fn translate_span(span: Span, sources: &HashMap) -> Span { } /// Render a diagnostic using Ariadne (pretty CLI output). +/// +/// Takes a pre-built [`SourceCache`] so batch rendering shares one line-index +/// computation across all diagnostics. fn render_ariadne( diagnostic: &Diagnostic, sources: &HashMap, - file_paths: &HashMap, + cache: &mut SourceCache, color: bool, ) -> String { // BAML CLI uses lowercase `error:` / `warning:` everywhere (matching @@ -321,7 +338,7 @@ fn render_ariadne( .finish(); // Render to string using SourceCache for proper filename display. - let rendered = render_report_to_string(&report, sources, file_paths); + let rendered = render_report_to_string(&report, cache); // See the rant above the `report_kind` match — ariadne paints the // `Custom` keyword unconditionally, so `with_color(false)` doesn't @@ -427,17 +444,10 @@ fn render_concise( } /// Render an ariadne Report to a String using `SourceCache` for proper filename display. -fn render_report_to_string( - report: &Report<'_, Span>, - sources: &HashMap, - file_paths: &HashMap, -) -> String { +fn render_report_to_string(report: &Report<'_, Span>, cache: &mut SourceCache) -> String { let mut output = Vec::new(); - // Use SourceCache for proper filename display - let mut cache = SourceCache::new(sources.clone(), file_paths.clone()); - - report.write(&mut cache, &mut output).unwrap_or_else(|_| { + report.write(cache, &mut output).unwrap_or_else(|_| { output.clear(); output.extend_from_slice(b""); });