[NOT READY] feat(profiling): add --show-from filter to profiling aggregate

src/commands/profiling.rs

@@ -57,6 +57,7 @@ pub async fn aggregate(
     to: String,
     limit: u32,
     aggregation_function: String,
+    show_from: Option<String>,
 ) -> Result<()> {
     let (from_iso, to_iso) = parse_window(&from, &to)?;
     // /profiling/api/v1/aggregate expects a flat body — query/from/to are siblings, not wrapped in filter{}.
@@ -68,12 +69,296 @@ pub async fn aggregate(
         "limit": limit,
         "aggregationFunction": aggregation_function,
     });
-    let resp = client::raw_post(cfg, "/profiling/api/v1/aggregate", body)
+    let mut resp = client::raw_post(cfg, "/profiling/api/v1/aggregate", body)
         .await
         .map_err(|e| anyhow::anyhow!("failed to aggregate profiles: {e:?}"))?;
+    if let Some(name) = show_from {
+        apply_show_from(&mut resp, &name)?;
+        prune_aggregate_response(&mut resp);
+    }
     formatter::output(cfg, &resp)
 }
 
+// When --show-from has trimmed `flameGraph` to a small subtree, the rest of the
+// response (frames, strings, metadata) is mostly dead weight. Compact `frames`
+// and `strings` to just the entries the trimmed flame graph still references,
+// remap IDs accordingly, and drop heavy fields the UI uses but a CLI consumer
+// almost never wants alongside a single-function view.
+fn prune_aggregate_response(resp: &mut serde_json::Value) {
+    use std::collections::{HashMap, HashSet};
+
+    let Some(flame) = resp.get("flameGraph").cloned() else {
+        return;
+    };
+    let Some(frames) = resp.get("frames").and_then(|f| f.as_array()).cloned() else {
+        return;
+    };
+    let Some(strings) = resp.get("strings").and_then(|s| s.as_array()).cloned() else {
+        return;
+    };
+
+    let mut used_frames: HashSet<i64> = HashSet::new();
+    collect_used_frame_ids(&flame, &mut used_frames);
+    if used_frames.is_empty() {
+        return;
+    }
+
+    // Indices 2..=5 of a frame are string-table refs (library, package, function, file).
+    let mut used_strings: HashSet<i64> = HashSet::new();
+    for &fid in &used_frames {
+        let Some(f) = frames.get(fid as usize).and_then(|f| f.as_array()) else {
+            continue;
+        };
+        for idx in 2usize..=5 {
+            if let Some(sid) = f.get(idx).and_then(|v| v.as_i64()) {
+                used_strings.insert(sid);
+            }
+        }
+    }
+
+    let mut sorted_frames: Vec<i64> = used_frames.iter().copied().collect();
+    sorted_frames.sort_unstable();
+    let frame_remap: HashMap<i64, i64> = sorted_frames
+        .iter()
+        .enumerate()
+        .map(|(i, &old)| (old, i as i64))
+        .collect();
+
+    let mut sorted_strings: Vec<i64> = used_strings.iter().copied().collect();
+    sorted_strings.sort_unstable();
+    let string_remap: HashMap<i64, i64> = sorted_strings
+        .iter()
+        .enumerate()
+        .map(|(i, &old)| (old, i as i64))
+        .collect();
+
+    let new_frames: Vec<serde_json::Value> = sorted_frames
+        .iter()
+        .filter_map(|&old_fid| {
+            let mut f = frames.get(old_fid as usize).cloned()?;
+            if let Some(arr) = f.as_array_mut() {
+                for idx in 2usize..=5 {
+                    if let Some(slot) = arr.get_mut(idx) {
+                        if let Some(old_sid) = slot.as_i64() {
+                            if let Some(&new_sid) = string_remap.get(&old_sid) {
+                                *slot = json!(new_sid);
+                            }
+                        }
+                    }
+                }
+            }
+            Some(f)
+        })
+        .collect();
+
+    let new_strings: Vec<serde_json::Value> = sorted_strings
+        .iter()
+        .filter_map(|&old_sid| strings.get(old_sid as usize).cloned())
+        .collect();
+
+    let mut new_flame = flame;
+    remap_frame_ids(&mut new_flame, &frame_remap);
+
+    resp["flameGraph"] = new_flame;
+    resp["frames"] = json!(new_frames);
+    resp["strings"] = json!(new_strings);
+
+    // These are large UI-oriented fields with no relationship to the
+    // post-filter subtree; strip to keep terminal output usable.
+    if let Some(obj) = resp.as_object_mut() {
+        for k in [
+            "metadata",
+            "frameSchemas",
+            "endpointCounts",
+            "endpointValues",
+            "summaryTable",
+            "summaryValues",
+            "summaryDurations",
+            "availableAttributes",
+            "featureUpgrades",
+            "languageFrameCounts",
+            "profileIds",
+            "emptyStateReason",
+        ] {
+            obj.remove(k);
+        }
+    }
+}
+
+fn collect_used_frame_ids(node: &serde_json::Value, out: &mut std::collections::HashSet<i64>) {
+    let Some(arr) = node.as_array() else {
+        return;
+    };
+    if let Some(fid) = arr.first().and_then(|v| v.as_i64()) {
+        out.insert(fid);
+    }
+    if let Some(children) = arr.get(3).and_then(|c| c.as_array()) {
+        for c in children {
+            collect_used_frame_ids(c, out);
+        }
+    }
+}
+
+fn remap_frame_ids(node: &mut serde_json::Value, remap: &std::collections::HashMap<i64, i64>) {
+    let Some(arr) = node.as_array_mut() else {
+        return;
+    };
+    if let Some(slot) = arr.first_mut() {
+        if let Some(old) = slot.as_i64() {
+            if let Some(&new) = remap.get(&old) {
+                *slot = json!(new);
+            }
+        }
+    }
+    if let Some(children) = arr.get_mut(3).and_then(|c| c.as_array_mut()) {
+        for c in children {
+            remap_frame_ids(c, remap);
+        }
+    }
+}
+
+// Client-side equivalent of the UI's `show_from(<function>)` flame-graph filter.
+// Replaces `data.flameGraph` with a synthetic root whose children are the
+// topmost subtrees rooted at frames with the given function name (exact match).
+// Frame names are resolved via `data.frames[i][4]` (the `function` field of
+// `frameSchema`) → `data.strings[id]`. `data.frames` and `data.strings` are
+// left untouched so callers can still resolve frame metadata.
+fn apply_show_from(resp: &mut serde_json::Value, function_name: &str) -> Result<()> {
+    use std::collections::HashSet;
+
+    let strings = resp
+        .get("strings")
+        .and_then(|s| s.as_array())
+        .ok_or_else(|| anyhow::anyhow!("response has no 'strings' array"))?;
+    let target_string_ids: HashSet<i64> = strings
+        .iter()
+        .enumerate()
+        .filter_map(|(i, s)| {
+            s.as_str()
+                .filter(|st| *st == function_name)
+                .map(|_| i as i64)
+        })
+        .collect();
+    if target_string_ids.is_empty() {
+        anyhow::bail!(
+            "--show-from: no string in 'strings' matches function name {function_name:?}"
+        );
+    }
+
+    let frames_arr: Vec<serde_json::Value> = resp
+        .get("frames")
+        .and_then(|f| f.as_array())
+        .ok_or_else(|| anyhow::anyhow!("response has no 'frames' array"))?
+        .clone();
+    let target_frame_ids: HashSet<i64> = frames_arr
+        .iter()
+        .enumerate()
+        .filter_map(|(i, f)| {
+            let arr = f.as_array()?;
+            let fname_id = arr.get(4)?.as_i64()?;
+            if target_string_ids.contains(&fname_id) {
+                Some(i as i64)
+            } else {
+                None
+            }
+        })
+        .collect();
+    if target_frame_ids.is_empty() {
+        anyhow::bail!(
+            "--show-from: no frame in 'frames' references function name {function_name:?}"
+        );
+    }
+
+    let flame = resp
+        .get("flameGraph")
+        .cloned()
+        .ok_or_else(|| anyhow::anyhow!("response has no 'flameGraph'"))?;
+    let mut matches = Vec::new();
+    collect_show_from_subtrees(&flame, &target_frame_ids, &mut matches);
+    if matches.is_empty() {
+        anyhow::bail!(
+            "--show-from: no node in 'flameGraph' references function name {function_name:?}"
+        );
+    }
+
+    resp["flameGraph"] = merge_subtrees_by_function(matches, &frames_arr);
+    Ok(())
+}
+
+// Merge a set of flame-graph subtrees into a single subtree by collapsing
+// children that share the same function name (frames[fid][4] -> string id).
+// Datadog's flame graph emits a distinct `frame_id` per (function, file, line)
+// tuple, so a single logical function can appear under many frame IDs after
+// inlining/generics. The UI merges siblings by display name; we mirror that.
+fn merge_subtrees_by_function(
+    mut nodes: Vec<serde_json::Value>,
+    frames: &[serde_json::Value],
+) -> serde_json::Value {
+    if nodes.len() == 1 {
+        return nodes.pop().unwrap();
+    }
+    let representative_fid = nodes
+        .iter()
+        .filter_map(|n| n.as_array().and_then(|a| a.first()).and_then(|v| v.as_i64()))
+        .next()
+        .unwrap_or(0);
+    let total_value: i64 = nodes
+        .iter()
+        .filter_map(|n| n.as_array().and_then(|a| a.get(1)).and_then(|v| v.as_i64()))
+        .sum();
+    let mut all_children: Vec<serde_json::Value> = Vec::new();
+    for n in &nodes {
+        if let Some(children) = n.as_array().and_then(|a| a.get(3)).and_then(|c| c.as_array()) {
+            all_children.extend(children.iter().cloned());
+        }
+    }
+    // Group by function string id; preserve insertion order so output is stable.
+    use std::collections::BTreeMap;
+    let mut groups: BTreeMap<i64, Vec<serde_json::Value>> = BTreeMap::new();
+    let mut keyless: Vec<serde_json::Value> = Vec::new();
+    for c in all_children {
+        let key = c
+            .as_array()
+            .and_then(|a| a.first())
+            .and_then(|v| v.as_i64())
+            .and_then(|fid| frames.get(fid as usize))
+            .and_then(|f| f.as_array())
+            .and_then(|fa| fa.get(4))
+            .and_then(|v| v.as_i64());
+        match key {
+            Some(k) => groups.entry(k).or_default().push(c),
+            None => keyless.push(c),
+        }
+    }
+    let mut merged_children: Vec<serde_json::Value> = groups
+        .into_values()
+        .map(|grp| merge_subtrees_by_function(grp, frames))
+        .collect();
+    merged_children.append(&mut keyless);
+    json!([representative_fid, total_value, -1.0, merged_children])
+}
+
+fn collect_show_from_subtrees(
+    node: &serde_json::Value,
+    targets: &std::collections::HashSet<i64>,
+    out: &mut Vec<serde_json::Value>,
+) {
+    let Some(arr) = node.as_array() else {
+        return;
+    };
+    if let Some(fid) = arr.first().and_then(|v| v.as_i64()) {
+        if targets.contains(&fid) {
+            out.push(node.clone());
+            return; // topmost match only — don't descend into nested re-entries
+        }
+    }
+    if let Some(children) = arr.get(3).and_then(|c| c.as_array()) {
+        for c in children {
+            collect_show_from_subtrees(c, targets, out);
+        }
+    }
+}
+
 pub async fn analysis(cfg: &Config, profile_id: &str, event_id: Option<String>) -> Result<()> {
     let path = format!("/profiling/api/v1/profiles/{profile_id}/analysis");
     let query: Vec<(&str, &str)> = match event_id.as_deref() {

src/main.rs

@@ -8868,6 +8868,11 @@ enum ProfilingActions {
         limit: u32,
         #[arg(long, default_value = "sum", help = "Aggregation function: sum or avg")]
         aggregation_function: String,
+        #[arg(
+            long,
+            help = "Filter the flame graph to subtrees rooted at frames whose function name matches exactly (UI equivalent: show_from(<name>))"
+        )]
+        show_from: Option<String>,
     },
     /// Get automated analysis for a profile
     Analysis {
@@ -13984,6 +13989,7 @@ async fn main_inner() -> anyhow::Result<()> {
                     to,
                     limit,
                     aggregation_function,
+                    show_from,
                 } => {
                     commands::profiling::aggregate(
                         &cfg,
@@ -13993,6 +13999,7 @@ async fn main_inner() -> anyhow::Result<()> {
                         to,
                         limit,
                         aggregation_function,
+                        show_from,
                     )
                     .await?;
                 }