forge/profile/

samply.rs

// #t(file: rust_alloc_in_loop) CLI symbol resolution caches build owned strings; cold path.
// #t(file: rust_println) CLI binary, terminal output is the purpose.

#![expect(
    clippy::cast_possible_truncation,
    reason = "profiler JSON indices and address tables are bounded by in-memory collection sizes"
)]

use std::{
    io::Read,
    process::{Command, Stdio},
};

use tokio::{
    io::{AsyncBufReadExt, BufReader},
    process::Command as TokioCommand,
};

/// addr2line outputs two lines per address (function name + location).
const ADDR2LINE_LINES_PER_ADDR: usize = 2;
const ATOS_IN_PREFIX_LEN: usize = " (in ".len();
const HEX_RADIX: u32 = 16;

use anyhow::{Context as _, anyhow};
use flate2::read::GzDecoder;
use wowlab_common::output;
use wowlab_fs::{
    directory::{self, EntryKind},
    file,
    path::{Path, PathBuf},
};
use wowlab_types::sim::{FastMap, FastSet};

#[derive(Debug)]
pub(super) struct RawProfile {
    pub samples: Vec<StackSample>,
    pub strings: Vec<String>,
    pub total_weight: u64,
}

#[derive(Debug)]
pub(super) struct StackSample {
    pub frames: Vec<usize>,
    pub weight: u64,
}

#[derive(Debug, thiserror::Error)]
#[error("{cause}{details}", details = SamplyStderr(stderr))]
pub(super) struct RecordError {
    pub cause: String,
    pub stderr: Vec<String>,
}

struct SamplyStderr<'a>(&'a [String]);

impl std::fmt::Display for SamplyStderr<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if !self.0.is_empty() {
            write!(f, "\n--- samply stderr ---")?;

            for line in self.0 {
                write!(f, "\n{line}")?;
            }
        }

        Ok(())
    }
}

pub(super) fn record(binary: &Path, args: &[String], output: &Path) -> Result<(), RecordError> {
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_io()
        .build()
        .map_err(|e| RecordError {
            cause: format!("failed to build tokio runtime: {e}"),
            stderr: Vec::new(),
        })?;

    runtime.block_on(record_async(binary, args, output))
}

async fn record_async(binary: &Path, args: &[String], output: &Path) -> Result<(), RecordError> {
    let mut child = TokioCommand::new("samply")
        .arg("record")
        .arg("-s")
        .arg("-o")
        .arg(output.as_os_str())
        .arg(binary.as_os_str())
        .args(args)
        .env("RUST_LOG", "error")
        .stdout(Stdio::null())
        .stderr(Stdio::piped())
        .spawn()
        .map_err(|e| RecordError {
            cause: format!("failed to run samply: {e}"),
            stderr: Vec::new(),
        })?;

    let stderr_pipe = child.stderr.take().ok_or_else(|| RecordError {
        cause: "samply child has no stderr pipe".to_string(),
        stderr: Vec::new(),
    })?;

    let tee_handle = tokio::spawn(async move {
        let mut buffer: Vec<String> = Vec::new();
        let mut lines = BufReader::new(stderr_pipe).lines();

        while let Ok(Some(line)) = lines.next_line().await {
            eprintln!("{line}");
            buffer.push(line);
            tokio::task::yield_now().await;
        }

        buffer
    });

    let status = child.wait().await.map_err(|e| RecordError {
        cause: format!("failed to wait on samply: {e}"),
        stderr: Vec::new(),
    })?;

    let stderr_buf = tee_handle.await.unwrap_or_default();

    if !status.success() {
        return Err(RecordError {
            cause: "samply exited with non-zero status".to_string(),
            stderr: stderr_buf,
        });
    }

    match directory::inspect(output) {
        Ok(Some(entry)) if entry.kind() == EntryKind::File => {}
        Ok(Some(_)) => {
            return Err(RecordError {
                cause: format!("profile output is not a regular file: {}", output.display()),
                stderr: stderr_buf,
            });
        }
        Ok(None) => {
            return Err(RecordError {
                cause: format!("profile not found at {}", output.display()),
                stderr: stderr_buf,
            });
        }
        Err(error) => {
            return Err(RecordError {
                cause: format!("failed to inspect profile output: {error}"),
                stderr: stderr_buf,
            });
        }
    }

    Ok(())
}

