feat(ast): Add graph output

2026-05-09 12:29:59 +08:00
parent 567057fd76
commit 3c728fb2b8
7 changed files with 330 additions and 3 deletions
@@ -0,0 +1,155 @@
+use petgraph::dot::{Config, Dot};
+use petgraph::graph::{Graph, NodeIndex};
+
+use crate::ast::types::{
+    BlockStmt, CompileUnit, Expr, ExprValue, FuncDeclStmt, GlobalDeclStmt, Param, ReturnStmt,
+    Statement, VarDeclStmt, VarDeclStmtValue,
+};
+
+pub type AstGraph = Graph<String, String>;
+
+pub trait AstGraphExt {
+    fn to_graph(&self) -> AstGraph;
+    fn to_dot(&self) -> String {
+        format!("{}", Dot::with_config(&self.to_graph(), &[Config::EdgeNoLabel]))
+    }
+}
+
+impl AstGraphExt for CompileUnit {
+    fn to_graph(&self) -> AstGraph {
+        let mut builder = AstGraphBuilder::new();
+        builder.add_compile_unit(self);
+        builder.graph
+    }
+}
+
+struct AstGraphBuilder {
+    graph: AstGraph,
+}
+
+impl AstGraphBuilder {
+    fn new() -> Self {
+        Self { graph: Graph::new() }
+    }
+
+    fn node(&mut self, label: impl Into<String>) -> NodeIndex {
+        self.graph.add_node(label.into())
+    }
+
+    fn child(&mut self, parent: NodeIndex, label: impl Into<String>) -> NodeIndex {
+        let child = self.node(label);
+        self.graph.add_edge(parent, child, String::new());
+        child
+    }
+
+    fn add_compile_unit(&mut self, compile_unit: &CompileUnit) -> NodeIndex {
+        let root = self.node(compile_unit.to_string());
+        for decl in &compile_unit.global_decls {
+            self.add_global_decl(root, decl);
+        }
+        root
+    }
+
+    fn add_global_decl(&mut self, parent: NodeIndex, decl: &GlobalDeclStmt) -> NodeIndex {
+        match decl {
+            GlobalDeclStmt::VarDecl(var_decl) => {
+                let node = self.child(parent, decl.to_string());
+                self.add_var_decl(node, var_decl);
+                node
+            }
+            GlobalDeclStmt::FuncDecl(func_decl) => {
+                let node = self.child(parent, decl.to_string());
+                self.add_func_decl(node, func_decl);
+                node
+            }
+        }
+    }
+
+    fn add_var_decl(&mut self, parent: NodeIndex, var_decl: &VarDeclStmt) -> NodeIndex {
+        let node = self.child(parent, var_decl.to_string());
+        for value in &var_decl.values {
+            self.add_var_decl_value(node, value);
+        }
+        node
+    }
+
+    fn add_var_decl_value(&mut self, parent: NodeIndex, value: &VarDeclStmtValue) -> NodeIndex {
+        self.child(parent, value.to_string())
+    }
+
+    fn add_func_decl(&mut self, parent: NodeIndex, func_decl: &FuncDeclStmt) -> NodeIndex {
+        let node = self.child(parent, func_decl.to_string());
+        let params = self.child(node, "Params");
+        for param in &func_decl.params {
+            self.add_param(params, param);
+        }
+        self.add_block_stmt(node, &func_decl.body);
+        node
+    }
+
+    fn add_param(&mut self, parent: NodeIndex, param: &Param) -> NodeIndex {
+        self.child(parent, param.to_string())
+    }
+
+    fn add_block_stmt(&mut self, parent: NodeIndex, block_stmt: &BlockStmt) -> NodeIndex {
+        let node = self.child(parent, block_stmt.to_string());
+        for stmt in &block_stmt.statements {
+            self.add_statement(node, stmt);
+        }
+        node
+    }
+
+    fn add_statement(&mut self, parent: NodeIndex, stmt: &Statement) -> NodeIndex {
+        match stmt {
+            Statement::Return(return_stmt) => {
+                let node = self.child(parent, stmt.to_string());
+                self.add_return_stmt(node, return_stmt);
+                node
+            }
+            Statement::Block(block_stmt) => self.add_block_stmt(parent, block_stmt),
+            Statement::Expr(expr) => {
+                let node = self.child(parent, stmt.to_string());
+                self.add_expr(node, expr);
+                node
+            }
+            Statement::VarDecl(var_decl) => {
+                let node = self.child(parent, stmt.to_string());
+                self.add_var_decl(node, var_decl);
+                node
+            }
+        }
+    }
+
+    fn add_return_stmt(&mut self, parent: NodeIndex, return_stmt: &ReturnStmt) -> NodeIndex {
+        match &return_stmt.value {
+            Some(expr) => self.add_expr(parent, expr),
+            None => self.child(parent, "Void"),
+        }
+    }
+
+    fn add_expr(&mut self, parent: NodeIndex, expr: &Expr) -> NodeIndex {
+        match &expr.value {
+            ExprValue::IntLit(_) | ExprValue::Var(_) => self.child(parent, expr.value.to_string()),
+            ExprValue::BinaryOp { lhs, op: _, rhs } => {
+                let node = self.child(parent, expr.value.to_string());
+                self.add_expr(node, lhs);
+                self.add_expr(node, rhs);
+                node
+            }
+            ExprValue::FuncCall(_, args) => {
+                let node = self.child(parent, expr.value.to_string());
+                let args_node = self.child(node, "Args");
+                for arg in args {
+                    self.add_expr(args_node, arg);
+                }
+                node
+            }
+            ExprValue::Assign { lvalue, rvalue } => {
+                let node = self.child(parent, expr.value.to_string());
+                self.add_expr(node, lvalue);
+                self.add_expr(node, rvalue);
+                node
+            }
+        }
+    }
+}
@@ -1 +1,2 @@
+pub mod graph;
 pub mod types;
