From 27aa6d04abdbcee509f4e02de39728178b07aaa4 Mon Sep 17 00:00:00 2001
From: willson walter <whtoo@qq.com>
Date: Wed, 29 Nov 2023 22:17:39 +0800
Subject: [PATCH] - Refactor CFG generation to use BasicBlocks - Build CFG from
 BasicBlocks instead of LinearIRBlocks - Make Label and LinearIRBlock
 comparable - Add JMP interface for jump instructions - Refactor CJMP and JMP
 to implement JMP interface - Update control flow analysis visitors to use
 generics

---
 ep18/src/main/resources/t.vm                  |   3 +-
 ep20/README.md                                | 121 ++++++++++--------
 .../org/teachfx/antlr4/ep20/Compiler.java     |   9 +-
 .../org/teachfx/antlr4/ep20/ir/IRVisitor.java |   2 +-
 .../org/teachfx/antlr4/ep20/ir/JMPInstr.java  |   7 +
 .../org/teachfx/antlr4/ep20/ir/stmt/CJMP.java |   8 +-
 .../org/teachfx/antlr4/ep20/ir/stmt/JMP.java  |   8 +-
 .../teachfx/antlr4/ep20/ir/stmt/Label.java    |  21 ++-
 .../antlr4/ep20/pass/cfg/BasicBlock.java      |  67 +++++-----
 .../org/teachfx/antlr4/ep20/pass/cfg/CFG.java |  30 +++--
 .../antlr4/ep20/pass/cfg/CFGBuilder.java      |  35 ++---
 .../ep20/pass/cfg/ControlFlowAnalysis.java    |   2 +-
 .../antlr4/ep20/pass/cfg/LinearIRBlock.java   |  81 +++++++-----
 .../ep20/pass/cfg/LivenessAnalysis.java       |   2 +-
 .../org/teachfx/antlr4/ep20/pass/cfg/Loc.java |   9 +-
 .../ep20/pass/codegen/CymbolAssembler.java    |   7 +-
 .../antlr4/ep20/pass/ir/CymbolIRBuilder.java  |  14 +-
 17 files changed, 244 insertions(+), 182 deletions(-)
 create mode 100644 ep20/src/main/java/org/teachfx/antlr4/ep20/ir/JMPInstr.java

diff --git a/ep18/src/main/resources/t.vm b/ep18/src/main/resources/t.vm
index 6d17873..d51d757 100644
--- a/ep18/src/main/resources/t.vm
+++ b/ep18/src/main/resources/t.vm
@@ -25,11 +25,10 @@ L7:
     load 0
     iconst 7
     ieq
-    brf L10
+    brf L8
 L9:
     iconst 7
     br L3
-L10:
 L8:
     sconst "break"
     print
