From fd1043e09b05d6bd8274539e02b700f59ec2bba8 Mon Sep 17 00:00:00 2001
From: Minseong Jang <minseong.jang@cp.kaist.ac.kr>
Date: Thu, 27 Jan 2022 15:21:56 +0900
Subject: [PATCH] Update IR visualizer

---
 src/ir/visualize.rs | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/ir/visualize.rs b/src/ir/visualize.rs
index f047dfe..c5ecf68 100644
--- a/src/ir/visualize.rs
+++ b/src/ir/visualize.rs
@@ -48,9 +48,8 @@ impl Translate<TranslationUnit> for Visualizer {
                 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);
+                            let to = self.translate_callee(name, callee)?;
 
                             edges.push(format!("{} -> {};", from, to));
                         }
@@ -93,11 +92,15 @@ impl Visualizer {
     }
 
     #[inline]
-    fn translate_callee(&mut self, callee: &Operand) -> Result<String, ()> {
+    fn translate_callee(&mut self, name: &str, callee: &Operand) -> Result<String, ()> {
         match callee {
             Operand::Constant(_constant @ Constant::GlobalVariable { name, .. }) => {
-                Ok(name.clone())
+                Ok(self.get_function_first_instruction(name))
             }
+            Operand::Register {
+                rid: _rid @ RegisterId::Temp { bid, iid },
+                ..
+            } => Ok(self.translate_instruction_node(name, *bid, *iid)),
             _ => todo!("translate_callee: operand {:?}", callee),
         }
     }