Remove unnecessary stack frame suffix (#24)

src/lib.rs

@@ -23,6 +23,29 @@ pub struct ParseConfig {
     pub custom_parsers: Vec<Box<dyn crate::parsers::StructuredLogParser>>,
 }
 
+fn maybe_remove_suffix(frames: &mut Vec<FrameSummary>) {
+    let target_frames = [
+        ("torch/_dynamo/convert_frame.py", "catch_errors"),
+        ("torch/_dynamo/convert_frame.py", "_convert_frame"),
+        ("torch/_dynamo/convert_frame.py", "_convert_frame_assert"),
+    ];
+
+    let len = frames.len();
+    if len >= target_frames.len() {
+        let suffix = &frames[len - target_frames.len()..];
+        if suffix
+            .iter()
+            .zip(target_frames.iter())
+            .all(|(frame, target)| {
+                simplify_filename(unintern_str(frame.filename).as_ref()) == target.0
+                    && frame.name == target.1
+            })
+        {
+            frames.truncate(len - 3);
+        }
+    }
+}
+
 pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOutput> {
     let strict = config.strict;
     let file = File::open(path)?;
@@ -263,7 +286,8 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
         }
 
         if let Some(m) = e.dynamo_start {
-            if let Some(stack) = m.stack {
+            if let Some(mut stack) = m.stack {
+                maybe_remove_suffix(&mut stack);
                 stack_trie.insert(
                     stack,
                     e.compile_id.map_or("(unknown) ".to_string(), |c| {

src/types.rs

@@ -108,12 +108,12 @@ pub struct Stats {
 
 #[derive(Debug, Hash, Eq, PartialEq, Deserialize, Serialize)]
 pub struct FrameSummary {
-    filename: u32,
-    line: i32,
-    name: String,
+    pub filename: u32,
+    pub line: i32,
+    pub name: String,
 }
 
-fn simplify_filename<'a>(filename: &'a str) -> &'a str {
+pub fn simplify_filename<'a>(filename: &'a str) -> &'a str {
     let parts: Vec<&'a str> = filename.split("#link-tree/").collect();
     if parts.len() > 1 {
         return parts[1];
@@ -127,6 +127,14 @@ fn simplify_filename<'a>(filename: &'a str) -> &'a str {
     return filename;
 }
 
+pub fn unintern_str(interned_str: u32) -> String {
+    let intern_table = INTERN_TABLE.lock().unwrap();
+    let filename = intern_table
+        .get(&interned_str)
+        .map_or("(unknown)", |s| s.as_str());
+    return filename.to_string();
+}
+
 impl fmt::Display for FrameSummary {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let intern_table = INTERN_TABLE.lock().unwrap();