Implement IR visualizer

bin/kecc.rs

@@ -1,13 +1,17 @@
 use clap::Parser;
 
 use std::ffi::OsStr;
+use std::io::Write;
+use std::os::unix::io::{FromRawFd, IntoRawFd};
 use std::path::Path;
+use std::process::{Command, Stdio};
 
 use lang_c::ast::TranslationUnit;
+use tempfile::tempdir;
 
 use kecc::{
-    ir, ok_or_exit, write, Asmgen, Deadcode, Gvn, IrParse, Irgen, Mem2reg, Optimize, Parse,
-    SimplifyCfg, Translate, O1,
+    ir, ok_or_exit, write, Asmgen, Deadcode, Gvn, IrParse, IrVisualizer, Irgen, Mem2reg, Optimize,
+    Parse, SimplifyCfg, Translate, O1,
 };
 
 #[derive(Debug, Parser)]
@@ -37,6 +41,10 @@ struct KeccCli {
     #[clap(long)]
     irrun: bool,
 
+    /// Visualizes IR
+    #[clap(long, value_name = "FILE")]
+    irvisualize: Option<String>,
+
     /// Optimizes IR
     #[clap(short = 'O', long)]
     optimize: bool,
@@ -133,6 +141,46 @@ fn compile_ir(
         return;
     }
 
+    if let Some(path) = &matches.irvisualize {
+        assert_eq!(
+            Path::new(&path).extension(),
+            Some(std::ffi::OsStr::new("png"))
+        );
+        let img_path = Path::new(path);
+        let dot = IrVisualizer::default()
+            .translate(input)
+            .expect("ir visualize failed");
+
+        let temp_dir = tempdir().expect("temp dir creation failed");
+        let dot_path = temp_dir.path().join("temp.dot");
+        let dot_path_str = dot_path.as_path().display().to_string();
+        let img_path_str = img_path.display().to_string();
+
+        // Create the dot file
+        let mut buffer =
+            ::std::fs::File::create(dot_path.as_path()).expect("need to success creating file");
+        buffer
+            .write(dot.as_bytes())
+            .expect("failed to write to dot file");
+
+        // Create the image file
+        let img = ::std::fs::File::create(&img_path_str).expect("need to success creating file");
+
+        // Translate dot file into image
+        if !Command::new("dot")
+            .args(&["-Tpng", &dot_path_str])
+            .stdout(unsafe { Stdio::from_raw_fd(img.into_raw_fd()) })
+            .status()
+            .unwrap()
+            .success()
+        {
+            panic!("failed to save image file");
+        }
+
+        drop(buffer);
+        temp_dir.close().expect("temp dir deletion failed");
+    }
+
     if matches.optimize {
         O1::default().optimize(input);
     } else {

src/ir/mod.rs

@@ -5,6 +5,7 @@ mod equiv;
 mod interp;
 #[allow(clippy::all)]
 mod parse;
+mod visualize;
 mod write_ir;
 
 use core::convert::TryFrom;
@@ -19,6 +20,7 @@ use std::hash::{Hash, Hasher};
 pub use dtype::{Dtype, DtypeError, HasDtype};
 pub use interp::{interp, Value};
 pub use parse::Parse;
+pub use visualize::Visualizer;
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct TranslationUnit {

src/ir/visualize.rs

@@ -0,0 +1,229 @@
+//! Visualize IR.
+
+use std::collections::HashMap;
+
+use crate::ir::*;
+use crate::some_or;
+use crate::write_base::*;
+use crate::Translate;
+
+#[derive(Default, Debug)]
+pub struct Visualizer {
+    /// First instruction in the function.
+    function_first_instruction: HashMap<String, String>,
+
+    /// First instruction in the block.
+    block_first_instruction: HashMap<(String, BlockId), String>,
+}
+
+impl Translate<TranslationUnit> for Visualizer {
+    type Target = String;
+    type Error = ();
+
+    fn translate(&mut self, source: &TranslationUnit) -> Result<Self::Target, Self::Error> {
+        let mut subgraphs = Vec::new();
+
+        // TODO: Add variables and structs information
+        for (name, decl) in &source.decls {
+            match decl {
+                Declaration::Variable { .. } => {}
+                Declaration::Function {
+                    signature,
+                    definition,
+                } => {
+                    let definition = some_or!(definition, continue);
+                    let subgraph = self.translate_function(name, signature, definition)?;
+                    subgraphs.push(subgraph);
+                }
+            }
+        }
+
+        let mut edges = Vec::new();
+
+        // Add edges between subgraphs
+        for (name, decl) in &source.decls {
+            if let Declaration::Function { definition, .. } = decl {
+                let definition = some_or!(definition, continue);
+
+                for (bid, block) in &definition.blocks {
+                    for (iid, instruction) in block.instructions.iter().enumerate() {
+                        if let Instruction::Call { callee, .. } = &instruction.inner {
+                            let callee = self.translate_callee(callee)?;
+                            let from = self.translate_instruction_node(name, *bid, iid);
+                            let to = self.get_function_first_instruction(&callee);
+
+                            edges.push(format!("{} -> {};", from, to));
+                        }
+                    }
+                }
+            }
+        }
+
+        let inner = vec![subgraphs, edges].concat().join("\n");
+
+        Ok(format!("digraph G {{\n{}\n}}", inner))
+    }
+}
+
+impl Visualizer {
+    #[inline]
+    fn get_function_first_instruction(&self, name: &str) -> String {
+        self.function_first_instruction
+            .get(name)
+            .expect("failed to get first instruction from function")
+            .clone()
+    }
+
+    #[inline]
+    fn get_block_first_instruction(&self, name: &str, bid: BlockId) -> String {
+        self.block_first_instruction
+            .get(&(name.to_string(), bid))
+            .expect("failed to get first instruction from block")
+            .clone()
+    }
+
+    #[inline]
+    fn translate_instruction_node(&self, name: &str, bid: BlockId, iid: usize) -> String {
+        format!("\"{}:{}:i{}\"", name, bid, iid)
+    }
+
+    #[inline]
+    fn translate_block_exit_node(&self, name: &str, bid: BlockId) -> String {
+        format!("\"{}:{}:exit\"", name, bid)
+    }
+
+    #[inline]
+    fn translate_callee(&mut self, callee: &Operand) -> Result<String, ()> {
+        match callee {
+            Operand::Constant(_constant @ Constant::GlobalVariable { name, .. }) => {
+                Ok(name.clone())
+            }
+            _ => todo!("translate_callee: operand {:?}", callee),
+        }
+    }
+
+    fn translate_function(
+        &mut self,
+        name: &str,
+        signature: &FunctionSignature,
+        definition: &FunctionDefinition,
+    ) -> Result<String, ()> {
+        let mut subgraphs = Vec::new();
+
+        for (bid, block) in &definition.blocks {
+            let subgraph = self.translate_block(name, bid, block)?;
+            subgraphs.push(subgraph);
+        }
+
+        let _ = self.function_first_instruction.insert(
+            name.to_string(),
+            self.get_block_first_instruction(name, definition.bid_init),
+        );
+
+        let params = signature
+            .params
+            .iter()
+            .map(|p| p.to_string())
+            .collect::<Vec<_>>()
+            .join(", ");
+        let label = format!("label=\"fun {} @{} ({})\";", signature.ret, name, params);
+
+        let mut edges = Vec::new();
+
+        // Add edges for block exit
+        for (bid, block) in &definition.blocks {
+            match &block.exit {
+                BlockExit::Jump { arg } => {
+                    edges.push(format!(
+                        "{} -> {};",
+                        self.translate_block_exit_node(name, *bid),
+                        self.get_block_first_instruction(name, arg.bid)
+                    ));
+                }
+                BlockExit::ConditionalJump {
+                    arg_then, arg_else, ..
+                } => {
+                    edges.push(format!(
+                        "{} -> {} [label=\"true\"];",
+                        self.translate_block_exit_node(name, *bid),
+                        self.get_block_first_instruction(name, arg_then.bid)
+                    ));
+                    edges.push(format!(
+                        "{} -> {} [label=\"false\"];",
+                        self.translate_block_exit_node(name, *bid),
+                        self.get_block_first_instruction(name, arg_else.bid)
+                    ));
+                }
+                BlockExit::Switch { default, cases, .. } => {
+                    edges.push(format!(
+                        "{} -> {} [label=\"default\"];",
+                        self.translate_block_exit_node(name, *bid),
+                        self.get_block_first_instruction(name, default.bid)
+                    ));
+                    for (constant, arg) in cases {
+                        edges.push(format!(
+                            "{} -> {} [label=\"{}\"];",
+                            self.translate_block_exit_node(name, *bid),
+                            self.get_block_first_instruction(name, arg.bid),
+                            constant,
+                        ));
+                    }
+                }
+                _ => {}
+            }
+        }
+
+        // TODO: Add init information (bid_init, allocations)
+        let inner = vec![subgraphs, vec![label], edges].concat().join("\n");
+
+        Ok(format!("subgraph \"cluster.{}\" {{\n{}\n}}", name, inner))
+    }
+
+    fn translate_block(&mut self, name: &str, bid: &BlockId, block: &Block) -> Result<String, ()> {
+        let mut header = Vec::new();
+        header.push("style=filled;".to_string());
+        header.push("color=lightgrey;".to_string());
+        header.push("node [shape=record];".to_string());
+        header.push(format!("label=\"{}\";", bid));
+
+        let mut nodes = Vec::new();
+
+        // Add nodes for instructions
+        for (iid, instruction) in block.instructions.iter().enumerate() {
+            nodes.push(format!(
+                "{} [label=\"{}\"]",
+                self.translate_instruction_node(name, *bid, iid),
+                instruction.write_string()
+            ));
+        }
+
+        // Add node for block exit
+        nodes.push(format!(
+            "{} [label=\"{}\"]",
+            self.translate_block_exit_node(name, *bid),
+            block.exit.write_string()
+        ));
+
+        let edges = (0..block.instructions.len())
+            .map(|iid| self.translate_instruction_node(name, *bid, iid))
+            .chain([self.translate_block_exit_node(name, *bid)].into_iter())
+            .collect::<Vec<String>>()
+            .join(" -> ");
+
+        let first_instruction = if block.instructions.is_empty() {
+            self.translate_block_exit_node(name, *bid)
+        } else {
+            self.translate_instruction_node(name, *bid, 0)
+        };
+        let _ = self
+            .block_first_instruction
+            .insert((name.to_string(), *bid), first_instruction);
+
+        let inner = vec![header, nodes, vec![edges]].concat().join("\n");
+
+        Ok(format!(
+            "subgraph \"cluster.{}.{}\" {{\n{}\n}}",
+            name, bid, inner
+        ))
+    }
+}

src/lib.rs

@@ -8,12 +8,12 @@
 #![deny(deprecated_in_future)]
 #![deny(elided_lifetimes_in_paths)]
 #![deny(explicit_outlives_requirements)]
-#![deny(invalid_html_tags)]
+#![deny(rustdoc::invalid_html_tags)]
 #![deny(keyword_idents)]
 #![deny(macro_use_extern_crate)]
 #![deny(missing_debug_implementations)]
 // #![deny(missing_docs)] TODO
-#![deny(missing_doc_code_examples)]
+#![deny(rustdoc::missing_doc_code_examples)]
 #![deny(non_ascii_idents)]
 #![deny(pointer_structural_match)]
 // #![deny(single_use_lifetimes)]
@@ -58,6 +58,7 @@ pub use write_base::write;
 
 pub use c::Parse;
 pub use ir::Parse as IrParse;
+pub use ir::Visualizer as IrVisualizer;
 
 pub use asmgen::Asmgen;
 pub use irgen::Irgen;