up(semantics): add method to print dot graphs

crates/oxc_semantic/examples/cfg.rs

@@ -1,11 +1,10 @@
-use std::{collections::HashMap, env, path::Path, sync::Arc};
+use std::{collections::HashMap, env, fs, io, path::Path, sync::Arc};
 
 use itertools::Itertools;
 use oxc_allocator::Allocator;
 use oxc_parser::Parser;
-use oxc_semantic::{print_basic_block, SemanticBuilder};
+use oxc_semantic::SemanticBuilder;
 use oxc_span::SourceType;
-use petgraph::dot::{Config, Dot};
 
 // Instruction:
 // 1. create a `test.js`,
@@ -16,7 +15,7 @@ use petgraph::dot::{Config, Dot};
 //    - CFG blocks (test.cfg.txt)
 //    - CFG graph (test.dot)
 
-fn main() -> std::io::Result<()> {
+fn main() -> io::Result<()> {
     let test_file_name = env::args().nth(1).unwrap_or_else(|| "test.js".to_string());
     let ast_file_name = env::args().nth(1).unwrap_or_else(|| "test.ast.txt".to_string());
     let cfg_file_name = env::args().nth(1).unwrap_or_else(|| "test.cfg.txt".to_string());
@@ -27,13 +26,13 @@ fn main() -> std::io::Result<()> {
     let cfg_file_path = Path::new(&cfg_file_name);
     let dot_file_path = Path::new(&dot_file_name);
 
-    let source_text = Arc::new(std::fs::read_to_string(test_file_path)?);
+    let source_text = Arc::new(fs::read_to_string(test_file_path)?);
     let allocator = Allocator::default();
     let source_type = SourceType::from_path(test_file_path).unwrap();
     let ret = Parser::new(&allocator, &source_text, source_type).parse();
 
     let program = allocator.alloc(ret.program);
-    std::fs::write(ast_file_path, format!("{:#?}", &program))?;
+    fs::write(ast_file_path, format!("{:#?}", &program))?;
     println!("Wrote AST to: {}", &ast_file_name);
 
     let semantic = SemanticBuilder::new(&source_text, source_type)
@@ -64,12 +63,11 @@ fn main() -> std::io::Result<()> {
         .cfg()
         .basic_blocks
         .iter()
-        .map(print_basic_block)
         .enumerate()
-        .map(|(i, it)| {
+        .map(|(i, basic_block)| {
             format!(
                 "bb{i}: {{\n{}\n---\n{}\n}}",
-                it.lines().map(|x| format!("\t{}", x.trim())).join("\n"),
+                basic_block.iter().map(|el| format!("\t{el}")).join("\n"),
                 ast_nodes_by_block
                     .get(&i)
                     .map(|nodes| {
@@ -80,32 +78,11 @@ fn main() -> std::io::Result<()> {
         })
         .join("\n\n");
 
-    std::fs::write(cfg_file_path, basic_blocks_printed)?;
+    fs::write(cfg_file_path, basic_blocks_printed)?;
     println!("Wrote CFG blocks to: {}", &cfg_file_name);
 
-    let cfg_dot_diagram = format!(
-        "{:?}",
-        Dot::with_attr_getters(
-            &semantic.semantic.cfg().graph,
-            &[Config::EdgeNoLabel, Config::NodeNoLabel],
-            &|_graph, edge| format!("label = {:?}", edge.weight()),
-            &|_graph, node| format!(
-                "xlabel = {:?}, label = {:?}",
-                format!(
-                    "bb{} [{}]",
-                    node.1,
-                    print_basic_block(&semantic.semantic.cfg().basic_blocks[*node.1],).trim()
-                ),
-                ast_nodes_by_block
-                    .get(node.1)
-                    .map(|nodes| {
-                        nodes.iter().map(|node| format!("{}", node.kind().debug_name())).join("\n")
-                    })
-                    .unwrap_or_default()
-            )
-        )
-    );
-    std::fs::write(dot_file_path, cfg_dot_diagram)?;
+    let cfg_dot_diagram = format!("{:?}", semantic.semantic.cfg().dot());
+    fs::write(dot_file_path, cfg_dot_diagram)?;
     println!("Wrote CFG dot diagram to: {}", &dot_file_name);
 
     Ok(())

crates/oxc_semantic/src/control_flow/dot.rs

@@ -0,0 +1,65 @@
+use petgraph::{dot, stable_graph::NodeIndex};
+
+use crate::ControlFlowGraph;
+
+use std::fmt::{self, Debug, Display};
+
+impl ControlFlowGraph {
+    /// Returns an object that implements [`Display`] for printing a
+    /// [`GraphViz DOT`] representation of the control flow graph.
+    ///
+    /// [`GraphViz DOT`]: https://graphviz.org/doc/info/lang.html
+    pub fn dot(&self) -> Dot<'_> {
+        // Exposing our own struct instead of petgraph's allows us to control
+        // our API surface.
+        Dot { cfg: self }
+    }
+
+    pub(crate) fn fmt_dot(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        const CONFIG: &[dot::Config] = &[dot::Config::EdgeNoLabel, dot::Config::NodeNoLabel];
+
+        let get_node_attributes = |_graph, node: (NodeIndex, &usize)| {
+            let block = &self.basic_blocks[*node.1];
+
+            format!(
+                "label = \"{}\"",
+                block.iter().map(|el| format!("{el}")).collect::<Vec<_>>().join("\\n")
+            )
+        };
+
+        let dot = dot::Dot::with_attr_getters(
+            &self.graph,
+            CONFIG,
+            &|_graph, _edge| String::new(),
+            // todo: We currently do not print edge types into cfg dot diagram
+            // so they aren't snapshotted, but we could by uncommenting this.
+            // &|_graph, edge| format!("label = {:?}", edge.weight()),
+            // &self.node_attributes,
+            &get_node_attributes,
+        );
+
+        dot.fmt(f)
+    }
+}
+
+/// Helper struct for rendering [`DOT`] diagrams of a [`ControlFlowGraph`].
+/// Returned from [`ControlFlowGraph::dot`].
+///
+/// [`DOT`]: https://graphviz.org/doc/info/lang.html
+///
+#[derive(Clone, Copy)]
+pub struct Dot<'a> {
+    cfg: &'a ControlFlowGraph,
+}
+
+impl<'a> Debug for Dot<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        self.cfg.fmt_dot(f)
+    }
+}
+
+impl<'a> Display for Dot<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        self.cfg.fmt_dot(f)
+    }
+}

crates/oxc_semantic/src/control_flow/mod.rs

@@ -1,4 +1,7 @@
 mod builder;
+mod dot;
+
+use core::fmt;
 
 use oxc_span::CompactStr;
 use oxc_syntax::operator::{
@@ -156,43 +159,30 @@ pub struct PreservedExpressionState {
     pub store_final_assignments_into_this_array: Vec<Vec<Register>>,
 }
 
-#[must_use]
-fn print_register(register: Register) -> String {
-    match &register {
-        Register::Index(i) => format!("${i}"),
-        Register::Return => "$return".into(),
+impl fmt::Display for Register {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Index(i) => write!(f, "${i}"),
+            Self::Return => write!(f, "$return"),
+        }
     }
 }
 
-#[must_use]
-pub fn print_basic_block(basic_block_elements: &Vec<BasicBlockElement>) -> String {
-    let mut output = String::new();
-    for basic_block in basic_block_elements {
-        match basic_block {
-            BasicBlockElement::Unreachable => output.push_str("Unreachable()\n"),
-            BasicBlockElement::Throw(reg) => {
-                output.push_str(&format!("throw {}\n", print_register(*reg)));
-            }
-
-            BasicBlockElement::Break(Some(reg)) => {
-                output.push_str(&format!("break {}\n", print_register(*reg)));
-            }
-            BasicBlockElement::Break(None) => {
-                output.push_str("break");
-            }
-            BasicBlockElement::Assignment(to, with) => {
-                output.push_str(&format!("{} = ", print_register(*to)));
-
-                match with {
-                    AssignmentValue::ImplicitUndefined => {
-                        output.push_str("<implicit undefined>");
-                    }
-                    AssignmentValue::NotImplicitUndefined => output.push_str("<value>"),
-                }
-
-                output.push('\n');
+impl fmt::Display for BasicBlockElement {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::Unreachable => f.write_str("Unreachable()"),
+            Self::Throw(reg) => write!(f, "throw {reg}"),
+            Self::Break(Some(reg)) => write!(f, "break {reg}"),
+            Self::Break(None) => f.write_str("break"),
+            Self::Assignment(to, with) => {
+                // output.push_str(&format!("{} = ", print_register(*to)));
+                let value = match with {
+                    AssignmentValue::ImplicitUndefined => "<implicit undefined>",
+                    AssignmentValue::NotImplicitUndefined => "<value>",
+                };
+                write!(f, "{to} = {value}")
             }
         }
     }
-    output
 }

crates/oxc_semantic/src/lib.rs

@@ -31,10 +31,10 @@ use rustc_hash::FxHashSet;
 
 pub use crate::{
     control_flow::{
-        print_basic_block, AssignmentValue, BasicBlockElement, BasicBlockId, BinaryAssignmentValue,
-        BinaryOp, CallType, CalleeWithArgumentsAssignmentValue, CollectionAssignmentValue,
-        ControlFlowGraph, EdgeType, ObjectPropertyAccessAssignmentValue, Register,
-        UnaryExpressioneAssignmentValue, UpdateAssignmentValue,
+        AssignmentValue, BasicBlockElement, BasicBlockId, BinaryAssignmentValue, BinaryOp,
+        CallType, CalleeWithArgumentsAssignmentValue, CollectionAssignmentValue, ControlFlowGraph,
+        EdgeType, ObjectPropertyAccessAssignmentValue, Register, UnaryExpressioneAssignmentValue,
+        UpdateAssignmentValue,
     },
     node::{AstNode, AstNodeId, AstNodes},
     reference::{Reference, ReferenceFlag, ReferenceId},

crates/oxc_semantic/tests/integration/util/mod.rs

@@ -6,12 +6,11 @@ use std::{path::PathBuf, sync::Arc};
 use itertools::Itertools;
 use oxc_allocator::Allocator;
 use oxc_diagnostics::{Error, NamedSource, OxcDiagnostic};
-use oxc_semantic::{print_basic_block, Semantic, SemanticBuilder};
+use oxc_semantic::{Semantic, SemanticBuilder};
 use oxc_span::SourceType;
 
 pub use class_tester::ClassTester;
 pub use expect::Expect;
-use petgraph::dot::{Config, Dot};
 pub use symbol_tester::SymbolTester;
 
 pub struct SemanticTester<'a> {
@@ -111,34 +110,19 @@ impl<'a> SemanticTester<'a> {
             .cfg()
             .basic_blocks
             .iter()
-            .map(print_basic_block)
             .enumerate()
-            .map(|(i, it)| {
+            .map(|(i, basic_block)| {
                 format!(
                     "bb{i}: {{\n{}\n}}",
-                    it.lines().map(|x| format!("\t{}", x.trim())).join("\n")
+                    basic_block.iter().map(|block_el| format!("\t{block_el}")).join("\n")
                 )
             })
             .join("\n\n")
     }
 
     pub fn cfg_dot_diagram(&self) -> String {
         let built = self.build();
-        format!(
-            "{:?}",
-            Dot::with_attr_getters(
-                &built.cfg().graph,
-                &[Config::EdgeNoLabel, Config::NodeNoLabel],
-                &|_graph, _edge| String::new(),
-                // todo: We currently do not print edge types into cfg dot diagram
-                // so they aren't snapshotted, but we could by uncommenting this.
-                // &|_graph, edge| format!("label = {:?}", edge.weight()),
-                &|_graph, node| format!(
-                    "label = {:?}",
-                    print_basic_block(&built.cfg().basic_blocks[*node.1],).trim()
-                )
-            )
-        )
+        format!("{:?}", built.cfg().dot())
     }
 
     /// Tests that a symbol with the given name exists at the top-level scope and provides a