@@ -1,4 +1,5 @@
 use crate::{diagnostic::span::Span, frontend::types::{TokenValue, TypeIdent}};
+use std::fmt;

 pub struct CompileUnit {
    pub global_decls: Vec<GlobalDeclStmt>,
@@ -109,4 +110,112 @@ pub struct Param {
    pub name: String,
    pub param_type: Type,
    pub span: Span,
-}
+}
+
+impl fmt::Display for CompileUnit {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "CompileUnit")
+    }
+}
+
+impl fmt::Display for GlobalDeclStmt {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            GlobalDeclStmt::VarDecl(_) => write!(f, "GlobalVarDecl"),
+            GlobalDeclStmt::FuncDecl(_) => write!(f, "FuncDecl"),
+        }
+    }
+}
+
+impl fmt::Display for VarDeclStmt {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "VarDecl")
+    }
+}
+
+impl fmt::Display for VarDeclStmtValue {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{} {}", self.var_type, self.name)
+    }
+}
+
+impl fmt::Display for FuncDeclStmt {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{} {}", self.return_type, self.name)
+    }
+}
+
+impl fmt::Display for BlockStmt {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "Block")
+    }
+}
+
+impl fmt::Display for Statement {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Statement::Return(_) => write!(f, "ReturnStmt"),
+            Statement::Block(_) => write!(f, "BlockStmt"),
+            Statement::Expr(_) => write!(f, "ExprStmt"),
+            Statement::VarDecl(_) => write!(f, "VarDeclStmt"),
+        }
+    }
+}
+
+impl fmt::Display for ReturnStmt {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "ReturnStmt")
+    }
+}
+
+impl fmt::Display for Expr {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.value)
+    }
+}
+
+impl fmt::Display for ExprValue {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            ExprValue::IntLit(value) => write!(f, "IntLit({})", value),
+            ExprValue::Var(name) => write!(f, "Var({})", name),
+            ExprValue::BinaryOp { op, .. } => write!(f, "BinaryOp({})", op),
+            ExprValue::FuncCall(name, _) => write!(f, "FuncCall({})", name),
+            ExprValue::Assign { .. } => write!(f, "Assign"),
+        }
+    }
+}
+
+impl fmt::Display for Param {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{} {}", self.param_type, self.name)
+    }
+}
+
+impl fmt::Display for BinaryOp {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let op = match self {
+            BinaryOp::Add => "+",
+            BinaryOp::Sub => "-",
+            BinaryOp::Mul => "*",
+            BinaryOp::Div => "/",
+            BinaryOp::Mod => "%",
+            BinaryOp::Equal => "==",
+            BinaryOp::NotEqual => "!=",
+            BinaryOp::Less => "<",
+            BinaryOp::LessEqual => "<=",
+            BinaryOp::Greater => ">",
+            BinaryOp::GreaterEqual => ">=",
+        };
+        write!(f, "{}", op)
+    }
+}
+
+impl fmt::Display for Type {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Type::Int => write!(f, "int"),
+            Type::Void => write!(f, "void"),
+        }
+    }
+}