Add ast and hir stats to print-codegen-stats-json flag

Cargo.lock

@@ -4233,6 +4233,7 @@ dependencies = [
  "rustc_trait_selection",
  "rustc_traits",
  "rustc_ty_utils",
+ "serde_json",
  "tracing",
 ]
 

compiler/rustc_interface/Cargo.toml

@@ -45,6 +45,7 @@ rustc_thread_pool = { path = "../rustc_thread_pool" }
 rustc_trait_selection = { path = "../rustc_trait_selection" }
 rustc_traits = { path = "../rustc_traits" }
 rustc_ty_utils = { path = "../rustc_ty_utils" }
+serde_json = "1.0.59"
 tracing = "0.1"
 # tidy-alphabetical-end
 

compiler/rustc_interface/src/passes.rs

@@ -406,6 +406,10 @@ fn early_lint_checks(tcx: TyCtxt<'_>, (): ()) {
         input_stats::print_ast_stats(tcx, krate);
     }
 
+    if sess.print_llvm_stats_json().is_some() {
+        input_stats::collect_ast_stats(tcx, krate);
+    }
+
     // Needs to go *after* expansion to be able to check the results of macro expansion.
     sess.time("complete_gated_feature_checking", || {
         rustc_ast_passes::feature_gate::check_crate(krate, sess, tcx.features());
@@ -1074,6 +1078,9 @@ fn run_required_analyses(tcx: TyCtxt<'_>) {
     if tcx.sess.opts.unstable_opts.input_stats {
         rustc_passes::input_stats::print_hir_stats(tcx);
     }
+    if tcx.sess.print_llvm_stats_json().is_some() {
+        rustc_passes::input_stats::collect_hir_stats(tcx);
+    }
     // When using rustdoc's "jump to def" feature, it enters this code and `check_crate`
     // is not defined. So we need to cfg it out.
     #[cfg(all(not(doc), debug_assertions))]

compiler/rustc_interface/src/queries.rs

@@ -73,14 +73,47 @@ impl Linker {
         if let Some(out_path) = sess.print_llvm_stats_json() {
             let llvm_stats_json = codegen_backend.print_statistics_json();
 
+            if llvm_stats_json.is_empty() {
+                sess.dcx().warn(format!(
+                    "requested to print LLVM statistics to JSON file {}, but the codegen backend did not provide any statistics",
+                    out_path,
+                ));
+            }
+
+            let mut merged_stats = serde_json::Map::new();
+
+            // Parse LLVM stats if present
             if !llvm_stats_json.is_empty() {
-                if let Err(e) = std::fs::write(&out_path, llvm_stats_json) {
-                    sess.dcx().err(format!("failed to write stats to {}: {}", out_path, e));
+                match serde_json::from_str::<serde_json::Value>(&llvm_stats_json) {
+                    Ok(serde_json::Value::Object(map)) => {
+                        merged_stats = map;
+                    }
+                    Ok(_) => {
+                        sess.dcx().warn("LLVM statistics JSON was not a valid JSON object");
+                    }
+                    Err(e) => {
+                        sess.dcx().warn(format!("failed to parse LLVM statistics JSON: {}", e));
+                    }
+                }
+            }
+
+            // Append frontend stats
+            let frontend_stats = sess.frontend_stats.lock();
+            for (key, value) in frontend_stats.iter() {
+                merged_stats.insert(key.clone(), serde_json::Value::Number((*value).into()));
+            }
+
+            if !merged_stats.is_empty() {
+                if let Ok(final_json) =
+                    serde_json::to_string_pretty(&serde_json::Value::Object(merged_stats))
+                {
+                    if let Err(e) = std::fs::write(&out_path, final_json) {
+                        sess.dcx().err(format!("failed to write stats to {}: {}", out_path, e));
+                    }
                 }
             } else {
                 sess.dcx().warn(format!(
-                    "requested to print LLVM statistics to JSON file {}, but the codegen backend \
-                    did not provide any statistics",
+                    "requested to print statistics to JSON file {}, but no statistics were collected",
                     out_path,
                 ));
             }

compiler/rustc_passes/src/input_stats.rs

@@ -197,6 +197,63 @@ impl<'k> StatCollector<'k> {
         _ = writeln!(s, "{prefix} {}", "=".repeat(banner_w));
         eprint!("{s}");
     }
+    fn to_json_entries(&self, prefix: &str) -> Vec<(String, usize)> {
+        let mut out = Vec::new();
+        // We will soon sort, so the initial order does not matter.
+        #[allow(rustc::potential_query_instability)]
+        let mut keys: Vec<_> = self.nodes.keys().collect();
+        keys.sort(); // Ensure deterministic output
+        let mut total_size = 0;
+        let mut total_count = 0;
+        for label in keys {
+            let node = &self.nodes[label];
+            total_size += node.stats.accum_size();
+            total_count += node.stats.count;
+            out.push((format!("{prefix}.{label}.count"), node.stats.count));
+            out.push((format!("{prefix}.{label}.size"), node.stats.size));
+            out.push((format!("{prefix}.{label}.cumulative_size"), node.stats.accum_size()));
+            // We will soon sort, so the initial order does not matter.
+            #[allow(rustc::potential_query_instability)]
+            let mut sub_keys: Vec<_> = node.subnodes.keys().collect();
+            sub_keys.sort();
+            for sub_label in sub_keys {
+                let subnode = &node.subnodes[sub_label];
+                out.push((format!("{prefix}.{label}.{sub_label}.count"), subnode.count));
+                out.push((format!("{prefix}.{label}.{sub_label}.size"), subnode.size));
+                out.push((
+                    format!("{prefix}.{label}.{sub_label}.cumulative_size"),
+                    subnode.accum_size(),
+                ));
+            }
+        }
+        out.push((format!("{prefix}.total.count"), total_count));
+        out.push((format!("{prefix}.total.size"), total_size));
+        out
+    }
+}
+
+pub fn collect_hir_stats(tcx: TyCtxt<'_>) {
+    if tcx.sess.print_llvm_stats_json().is_none() {
+        return;
+    }
+    let mut collector =
+        StatCollector { tcx: Some(tcx), nodes: FxHashMap::default(), seen: FxHashSet::default() };
+    tcx.hir_walk_toplevel_module(&mut collector);
+    tcx.hir_walk_attributes(&mut collector);
+    let mut stats = tcx.sess.frontend_stats.lock();
+    stats.extend(collector.to_json_entries("hir-stats"));
+}
+
+pub fn collect_ast_stats(tcx: TyCtxt<'_>, krate: &ast::Crate) {
+    if tcx.sess.print_llvm_stats_json().is_none() {
+        return;
+    }
+    use rustc_ast::visit::Visitor;
+    let mut collector =
+        StatCollector { tcx: None, nodes: FxHashMap::default(), seen: FxHashSet::default() };
+    collector.visit_crate(krate);
+    let mut stats = tcx.sess.frontend_stats.lock();
+    stats.extend(collector.to_json_entries("ast-stats"));
 }
 
 // Used to avoid boilerplate for types with many variants.

compiler/rustc_session/src/session.rs

@@ -178,6 +178,9 @@ pub struct Session {
     ///
     /// The value is the `DepNodeIndex` of the node encodes the used feature.
     pub used_features: Lock<FxHashMap<Symbol, u32>>,
+
+    /// Frontend stats collected for JSON output.
+    pub frontend_stats: Lock<Vec<(String, usize)>>,
 }
 
 #[derive(Clone, Copy)]
@@ -1127,6 +1130,7 @@ pub fn build_session(
         thin_lto_supported: true,                  // filled by `run_compiler`
         mir_opt_bisect_eval_count: AtomicUsize::new(0),
         used_features: Lock::default(),
+        frontend_stats: Lock::default(),
     };
 
     validate_commandline_args_with_session_available(&sess);