pub(super) fn parse(path: &Path) -> anyhow::Result<RawProfile> {
    let profile = file::open(path).context("cannot open profile")?;
    let mut decoder = GzDecoder::new(profile);
    let mut json_str = String::new();

    decoder
        .read_to_string(&mut json_str)
        .context("cannot decompress profile")?;

    let data: serde_json::Value =
        serde_json::from_str(&json_str).context("invalid profile JSON")?;

    let threads = data
        .get("threads")
        .and_then(|t| t.as_array())
        .ok_or_else(|| anyhow!("no threads in profile"))?;

    let usable_threads = threads.iter().filter(|t| {
        t.get("name")
            .and_then(|n| n.as_str())
            .is_none_or(|n| n != "samply")
    });
    let thread = usable_threads
        .max_by_key(|t| {
            t.get("samples")
                .and_then(|s| s.get("stack"))
                .and_then(|s| s.as_array())
                .map_or(0, Vec::len)
        })
        .ok_or_else(|| anyhow!("no usable thread in profile"))?;

    let strings: Vec<String> = thread
        .get("stringArray")
        .and_then(|a| a.as_array())
        .ok_or_else(|| anyhow!("missing stringArray"))?
        .iter()
        .map(|v| v.as_str().unwrap_or("").to_string())
        .collect();

    let func_names: Vec<usize> = json_usize_array(
        thread
            .get("funcTable")
            .and_then(|t| t.get("name"))
            .ok_or_else(|| anyhow!("missing funcTable.name"))?,
    );
    let frame_funcs: Vec<usize> = json_usize_array(
        thread
            .get("frameTable")
            .and_then(|t| t.get("func"))
            .ok_or_else(|| anyhow!("missing frameTable.func"))?,
    );
    let stack_frames: Vec<usize> = json_usize_array(
        thread
            .get("stackTable")
            .and_then(|t| t.get("frame"))
            .ok_or_else(|| anyhow!("missing stackTable.frame"))?,
    );
    let stack_prefix = thread.get("stackTable").and_then(|t| t.get("prefix"));
    let stack_prefix_values = stack_prefix
        .and_then(|a| a.as_array())
        .ok_or_else(|| anyhow!("missing stackTable.prefix"))?;
    let stack_prefixes: Vec<Option<usize>> = stack_prefix_values
        .iter()
        .map(|v| v.as_u64().map(|n| n as usize))
        .collect();

    let sample_stack = thread.get("samples").and_then(|s| s.get("stack"));
    let stack_values = sample_stack
        .and_then(|a| a.as_array())
        .ok_or_else(|| anyhow!("missing samples.stack"))?;
    let stacks: Vec<Option<usize>> = stack_values
        .iter()
        .map(|v| v.as_u64().map(|n| n as usize))
        .collect();

    let weight_values = thread
        .get("samples")
        .and_then(|s| s.get("weight"))
        .and_then(|a| a.as_array());
    let weights: Vec<u64> = weight_values.map_or_else(
        || vec![1; stacks.len()],
        |arr| arr.iter().map(|v| v.as_u64().unwrap_or(1)).collect(),
    );

    let total_weight: u64 = weights.iter().sum();

    let mut samples = Vec::with_capacity(stacks.len());
    let mut seen = FastSet::default();

    for (i, stack_idx) in stacks.iter().enumerate() {
        let Some(mut curr) = *stack_idx else {
            continue;
        };
        // BOUNDS: `weights` always has the same length as `stacks` (parsed from same array or vec![1; stacks.len()])
        let weight = weights[i];
        let mut frames = Vec::new();

        seen.clear();

        while seen.insert(curr) {
            let Some(&frame_idx) = stack_frames.get(curr) else {
                break;
            };
            let Some(&func_idx) = frame_funcs.get(frame_idx) else {
                break;
            };
            let Some(&name_idx) = func_names.get(func_idx) else {
                break;
            };

            frames.push(name_idx);

            match stack_prefixes.get(curr) {
                Some(Some(parent)) => curr = *parent,
                _ => break,
            }
        }

        samples.push(StackSample { frames, weight });
    }

    Ok(RawProfile {
        samples,
        strings,
        total_weight,
    })
}