diff --git a/ep20/README.md b/ep20/README.md
index 722cdc4..9e5fa75 100644
--- a/ep20/README.md
+++ b/ep20/README.md
@@ -35,61 +35,70 @@ graph LR
 ### CFG
 
 ```mermaid
-graph LR;
-    .def dec1: args=1 ,locals=1[t0 = @0;
-t1 =  1 ;
-t0 SUB t1;
-jmp L1;
-]
-L1[ret;
-]
-.def main: args=0 ,locals=1[t0 =  10 ;
-@0 = t0;
-jmp L4;
-]
-L4[t0 = @0;
-t1 =  0 ;
-t0 GT t1;
-jmpIf t0,L5,L6;
-]
-L5[t0 = @0;
-t1 =  5 ;
-t0 GT t1;
-jmpIf t0,L7,L8;
-]
-L7[t0 = @0;
-call print(args:1);
-t0 = @0;
-t1 =  7 ;
-t0 EQ t1;
-jmpIf t0,L9,L10;
-]
-L9[t0 =  7 ;
-jmp L3;
-]
-L10[]
-L8[t0 =  "break" ;
-call print(args:1);
-t0 = @0;
-call dec1(args:1);
-@0 = t0;
-jmp L4;
-]
-L6[t0 =  0 ;
-jmp L3;
-]
-L3[halt;
-]
-L0 --> L1
-L2 --> L4
-L4 --> L5
-L5 --> L7
-L7 --> L9
-L9 --> L10
-L10 --> L8
-L8 --> L6
-L6 --> L3
-
-
+graph TD
+    subgraph L0
+        Q112
+    end
+    subgraph L2
+        Q0["t0 =  10 ;"]
+        Q1["@0 = t0;"]
+        Q2["jmp L4;"]
+    end
+    subgraph L3
+        Q3["t0 = @0;"]
+        Q4["t1 =  0 ;"]
+        Q5["t0 GT t1;"]
+        Q6["jmpIf t0,L5,L6;"]
+    end
+    subgraph L4
+        Q7["t0 = @0;"]
+        Q8["t1 =  5 ;"]
+        Q9["t0 GT t1;"]
+        Q10["jmpIf t0,L7,L8;"]
+    end
+    subgraph L5
+        Q11["t0 = @0;"]
+        Q12["call print(args:1);"]
+        Q13["t0 = @0;"]
+        Q14["t1 =  7 ;"]
+        Q15["t0 EQ t1;"]
+        Q16["jmpIf t0,L9,L10;"]
+    end
+    subgraph L6
+        Q17["t0 =  7 ;"]
+        Q18["jmp L3;"]
+    end
+    subgraph L7
+        Q000
+    end
+    subgraph L8
+        Q19["t0 =  'break' ;"]
+        Q20["call print(args:1);"]
+        Q21["t0 = @0;"]
+        Q22["call dec1(args:1);"]
+        Q23["@0 = t0;"]
+        Q24["jmp L4;"]
+    end
+    subgraph L9
+        Q25["t0 =  0 ;"]
+        Q26["jmp L3;"]
+    end
+    subgraph L10
+        Q27["halt;"]
+    end
+    
+    L0 --> L4
+    L4 --> L5
+    L4 --> L6
+    L5 --> L7
+    L5 --> L8
+    L7 --> L0
+    L7 --> L9
+    L9 --> L3
+    L9 --> L10
+    L10 --> L8
+    L8 --> L4
+    L8 --> L6
+    L6 --> L3
 ```
 
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/Compiler.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/Compiler.java
index 0caba73..41cc9a5 100644
--- a/ep20/src/main/java/org/teachfx/antlr4/ep20/Compiler.java
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/Compiler.java
@@ -48,12 +48,13 @@ public static void main(String[] args) throws IOException {
         var irBuilder = new CymbolIRBuilder();
 
         astRoot.accept(irBuilder);
-
-        var cfg = irBuilder.getCFG(irBuilder.prog.blockList);
-        logger.info("CFG:\n" + cfg.toString());
+        for(var block : irBuilder.prog.blockList) {
+            var cfg = irBuilder.getCFG(List.of(block));
+            logger.info("CFG:\n" + cfg.toString());
+        }
 
         var assembler = new CymbolAssembler();
-        irBuilder.prog.accept(assembler);
+        assembler.visit(irBuilder.prog.linearInstrs());
         saveToEp18Res(assembler.getAsmInfo());
         logger.info("\n%s".formatted(assembler.getAsmInfo()));
     }
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/IRVisitor.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/IRVisitor.java
index 2cb389e..f6d80a3 100644
--- a/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/IRVisitor.java
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/IRVisitor.java
@@ -26,7 +26,7 @@ public interface IRVisitor<S,E> {
     S visit(ReturnVal returnVal);
     default S visit(ExprStmt exprStmt) { return exprStmt.accept(this); }
 
-    S visit(Prog prog);
+    default S visit(Prog prog) { return null; }
 
     E visit(OperandSlot operandSlot);
 
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/JMPInstr.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/JMPInstr.java
new file mode 100644
index 0000000..e800b6f
--- /dev/null
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/JMPInstr.java
@@ -0,0 +1,7 @@
+package org.teachfx.antlr4.ep20.ir;
+
+import org.teachfx.antlr4.ep20.ir.stmt.Label;
+
+public interface JMPInstr {
+    public Label getTarget();
+}
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/stmt/CJMP.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/stmt/CJMP.java
index ec27d1c..4941e28 100644
--- a/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/stmt/CJMP.java
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/stmt/CJMP.java
@@ -1,10 +1,11 @@
 package org.teachfx.antlr4.ep20.ir.stmt;
 
 import org.teachfx.antlr4.ep20.ir.IRVisitor;
+import org.teachfx.antlr4.ep20.ir.JMPInstr;
 import org.teachfx.antlr4.ep20.ir.expr.VarSlot;
 import org.teachfx.antlr4.ep20.pass.cfg.LinearIRBlock;
 
-public class CJMP extends Stmt {
+public class CJMP extends Stmt implements JMPInstr {
     public VarSlot cond;
     private LinearIRBlock thenBlock;
     private LinearIRBlock elseBlock;
@@ -18,6 +19,11 @@ public CJMP(VarSlot cond, LinearIRBlock thenLabel, LinearIRBlock elseLabel) {
         elseBlock.refJMP(this);
     }
 
+    @Override
+    public Label getTarget() {
+        return elseBlock.getLabel();
+    }
+
     @Override
     public <S, E> S accept(IRVisitor<S, E> visitor) {
         return visitor.visit(this);
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/stmt/JMP.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/stmt/JMP.java
index 77b8853..9f4698f 100644
--- a/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/stmt/JMP.java
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/stmt/JMP.java
@@ -1,9 +1,10 @@
 package org.teachfx.antlr4.ep20.ir.stmt;
 
 import org.teachfx.antlr4.ep20.ir.IRVisitor;
+import org.teachfx.antlr4.ep20.ir.JMPInstr;
 import org.teachfx.antlr4.ep20.pass.cfg.LinearIRBlock;
 
-public class JMP extends Stmt
+public class JMP extends Stmt implements JMPInstr
 {
     @Override
     public <S, E> S accept(IRVisitor<S, E> visitor) {
@@ -17,6 +18,11 @@ public JMP(LinearIRBlock block)
     }
     public LinearIRBlock next;
 
+    @Override
+    public Label getTarget() {
+        return next.getLabel();
+    }
+
     @Override
     public StmtType getStmtType() {
         return StmtType.JMP;
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/stmt/Label.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/stmt/Label.java
index 07dd151..423a513 100644
--- a/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/stmt/Label.java
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/ir/stmt/Label.java
@@ -1,11 +1,12 @@
 package org.teachfx.antlr4.ep20.ir.stmt;
 
+import org.jetbrains.annotations.NotNull;
 import org.teachfx.antlr4.ep20.ir.IRVisitor;
 import org.teachfx.antlr4.ep20.symtab.scope.Scope;
 
 import java.util.Objects;
 
-public class Label extends Stmt {
+public class Label extends Stmt implements Comparable<Label> {
 
     private Stmt nextEntry = null;
 
@@ -23,6 +24,12 @@ public Label(String rawLabel, Scope scope) {
         }
     }
 
+    public Label(Scope scope,Integer ord) {
+        this.seq = ord;
+        this.rawLabel = "L%d".formatted(seq);
+        this.scope = scope;
+    }
+
     public Label(Scope scope) {
         this.scope = scope;
         this.seq = scope.getLabelSeq();
@@ -93,4 +100,16 @@ public Scope getScope() {
     public int getSeq() {
         return seq;
     }
+
+    @Override
+    public int hashCode() {
+        return getSeq();
+    }
+
+    @Override
+    public int compareTo(@NotNull Label o) {
+        var e = Integer.valueOf(getSeq());
+        var t = Integer.valueOf(o.getSeq());
+        return e.compareTo(t);
+    }
 }
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/BasicBlock.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/BasicBlock.java
index 2a09370..4447db3 100644
--- a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/BasicBlock.java
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/BasicBlock.java
@@ -14,29 +14,45 @@
 
 public class BasicBlock<I extends IRNode> implements Comparable<BasicBlock<I>>, Iterable<Loc<I>> {
 
-    @Override
-    public int compareTo(@NotNull BasicBlock<I> o) {
-        return this.id - o.id;
-    }
-
-
+    public final int id;
     // Generate codes
     public List<Loc<I>> codes;
+    public Kind kind;
+    // For data flow analysis
+    public Set<Operand> def;
+    public Set<Operand> liveUse;
+    public Set<Operand> liveIn;
+    public Set<Operand> liveOut;
+    protected Label label;
 
-    public final int id;
+    public BasicBlock(Kind kind, List<Loc<I>> codes, Label label) {
+        this.codes = codes;
+        this.label = label;
+        this.id = label.getSeq();
+        this.kind = kind;
+    }
 
-    protected Optional<Label> label;
+    @NotNull
+    @Contract("_ -> new")
+    public static BasicBlock<IRNode> buildFromLinearBlock(@NotNull LinearIRBlock block) {
+        return new BasicBlock<IRNode>(block.getKind(), block.getStmts().stream().map(Loc::new).toList(), block.getLabel());
+    }
 
-    public Kind kind;
+    @Override
+    public int compareTo(@NotNull BasicBlock<I> o) {
+        return this.id - o.id;
+    }
 
     @NotNull
     @Override
     public Iterator<Loc<I>> iterator() {
         return codes.iterator();
     }
+
     public Iterator<Loc<I>> backwardIterator() {
         return new Iterator<Loc<I>>() {
             private int index = codes.size() - 1;
+
             @Override
             public boolean hasNext() {
                 return index != -1;
@@ -51,24 +67,11 @@ public Loc<I> next() {
         };
     }
 
-    @NotNull
-    @Contract("_ -> new")
-    public static BasicBlock<IRNode> buildFromLinearBlock(@NotNull LinearIRBlock block) {
-        return new BasicBlock<IRNode>(block.getKind(), block.getOrd(), block.getStmts().stream().map(Loc::new).toList(), Optional.of(block.getLabel()));
-    }
-
-    public BasicBlock(Kind kind, int id, List<Loc<I>> codes, Optional<Label> label) {
-        this.codes = codes;
-        this.label = label;
-        this.id = id;
-        this.kind = kind;
-    }
-
     public int getId() {
         return id;
     }
 
-    public Optional<Label> getLabel() {
+    public Label getLabel() {
         return label;
     }
 
@@ -80,14 +83,16 @@ public String getOrdLabel() {
     public boolean isEmpty() {
         return codes.isEmpty();
     }
+
     public List<Loc<I>> allSeq() {
-        if (kind.equals(Kind.CONTINUOUS)){
+        if (kind.equals(Kind.CONTINUOUS)) {
             return codes;
         }
         return codes.subList(0, codes.size() - 1);
     }
 
     public List<Loc<I>> dropLabelSeq() {
+        if (codes.size() <= 1) return codes;
         return codes.subList(1, codes.size());
     }
 
@@ -95,13 +100,11 @@ public I getLastInstr() {
         return codes.get(codes.size() - 1).instr;
     }
 
-    // For data flow analysis
-    public Set<Operand> def;
-
-    public Set<Operand> liveUse;
-
-    public Set<Operand> liveIn;
-
-    public Set<Operand> liveOut;
+    public void mergeNearBlock(BasicBlock<I> nextBlock) {
+        /// remove last jump instr
+        codes.remove(codes.size() - 1);
+        /// merge instr
+        codes.addAll(nextBlock.codes);
+    }
 
 }
\ No newline at end of file
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/CFG.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/CFG.java
index 7f4d67d..06fb696 100644
--- a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/CFG.java
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/CFG.java
@@ -17,25 +17,27 @@ public class CFG<I extends IRNode> implements Iterable<BasicBlock<I>> {
     private final static Logger logger = LogManager.getLogger(CFG.class);
     public final List<BasicBlock<I>> nodes;
 
-    public final List<Pair<Integer,Integer>> edges;
+    public final List<Pair<Integer, Integer>> edges;
 
     private final List<Pair<Set<Integer>, Set<Integer>>> links;
 
-    public CFG(List<BasicBlock<I>> nodes, List<Pair<Integer, Integer>> edges) {
+    public CFG(Map<Integer, BasicBlock<I>> nodes, List<Pair<Integer, Integer>> edges) {
         // Generate init
-        this.nodes = new LinkedList<>(nodes);
+        var lastOrd = Integer.max(nodes.keySet().stream().max(Integer::compareTo).get(), nodes.size()) + 1;
+
+        this.nodes = new LinkedList<>(nodes.values());
         this.edges = edges;
 
         links = new ArrayList<>();
-        for (var i = 0;i < nodes.size();i++) {
+        for (var i = 0; i < lastOrd; i++) {
             links.add(Pair.of(new TreeSet<>(), new TreeSet<>()));
         }
 
         for (var edge : edges) {
-             var u = edge.getLeft();
-             var v = edge.getRight();
-             links.get(u).getRight().add(v);
-             links.get(v).getLeft().add(u);
+            var u = edge.getLeft();
+            var v = edge.getRight();
+            links.get(u).getRight().add(v);
+            links.get(v).getLeft().add(u);
         }
     }
 
@@ -65,21 +67,25 @@ public Iterator<BasicBlock<I>> iterator() {
         return nodes.iterator();
     }
 
+    public List<I> simplifyIRInstrs() {
+
+        return null;
+    }
 
     @Override
     public String toString() {
-        var graphRenderBuffer = new StringBuilder("graph LR\n");
+        var graphRenderBuffer = new StringBuilder("graph TD\n");
         AtomicInteger i = new AtomicInteger();
 
-        for(var node : nodes){
+        for (var node : nodes.stream().sorted((b1,b2)-> b2.id - b1.id).toList()) {
 
             graphRenderBuffer.append("subgraph ").append(node.getOrdLabel()).append("\n");
-            node.dropLabelSeq().stream().map(x -> x.instr.toString()).map(x ->  "Q"+ (i.getAndIncrement()) + "[\"" + x + ";\"]\n").forEach(graphRenderBuffer::append);
+            node.dropLabelSeq().stream().map(x -> x.instr.toString()).map(x -> "Q" + (i.getAndIncrement()) + "[\"" + x + ";\"]\n").forEach(graphRenderBuffer::append);
             graphRenderBuffer.append("end").append("\n");
 
         }
 
-        for (var edge : edges){
+        for (var edge : edges) {
 
             graphRenderBuffer.append("L").append(edge.getLeft()).append(" --> ").append("L").append(edge.getRight()).append("\n");
         }
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/CFGBuilder.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/CFGBuilder.java
index 8ca3a6e..00d8848 100644
--- a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/CFGBuilder.java
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/CFGBuilder.java
@@ -6,11 +6,13 @@
 import java.util.*;
 
 public class CFGBuilder {
-    private CFG<IRNode> cfg;
-    private List<BasicBlock<IRNode>> basicBlocks;
-    private List<Pair<Integer,Integer>> edges;
-    public CFGBuilder(List<LinearIRBlock> blockList){
-        basicBlocks =new ArrayList<>();
+    private static final Set<String> cachedEdgeLinks = new HashSet<>();
+    private final CFG<IRNode> cfg;
+    private final Map<Integer, BasicBlock<IRNode>> basicBlocks;
+    private final List<Pair<Integer, Integer>> edges;
+
+    public CFGBuilder(List<LinearIRBlock> blockList) {
+        basicBlocks = new HashMap<>();
         edges = new ArrayList<>();
         for (var funcLabelBlock : blockList) {
             build(funcLabelBlock);
@@ -18,28 +20,9 @@ public CFGBuilder(List<LinearIRBlock> blockList){
         cfg = new CFG<>(basicBlocks, edges);
     }
 
-    private static final Set<String> cachedEdgeLinks = new HashSet<>();
+    private void build(LinearIRBlock block) {
+        basicBlocks.put(block.getLabel().getSeq(), BasicBlock.buildFromLinearBlock(block));
 
-    private void build(LinearIRBlock block){
-
-        basicBlocks.add(BasicBlock.buildFromLinearBlock(block));
-        block.getJumpEntries().ifPresent(entries -> {
-            for(var dest : entries){
-                var key = "%d-%d".formatted(block.getOrd(), dest);
-                if (!cachedEdgeLinks.contains(key)) {
-                    edges.add(Pair.of(block.getOrd(), dest));
-                    cachedEdgeLinks.add("%d-%d".formatted(block.getOrd(), dest));
-                }
-            }
-        });
-        for(var successor : block.getSuccessors()){
-            var key = "%d-%d".formatted(block.getOrd(), successor.getOrd());
-            if (!cachedEdgeLinks.contains(key)) {
-                cachedEdgeLinks.add(key);
-                edges.add(Pair.of(block.getOrd(), successor.getOrd()));
-            }
-            build(successor);
-        }
     }
 
     public CFG<IRNode> getCFG() {
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/ControlFlowAnalysis.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/ControlFlowAnalysis.java
index 83fd763..c61a042 100644
--- a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/ControlFlowAnalysis.java
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/ControlFlowAnalysis.java
@@ -10,7 +10,7 @@
 import org.teachfx.antlr4.ep20.ir.expr.val.ConstVal;
 import org.teachfx.antlr4.ep20.ir.stmt.*;
 
-public class ControlFlowAnalysis implements IRVisitor<Void,Void> {
+public class ControlFlowAnalysis implements IRVisitor<Void, Void> {
 
     @Override
     public <T> Void visit(ConstVal<T> tConstVal) {
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/LinearIRBlock.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/LinearIRBlock.java
index 07b198f..30a44a9 100644
--- a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/LinearIRBlock.java
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/LinearIRBlock.java
@@ -1,23 +1,30 @@
 package org.teachfx.antlr4.ep20.pass.cfg;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.jetbrains.annotations.NotNull;
 import org.teachfx.antlr4.ep20.ir.IRNode;
+import org.teachfx.antlr4.ep20.ir.JMPInstr;
+import org.teachfx.antlr4.ep20.ir.expr.Operand;
 import org.teachfx.antlr4.ep20.ir.stmt.*;
 import org.teachfx.antlr4.ep20.symtab.scope.Scope;
 import org.teachfx.antlr4.ep20.utils.Kind;
 
 import java.util.*;
 
-public class LinearIRBlock {
+public class LinearIRBlock implements Comparable<LinearIRBlock> {
 
     // Fields
     private static int LABEL_SEQ = 0;
+    private static final Logger logger = LogManager.getLogger(LinearIRBlock.class);
     private Kind kind = Kind.CONTINUOUS;
     private int ord = 0;
     private ArrayList<IRNode> stmts;
     private List<LinearIRBlock> successors;
     private List<LinearIRBlock> predecessors;
+
     private Scope scope = null;
-    private List<IRNode> jmpRefMap = new ArrayList<>();
+    private List<JMPInstr> jmpRefMap = new ArrayList<>();
 
     // Constructor
     public LinearIRBlock() {
@@ -25,6 +32,17 @@ public LinearIRBlock() {
         successors = new ArrayList<>();
         predecessors = new ArrayList<>();
         ord = LABEL_SEQ++;
+        logger.info(ord);
+    }
+
+    public static boolean isBasicBlock(Stmt stmt) {
+        return !(stmt instanceof CJMP) && !(stmt instanceof JMP);
+    }
+
+    // Link Operations
+    public static void setLink(LinearIRBlock current, LinearIRBlock next) {
+        current.successors.add(next);
+        next.predecessors.add(current);
     }
 
     // Methods
@@ -46,10 +64,6 @@ private void updateKindByLastInstr(IRNode stmt) {
         }
     }
 
-    public static boolean isBasicBlock(Stmt stmt) {
-        return !(stmt instanceof CJMP) && !(stmt instanceof JMP);
-    }
-
     // Getters and Setters
     public List<IRNode> getStmts() {
         return stmts;
@@ -85,11 +99,11 @@ public void setScope(Scope scope) {
         this.scope = scope;
     }
 
-    public List<IRNode> getJmpRefMap() {
+    public List<JMPInstr> getJmpRefMap() {
         return jmpRefMap;
     }
 
-    public void setJmpRefMap(List<IRNode> jmpRefMap) {
+    public void setJmpRefMap(List<JMPInstr> jmpRefMap) {
         this.jmpRefMap = jmpRefMap;
     }
 
@@ -97,12 +111,6 @@ public int getOrd() {
         return ord;
     }
 
-    // Link Operations
-    public static void setLink(LinearIRBlock current, LinearIRBlock next) {
-        current.successors.add(next);
-        next.predecessors.add(current);
-    }
-
     public void setLink(LinearIRBlock next) {
         LinearIRBlock.setLink(this, next);
     }
@@ -112,7 +120,7 @@ public Kind getKind() {
     }
 
     // Jump Operations
-    public void refJMP(IRNode node) {
+    public void refJMP(JMPInstr node) {
         jmpRefMap.add(node);
     }
 
@@ -128,7 +136,7 @@ public String toString() {
         return "L" + ord;
     }
 
-    public String toSource(){
+    public String toSource() {
         var firstInstr = stmts.get(0);
 
         if (firstInstr instanceof FuncEntryLabel) {
@@ -140,32 +148,35 @@ public String toSource(){
 
     public Label getLabel() {
         if (stmts.isEmpty()) {
-            return new Label(toString(),scope);
+            return new Label(toString(), scope);
         }
 
-
         var firstInstr = stmts.get(0);
-        if (firstInstr instanceof FuncEntryLabel funcEntryLabel) {
-            return funcEntryLabel;
+        if (firstInstr instanceof Label label) {
+
+            return label;
+        } else {
+            stmts.add(0, new Label(scope, getOrd()));
+            firstInstr = stmts.get(0);
         }
 
-        return new Label(toString(),scope);
+        return (Label) firstInstr;
     }
 
-    public Optional<TreeSet<Integer>> getJumpEntries() {
-        var entries = new TreeSet<Integer>();
+    public Optional<TreeSet<LinearIRBlock>> getJumpEntries() {
+        var entries = new TreeSet<LinearIRBlock>();
         switch (kind) {
             case END_BY_CJMP -> {
-                CJMP cjmp = (CJMP) getStmts().getLast();
-                var tid = cjmp.getThenBlock().getOrd();
-                var eid = cjmp.getElseBlock().getOrd();
-                entries.add(tid);
-                entries.add(eid);
+                CJMP cjmp = (CJMP) stmts.get(stmts.size() - 1);
+                var targetLb = cjmp.getThenBlock();
+                var elseLb = cjmp.getElseBlock();
+                entries.add(targetLb);
+                entries.add(elseLb);
                 return Optional.of(entries);
             }
             case END_BY_JMP -> {
-                JMP cjmp = (JMP) getStmts().getLast();
-                var tid = cjmp.next.getOrd();
+                JMP cjmp = (JMP) stmts.get(stmts.size() - 1);
+                var tid = cjmp.next;
                 entries.add(tid);
                 return Optional.of(entries);
             }
@@ -175,4 +186,14 @@ public Optional<TreeSet<Integer>> getJumpEntries() {
         }
     }
 
+    @Override
+    public int compareTo(@NotNull LinearIRBlock o) {
+        return  ord - o.getOrd();
+    }
+
+    @Override
+    public int hashCode() {
+        return (ord * 177) % 71;
+    }
+
 }
\ No newline at end of file
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/LivenessAnalysis.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/LivenessAnalysis.java
index 4c35569..29d715b 100644
--- a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/LivenessAnalysis.java
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/LivenessAnalysis.java
@@ -10,7 +10,7 @@
 import org.teachfx.antlr4.ep20.ir.expr.val.ConstVal;
 import org.teachfx.antlr4.ep20.ir.stmt.*;
 
-public class LivenessAnalysis implements IRVisitor<Void,Void> {
+public class LivenessAnalysis implements IRVisitor<Void, Void> {
 
     @Override
     public Void visit(BinExpr node) {
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/Loc.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/Loc.java
index 9bcb013..18bc746 100644
--- a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/Loc.java
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/cfg/Loc.java
@@ -5,17 +5,16 @@
 
 import java.util.Set;
 
-public class Loc <I extends IRNode> {
+public class Loc<I extends IRNode> {
     // arrange properties
     public final I instr;
+    public Set<Operand> liveIn;
+    public Set<Operand> liveOut;
+
     public Loc(I instr) {
         this.instr = instr;
     }
 
-    public Set<Operand> liveIn;
-
-    public Set<Operand> liveOut;
-
     @Override
     public String toString() {
         return "Loc{" +
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/codegen/CymbolAssembler.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/codegen/CymbolAssembler.java
index 60d9e5f..2525550 100644
--- a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/codegen/CymbolAssembler.java
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/codegen/CymbolAssembler.java
@@ -1,5 +1,6 @@
 package org.teachfx.antlr4.ep20.pass.codegen;
 
+import org.teachfx.antlr4.ep20.ir.IRNode;
 import org.teachfx.antlr4.ep20.ir.IRVisitor;
 import org.teachfx.antlr4.ep20.ir.Prog;
 import org.teachfx.antlr4.ep20.ir.expr.CallFunc;
@@ -12,6 +13,7 @@
 import org.teachfx.antlr4.ep20.ir.stmt.*;
 
 import java.util.LinkedList;
+import java.util.List;
 import java.util.Objects;
 
 public class CymbolAssembler implements IRVisitor<Void,Void> {
@@ -30,9 +32,8 @@ public String getAsmInfo() {
     }
     private int indents = 0;
 
-    @Override
-    public Void visit(Prog prog) {
-        for (var instr: prog.linearInstrs()) {
+    public Void visit(List<IRNode> linearInstrs) {
+        for (var instr: linearInstrs) {
             if (instr instanceof Expr) {
                 ((Expr) instr).accept(this);
             } else  {
diff --git a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/ir/CymbolIRBuilder.java b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/ir/CymbolIRBuilder.java
index 052f576..7d13b1d 100644
--- a/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/ir/CymbolIRBuilder.java
+++ b/ep20/src/main/java/org/teachfx/antlr4/ep20/pass/ir/CymbolIRBuilder.java
@@ -24,6 +24,7 @@
 import org.teachfx.antlr4.ep20.pass.cfg.CFG;
 import org.teachfx.antlr4.ep20.pass.cfg.CFGBuilder;
 import org.teachfx.antlr4.ep20.pass.cfg.LinearIRBlock;
+import org.teachfx.antlr4.ep20.symtab.scope.Scope;
 import org.teachfx.antlr4.ep20.symtab.symbol.MethodSymbol;
 import org.teachfx.antlr4.ep20.symtab.symbol.VariableSymbol;
 
@@ -75,12 +76,14 @@ public Void visit(FuncDeclNode funcDeclNode) {
         var methodSymbol = (MethodSymbol) funcDeclNode.getRefSymbol();
         var entryLabel = new FuncEntryLabel(methodSymbol.getName(),methodSymbol.getArgs(),methodSymbol.getLocals(),methodSymbol);
         // Expand
-        forkNewBlock();
+        forkNewBlock(methodSymbol);
         var startBlock = currentBlock;
+
         evalExprStack = new Stack<>();
         getCurrentBlock().addStmt(entryLabel);
 
         var exitBlock = new LinearIRBlock();
+        exitBlock.setScope(methodSymbol);
         var exitEntry = new ReturnVal(null,methodSymbol);
 
         setExitHook(exitEntry,methodSymbol);
@@ -308,8 +311,9 @@ public Void visit(ContinueStmtNode continueStmtNode) {
         return null;
     }
 
-    public void forkNewBlock() {
+    public void forkNewBlock(Scope scope) {
         currentBlock = new LinearIRBlock();
+        currentBlock.setScope(scope);
         breakStack = new Stack<>();
         continueStack = new Stack<>();
     }
@@ -432,10 +436,8 @@ public CFG<IRNode> getCFG(List<LinearIRBlock> blocks) {
     }
 
     public void insertBlockLabel(LinearIRBlock startBlock) {
-        var firstInstr = startBlock.getLabel();
-        if (!(firstInstr instanceof FuncEntryLabel)) {
-            startBlock.getStmts().add(0,firstInstr);
-        }
+        startBlock.getLabel();
+
         for (var successor : startBlock.getSuccessors()) {
             insertBlockLabel(successor);
         }