From 18305eae8e48cdb437224702b7e753d3efefff23 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20Ber=C3=A1nek?= <berykubik@gmail.com>
Date: Fri, 14 Jul 2023 16:07:24 +0200
Subject: [PATCH] Print artifact sizes in `opt-dist`

---
 src/tools/opt-dist/src/main.rs      |  8 ++++++--
 src/tools/opt-dist/src/utils/io.rs  | 16 +++++++++++++++-
 src/tools/opt-dist/src/utils/mod.rs | 26 ++++++++++++++++++++++++--
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/src/tools/opt-dist/src/main.rs b/src/tools/opt-dist/src/main.rs
index 39fa7e1fb2afd..5e82430416b6b 100644
--- a/src/tools/opt-dist/src/main.rs
+++ b/src/tools/opt-dist/src/main.rs
@@ -7,7 +7,9 @@ use crate::tests::run_tests;
 use crate::timer::Timer;
 use crate::training::{gather_llvm_bolt_profiles, gather_llvm_profiles, gather_rustc_profiles};
 use crate::utils::io::reset_directory;
-use crate::utils::{clear_llvm_files, format_env_variables, print_free_disk_space};
+use crate::utils::{
+    clear_llvm_files, format_env_variables, print_binary_sizes, print_free_disk_space,
+};
 
 mod environment;
 mod exec;
@@ -170,6 +172,8 @@ fn main() -> anyhow::Result<()> {
     log::info!("Timer results\n{}", timer.format_stats());
 
     print_free_disk_space()?;
+    result.context("Optimized build pipeline has failed")?;
+    print_binary_sizes(env.as_ref())?;
 
-    result.context("Optimized build pipeline has failed")
+    Ok(())
 }
diff --git a/src/tools/opt-dist/src/utils/io.rs b/src/tools/opt-dist/src/utils/io.rs
index 43546e5fcfa67..aab078067afca 100644
--- a/src/tools/opt-dist/src/utils/io.rs
+++ b/src/tools/opt-dist/src/utils/io.rs
@@ -1,5 +1,5 @@
 use anyhow::Context;
-use camino::Utf8Path;
+use camino::{Utf8Path, Utf8PathBuf};
 use fs_extra::dir::CopyOptions;
 use std::fs::File;
 
@@ -46,3 +46,17 @@ pub fn unpack_archive(path: &Utf8Path, dest_dir: &Utf8Path) -> anyhow::Result<()
     archive.unpack(dest_dir.as_std_path())?;
     Ok(())
 }
+
+/// Returns paths in the given `dir` (non-recursively), optionally with the given `suffix`.
+/// The `suffix` should contain the leading dot.
+pub fn get_files_from_dir(
+    dir: &Utf8Path,
+    suffix: Option<&str>,
+) -> anyhow::Result<Vec<Utf8PathBuf>> {
+    let path = format!("{dir}/*{}", suffix.unwrap_or(""));
+
+    Ok(glob::glob(&path)?
+        .into_iter()
+        .map(|p| p.map(|p| Utf8PathBuf::from_path_buf(p).unwrap()))
+        .collect::<Result<Vec<_>, _>>()?)
+}
diff --git a/src/tools/opt-dist/src/utils/mod.rs b/src/tools/opt-dist/src/utils/mod.rs
index 9305d4989aa76..65e0e8cc4ef4a 100644
--- a/src/tools/opt-dist/src/utils/mod.rs
+++ b/src/tools/opt-dist/src/utils/mod.rs
@@ -1,8 +1,8 @@
 pub mod io;
 
 use crate::environment::Environment;
-use crate::utils::io::delete_directory;
-use humansize::BINARY;
+use crate::utils::io::{delete_directory, get_files_from_dir};
+use humansize::{format_size, BINARY};
 use sysinfo::{DiskExt, RefreshKind, System, SystemExt};
 
 pub fn format_env_variables() -> String {
@@ -25,6 +25,28 @@ pub fn print_free_disk_space() -> anyhow::Result<()> {
     Ok(())
 }
 
+pub fn print_binary_sizes(env: &dyn Environment) -> anyhow::Result<()> {
+    use std::fmt::Write;
+
+    let root = env.build_artifacts().join("stage2");
+
+    let mut files = get_files_from_dir(&root.join("bin"), None)?;
+    files.extend(get_files_from_dir(&root.join("lib"), Some(".so"))?);
+    files.sort_unstable();
+
+    let mut output = String::new();
+    for file in files {
+        let size = std::fs::metadata(file.as_std_path())?.len();
+        let size_formatted = format_size(size, BINARY);
+        let name = format!("{}:", file.file_name().unwrap());
+        writeln!(output, "{name:<50}{size_formatted:>10}")?;
+    }
+
+    log::info!("Rustc artifact size\n{output}");
+
+    Ok(())
+}
+
 pub fn clear_llvm_files(env: &dyn Environment) -> anyhow::Result<()> {
     // Bootstrap currently doesn't support rebuilding LLVM when PGO options
     // change (or any other llvm-related options); so just clear out the relevant