pub(super) fn resolve_symbols(
    strings: &[String],
    binary: &Path,
    quiet: bool,
) -> (FastMap<String, String>, u32) {
    let hex_addrs: Vec<&str> = strings
        .iter()
        .filter(|s| s.starts_with("0x"))
        .map(String::as_str)
        .collect();

    if hex_addrs.is_empty() {
        return (FastMap::default(), 0);
    }

    if !quiet {
        output::detail(&format!("Resolving {} addresses...", hex_addrs.len()));
    }

    let mut cache = FastMap::default();

    cache.reserve(hex_addrs.len());

    if cfg!(target_os = "macos") {
        cache = resolve_atos_batch(&hex_addrs, binary);
    } else {
        let batch = resolve_addr2line(&hex_addrs, binary);

        cache.extend(batch);
    }

    let unresolved =
        u32::try_from(cache.values().filter(|v| v.starts_with("0x")).count()).unwrap_or(u32::MAX);

    (cache, unresolved)
}

// Samply stores addresses as file offsets, but atos expects virtual addresses; add this vmaddr to convert.
fn macho_text_vmaddr(binary: &Path) -> Option<u64> {
    let out = Command::new("otool")
        .args(["-l"])
        .arg(binary.as_os_str())
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .output()
        .ok()?;

    let stdout = String::from_utf8_lossy(&out.stdout);
    let mut lines = stdout.lines().peekable();

    while let Some(line) = lines.next() {
        if line.trim().starts_with("segname __TEXT") {
            if let Some(vmaddr_line) = lines.next() {
                let trimmed = vmaddr_line.trim();

                if let Some(hex) = trimmed.strip_prefix("vmaddr ") {
                    return u64::from_str_radix(hex.trim().trim_start_matches("0x"), HEX_RADIX)
                        .ok();
                }
            }

            break;
        }
    }

    None
}

fn resolve_atos_batch(addrs: &[&str], binary: &Path) -> FastMap<String, String> {
    let mut result = FastMap::default();

    result.reserve(addrs.len());
    let vmaddr = macho_text_vmaddr(binary).unwrap_or(0);

    let abs_addrs: Vec<String> = addrs
        .iter()
        .map(|a| {
            let offset = u64::from_str_radix(a.trim_start_matches("0x"), HEX_RADIX).unwrap_or(0);

            format!("0x{:x}", offset + vmaddr)
        })
        .collect();

    let output = Command::new("atos")
        .arg("-o")
        .arg(binary.as_os_str())
        .args(&abs_addrs)
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .output();

    let Ok(output) = output else {
        // #t(block: rust_alloc_in_loop) building owned strings for fallback cache
        for a in addrs {
            result.insert(a.to_string(), a.to_string());
        }

        return result;
    };

    let stdout = String::from_utf8_lossy(&output.stdout);
    let lines: Vec<&str> = stdout.trim().lines().collect();

    for (i, addr) in addrs.iter().enumerate() {
        let sym = lines.get(i).unwrap_or(addr).trim();

        if sym.is_empty() || sym.starts_with("0x") {
            result.insert(addr.to_string(), addr.to_string());
            continue;
        }

        if let Some(idx) = sym.find(" (in ") {
            // BOUNDS: idx from find() is a valid byte offset within sym
            let func = &sym[..idx];
            // BOUNDS: idx from find() + ATOS_IN_PREFIX_LEN (len of " (in ") is within bounds
            let rest = &sym[idx + ATOS_IN_PREFIX_LEN..];
            let loc = rest.split(") ").last().unwrap_or("");

            result.insert(addr.to_string(), format!("{func} {loc}").trim().to_string());
        } else {
            result.insert(addr.to_string(), sym.to_string());
        }
    }

    result
}

