1 files changed, 43 insertions, 7 deletions
diff --git a/Intermediate.hs b/Intermediate.hs
index 95f6cc7..0feeb13 100644
--- a/Intermediate.hs
+++ b/Intermediate.hs
@@ -2,6 +2,7 @@ module Intermediate where
 
 import Data.List
 import qualified Data.Map.Strict as Map
+import Data.Maybe
 
 import AST
 
@@ -53,12 +54,44 @@ bidOf (BB i _ _) = i
 termOf :: BB -> Terminator
 termOf (BB _ _ t) = t
 
+class AllRefs a where
+    allRefs :: a -> [Ref]
+
+instance AllRefs InsCode where
+    allRefs (IAssign r) = [r]
+    allRefs (IParam _) = []
+    allRefs (IClosure _) = []
+    allRefs (IData _) = []
+    allRefs (ICallC r rs) = r : rs
+    allRefs (IAllocClo _ rs) = rs
+    allRefs (IDiscard r) = [r]
+
+instance AllRefs Terminator where
+    allRefs (IBr r _ _) = [r]
+    allRefs (IJmp _) = []
+    allRefs (IRet r) = [r]
+    allRefs (ITailC r rs) = r : rs
+    allRefs IExit = []
+    allRefs IUnknown = []
+
 
 instance Show IRProgram where
     show (IRProgram bbs gfds datas) = intercalate "\n" $
         ["IRPROGRAM", "Data Table:"] ++ map (("- " ++) . show) datas ++
         ["Global functions:"] ++ map (\(n, gfd) -> "- " ++ n ++ ": " ++ show gfd) (Map.assocs gfds) ++
-        ["Blocks:"] ++ [intercalate "\n" (map show bbs)]
+        ["Blocks:"] ++ [intercalate "\n" (map (genericShowBB icshow termshow) bbs)]
+      where
+        annotate s "" = s
+        annotate s a = s ++ "   ; " ++ a
+        refAnnot rs = intercalate ", " . catMaybes $
+                        [case Map.lookup name gfds of
+                             Nothing -> Nothing
+                             Just (GlobFuncDef i _ _) -> Just (name ++ " = BB " ++ show i)
+                        | RSClo name <- nub rs]
+        safeIndex l i = if 0 <= i && i < length l then Just (l !! i) else Nothing
+        icshow ins@(IData n) = annotate (show ins) (maybe "??" show (datas `safeIndex` n))
+        icshow ins = annotate (show ins) (refAnnot (allRefs ins))
+        termshow term = annotate (show term) (refAnnot (allRefs term))
 
 instance Show GlobFuncDef where
     show (GlobFuncDef bbid na []) = "BB " ++ show bbid ++ " (" ++ show na ++ ")"
@@ -66,12 +99,7 @@ instance Show GlobFuncDef where
         "BB " ++ show bbid ++ " (" ++ show na ++ ") (closure slots: " ++ intercalate ", " cs ++ ")"
 
 instance Show BB where
-    show (BB i inss term) =
-        "BB " ++ show i ++
-        concatMap (\(r, ic) -> case r of
-            RNone -> "\n  " ++ show ic
-            _ -> "\n  " ++ show r ++ " <- " ++ show ic) inss ++
-        "\n  " ++ show term
+    show = genericShowBB show show
 
 instance Show Ref where
     show (RConst n) = show n
@@ -95,3 +123,11 @@ instance Show Terminator where
     show (ITailC r as) = "tailc " ++ show r ++ " " ++ show as
     show IExit = "exit"
     show IUnknown = "<<UNKNOWN>>"
+
+genericShowBB :: (InsCode -> String) -> (Terminator -> String) -> BB -> String
+genericShowBB icshow termshow (BB i inss term) =
+    "BB " ++ show i ++
+    concatMap (\(r, ic) -> case r of
+        RNone -> "\n  " ++ icshow ic
+        _ -> "\n  " ++ show r ++ " <- " ++ icshow ic) inss ++
+    "\n  " ++ termshow term