Enhancements for xdot

hercules_ir/src/dot.rs

-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::env::temp_dir;
 use std::fmt::Write;
 use std::fs::File;
@@ -15,6 +15,7 @@ use crate::*;
  */
 pub fn xdot_module(
     module: &ir::Module,
+    func_selection: &Vec<bool>,
     reverse_postorders: &Vec<Vec<NodeID>>,
     typing: Option<&ModuleTyping>,
     doms: Option<&Vec<DomTree>>,
@@ -30,6 +31,7 @@ pub fn xdot_module(
     let mut contents = String::new();
     write_dot(
         &module,
+        func_selection,
         &reverse_postorders,
         typing,
         doms,
@@ -54,6 +56,7 @@ pub fn xdot_module(
  */
 pub fn write_dot<W: Write>(
     module: &ir::Module,
+    func_selection: &Vec<bool>,
     reverse_postorders: &Vec<Vec<NodeID>>,
     typing: Option<&ModuleTyping>,
     doms: Option<&Vec<DomTree>>,
@@ -64,7 +67,10 @@ pub fn write_dot<W: Write>(
 ) -> std::fmt::Result {
     write_digraph_header(w)?;
 
-    for function_id in (0..module.functions.len()).map(FunctionID::new) {
+    for function_id in (0..module.functions.len())
+        .filter(|idx| func_selection[*idx])
+        .map(FunctionID::new)
+    {
         let function = &module.functions[function_id.idx()];
         let reverse_postorder = &reverse_postorders[function_id.idx()];
         let mut reverse_postorder_node_numbers = vec![0; function.nodes.len()];
@@ -93,7 +99,9 @@ pub fn write_dot<W: Write>(
                 color,
                 module,
                 typing,
+                &function.labels[node_id.idx()],
                 &function.schedules[node_id.idx()],
+                &module.labels,
                 w,
             )?;
         }
@@ -254,7 +262,9 @@ fn write_node<W: Write>(
     color: &str,
     module: &Module,
     typing: Option<&ModuleTyping>,
+    label_set: &HashSet<LabelID>,
     schedules: &Vec<Schedule>,
+    labels: &Vec<String>,
     w: &mut W,
 ) -> std::fmt::Result {
     let node = &module.functions[function_id.idx()].nodes[node_id.idx()];
@@ -341,62 +351,39 @@ fn write_node<W: Write>(
     if let Some(ty) = typing.map(|typing| typing[function_id.idx()][node_id.idx()]) {
         module.write_type(ty, &mut tylabel)?;
     }
-
+    let mut iter = label_set.into_iter();
+    let label_set = if let Some(first) = iter.next() {
+        iter.fold(format!("{}", labels[first.idx()]), |b, i| {
+            format!("{}, {}", b, labels[i.idx()])
+        })
+    } else {
+        String::new()
+    };
     let mut iter = schedules.into_iter();
     let schedules = if let Some(first) = iter.next() {
         iter.fold(format!("{:?}", first), |b, i| format!("{}, {:?}", b, i))
     } else {
         String::new()
     };
-    if tylabel.is_empty() && schedules.is_empty() {
-        write!(
-            w,
-            "{}_{}_{} [xlabel={}, label=<{}>, color={}];\n",
-            node.lower_case_name(),
-            function_id.idx(),
-            node_id.idx(),
-            xlabel,
-            label,
-            color
-        )?;
-    } else if schedules.is_empty() {
-        write!(
-            w,
-            "{}_{}_{} [xlabel={}, label=<{}<BR /><FONT POINT-SIZE=\"8\">{}</FONT>>, color={}];\n",
-            node.lower_case_name(),
-            function_id.idx(),
-            node_id.idx(),
-            xlabel,
-            label,
-            tylabel,
-            color
-        )?;
-    } else if tylabel.is_empty() {
-        write!(
-            w,
-            "{}_{}_{} [xlabel={}, label=<{}<BR /><FONT POINT-SIZE=\"8\">{}</FONT>>, color={}];\n",
-            node.lower_case_name(),
-            function_id.idx(),
-            node_id.idx(),
-            xlabel,
-            label,
-            schedules,
-            color
-        )?;
-    } else {
-        write!(
-            w,
-            "{}_{}_{} [xlabel={}, label=<{}<BR /><FONT POINT-SIZE=\"8\">{}</FONT><BR /><FONT POINT-SIZE=\"8\">{}</FONT>>, color={}];\n",
-            node.lower_case_name(),
-            function_id.idx(),
-            node_id.idx(),
-            xlabel,
-            label,
-            tylabel,
-            schedules,
-            color
-        )?;
+    write!(
+        w,
+        "{}_{}_{} [xlabel={}, label=<{}",
+        node.lower_case_name(),
+        function_id.idx(),
+        node_id.idx(),
+        xlabel,
+        label,
+    )?;
+    if !tylabel.is_empty() {
+        write!(w, "<BR /><FONT POINT-SIZE=\"8\">{}</FONT>", tylabel,)?;
+    }
+    if !label_set.is_empty() {
+        write!(w, "<BR /><FONT POINT-SIZE=\"8\">{}</FONT>", label_set,)?;
+    }
+    if !schedules.is_empty() {
+        write!(w, "<BR /><FONT POINT-SIZE=\"8\">{}</FONT>", schedules,)?;
     }
+    write!(w, ">, color={}];\n", color)?;
     Ok(())
 }
 

juno_scheduler/src/pm.rs

@@ -2328,6 +2328,22 @@ fn run_pass(
                 None => true,
             };
 
+            let mut bool_selection = vec![];
+            if let Some(selection) = selection {
+                bool_selection = vec![false; pm.functions.len()];
+                for loc in selection {
+                    let CodeLocation::Function(id) = loc else {
+                        return Err(SchedulerError::PassError {
+                        pass: "xdot".to_string(),
+                        error: "expected coarse-grained selection (can't partially xdot a function)".to_string(),
+                    });
+                    };
+                    bool_selection[id.idx()] = true;
+                }
+            } else {
+                bool_selection = vec![true; pm.functions.len()];
+            }
+
             pm.make_reverse_postorders();
             if force_analyses {
                 pm.make_typing();
@@ -2345,6 +2361,7 @@ fn run_pass(
             pm.with_mod(|module| {
                 xdot_module(
                     module,
+                    &bool_selection,
                     &reverse_postorders,
                     typing.as_ref(),
                     doms.as_ref(),