fn resolve_addr2line(addrs: &[&str], binary: &Path) -> FastMap<String, String> {
    let mut result = FastMap::default();

    result.reserve(addrs.len());

    let addr2line = which_tool(&["addr2line", "gaddr2line"]);
    let cxxfilt = which_tool(&["c++filt", "gc++filt"]);

    let Some(addr2line) = addr2line else {
        // #t(block: rust_alloc_in_loop) building owned strings for symbol resolution cache
        for a in addrs {
            result.insert(a.to_string(), a.to_string());
        }

        return result;
    };

    let output = Command::new(&addr2line)
        .arg("-f")
        .arg("-e")
        .arg(binary.as_os_str())
        .args(addrs)
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .output();

    let Ok(output) = output else {
        // #t(block: rust_alloc_in_loop) building owned strings for symbol resolution cache
        for a in addrs {
            result.insert(a.to_string(), a.to_string());
        }

        return result;
    };

    let stdout = String::from_utf8_lossy(&output.stdout);
    let lines: Vec<&str> = stdout.lines().collect();

    let mut mangled: Vec<String> = Vec::with_capacity(addrs.len());

    for (i, _) in addrs.iter().enumerate() {
        let idx = i * ADDR2LINE_LINES_PER_ADDR;
        let name = lines.get(idx).copied().unwrap_or("??");

        mangled.push(name.to_string());
    }

    let demangled = if let Some(filt) = cxxfilt {
        let input = mangled.join("\n");
        let filt_out = Command::new(&filt)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::null())
            .spawn()
            .and_then(|mut child| {
                use std::io::Write;

                if let Some(ref mut stdin) = child.stdin {
                    let _ = stdin.write_all(input.as_bytes());
                }

                child.wait_with_output()
            });

        match filt_out {
            Ok(o) => String::from_utf8_lossy(&o.stdout)
                .lines()
                .map(String::from)
                .collect(),
            Err(_) => mangled.clone(),
        }
    } else {
        mangled.clone()
    };

    for (i, addr) in addrs.iter().enumerate() {
        // BOUNDS: `mangled` has exactly `addrs.len()` elements, built in the loop above
        let m = &mangled[i];

        if m == "??" || m.is_empty() {
            result.insert(addr.to_string(), addr.to_string());
        } else {
            let d = demangled.get(i).unwrap_or(m);

            result.insert(addr.to_string(), d.trim().to_string());
        }
    }

    result
}

fn which_tool(names: &[&str]) -> Option<PathBuf> {
    // #t(block: rust_alloc_in_loop) building owned path string from command output
    for name in names {
        if let Ok(output) = Command::new("which")
            .arg(name)
            .stdout(Stdio::piped())
            .stderr(Stdio::null())
            .output()
        {
            if output.status.success() {
                let path = String::from_utf8_lossy(&output.stdout).trim().to_string();

                if !path.is_empty() {
                    return Some(PathBuf::from(path));
                }
            }
        }
    }

    None
}

fn json_usize_array(val: &serde_json::Value) -> Vec<usize> {
    val.as_array()
        .map(|a| a.iter().map(|v| v.as_u64().unwrap_or(0) as usize).collect())
        .unwrap_or_default()
}

pub(super) fn check_samply() -> bool {
    which_tool(&["samply"]).is_some()
}