MLIR DAG rewriter w/ SSA IR, dominance and CFG analysis passes
Initial
author
vm.fail
date
(Jul 23, 2026, 3:40 PM UTC)
commit
ee8d504c
ee8d504ca88d5799ebf404057a28e90414ecd178
vm.fail
date
(Jul 23, 2026, 3:40 PM UTC)
commit
ee8d504c
ee8d504ca88d5799ebf404057a28e90414ecd178
+606
+34
Makefile
+34
Makefile
···
1
+
OCAMLFLAGS = -g
2
+
OCAMLC = ocamlc
3
+
4
+
LIB_DIR = lib
5
+
TEST_DIR = tests
6
+
7
+
LIB_FILES = ir.ml cfg.ml dominance.ml nemo.ml parser.ml printer.ml
8
+
TEST_FILES = test.ml
9
+
MAIN_FILES = main.ml
10
+
11
+
all: nemo test
12
+
13
+
nemo: $(LIB_FILES:%.ml=$(LIB_DIR)/%.cmo) $(MAIN_FILES:%.ml=%.cmo)
14
+
$(OCAMLC) $(OCAMLFLAGS) -o $@ $^
15
+
16
+
test: $(LIB_FILES:%.ml=$(LIB_DIR)/%.cmo) $(TEST_DIR)/test.cmo
17
+
$(OCAMLC) $(OCAMLFLAGS) -o $@ $^
18
+
19
+
$(LIB_DIR)/%.cmo: $(LIB_DIR)/%.ml
20
+
cd $(LIB_DIR) && $(OCAMLC) $(OCAMLFLAGS) -c $(notdir $<)
21
+
22
+
$(TEST_DIR)/%.cmo: $(TEST_DIR)/%.ml
23
+
cd $(TEST_DIR) && $(OCAMLC) $(OCAMLFLAGS) -c -I ../$(LIB_DIR) $(notdir $<)
24
+
25
+
%.cmo: %.ml
26
+
$(OCAMLC) $(OCAMLFLAGS) -I $(LIB_DIR) -c $<
27
+
28
+
clean:
29
+
rm -f nemo test
30
+
rm -f $(LIB_DIR)/*.cmo $(LIB_DIR)/*.cmi
31
+
rm -f $(TEST_DIR)/*.cmo $(TEST_DIR)/*.cmi
32
+
rm -f *.cmo *.cmi
33
+
34
+
.PHONY: all clean
+7
README.md
+7
README.md
assets/p1.png
assets/p1.png
This is a binary file and will not be displayed.
assets/p2.png
assets/p2.png
This is a binary file and will not be displayed.
+8
examples/basic.ir
+8
examples/basic.ir
+15
examples/phi.ir
+15
examples/phi.ir
+8
examples/simple.ir
+8
examples/simple.ir
+38
lib/cfg.ml
+38
lib/cfg.ml
···
1
+
open Ir
2
+
3
+
module IntMap = Map.Make(Int)
4
+
module IntSet = Set.Make(Int)
5
+
6
+
type cfg = {
7
+
blocks : block IntMap.t;
8
+
entry : value;
9
+
}
10
+
11
+
let build_cfg (f : func) =
12
+
let blocks = List.fold_left (fun m b -> IntMap.add b.id b m) IntMap.empty f.blocks in
13
+
let cfg = { blocks; entry = f.entry } in
14
+
let add_edges b =
15
+
let rec scan acc = function
16
+
| [] -> List.rev acc
17
+
| Ret _ :: _ -> List.rev acc
18
+
| BinOp (_, _, t) :: rest -> scan (t :: acc) rest
19
+
| Br t :: rest -> scan (t :: acc) rest
20
+
| CondBr (_, t1, t2) :: rest -> scan (t2 :: t1 :: acc) rest
21
+
| _ :: rest -> scan acc rest
22
+
in
23
+
b.succs <- scan [] b.insts
24
+
in
25
+
IntMap.iter (fun _ b -> add_edges b) blocks;
26
+
IntMap.iter (fun _ b ->
27
+
List.iter (fun s ->
28
+
try
29
+
let sb = IntMap.find s cfg.blocks in
30
+
sb.preds <- b.id :: sb.preds
31
+
with Not_found -> ()
32
+
) b.succs
33
+
) blocks;
34
+
cfg
35
+
36
+
let get_block cfg id = IntMap.find id cfg.blocks
37
+
let get_succs cfg id = (get_block cfg id).succs
38
+
let get_preds cfg id = (get_block cfg id).preds
+76
lib/dominance.ml
+76
lib/dominance.ml
···
1
+
open Ir
2
+
open Cfg
3
+
4
+
module IntMap = Map.Make(Int)
5
+
module IntSet = Set.Make(Int)
6
+
7
+
(* Immediate dominator, derived correctly from full dominator sets.
8
+
The idom of n is the unique proper dominator d of n that is itself
9
+
dominated by every other proper dominator of n (i.e. the deepest one).
10
+
This replaces the previous ad-hoc heuristic, which could mis-pick an outer
11
+
dominator when a node had 3+ dominators. *)
12
+
let idom cfg doms id =
13
+
let s = IntMap.find id doms in
14
+
let without = IntSet.diff s (IntSet.singleton id) in
15
+
if IntSet.is_empty without then None
16
+
else
17
+
let strict_doms = IntSet.elements without in
18
+
(* Find the d in strict_doms that every other strict dominator dominates. *)
19
+
let rec find = function
20
+
| [] -> None
21
+
| d :: rest ->
22
+
if List.for_all (fun m -> m = d || IntSet.mem d (IntMap.find m doms)) strict_doms
23
+
then Some d
24
+
else find rest
25
+
in
26
+
find strict_doms
27
+
28
+
let dominators cfg =
29
+
let all = IntMap.fold (fun k _ s -> IntSet.add k s) cfg.blocks IntSet.empty in
30
+
let init = IntSet.singleton cfg.entry in
31
+
let rec fixpoint doms =
32
+
let new_doms = IntMap.mapi (fun id _ ->
33
+
if id = cfg.entry then init
34
+
else
35
+
let b = IntMap.find id cfg.blocks in
36
+
match b.preds with
37
+
| [] -> IntSet.empty
38
+
| p :: ps ->
39
+
let inter = List.fold_left (fun acc pr ->
40
+
try IntSet.inter acc (IntMap.find pr doms)
41
+
with Not_found -> acc
42
+
) (try IntMap.find p doms with Not_found -> IntSet.empty) ps in
43
+
IntSet.add id inter
44
+
) cfg.blocks in
45
+
if IntMap.equal IntSet.equal doms new_doms then doms
46
+
else fixpoint new_doms
47
+
in
48
+
fixpoint (IntMap.map (fun _ -> all) cfg.blocks)
49
+
50
+
let dominates doms a b =
51
+
IntSet.mem a (IntMap.find b doms)
52
+
53
+
let dominance_frontiers cfg doms =
54
+
let df = ref IntMap.empty in
55
+
IntMap.iter (fun id _ -> df := IntMap.add id IntSet.empty !df) cfg.blocks;
56
+
IntMap.iter (fun id b ->
57
+
match b.preds with
58
+
| [] | [_] -> ()
59
+
| _ ->
60
+
let idom_y = idom cfg doms id in
61
+
List.iter (fun p ->
62
+
(* Runner p is a predecessor of Y (p <> Y). Add Y to DF at each node
63
+
from p up to (but not including) Y's idom. p itself is included. *)
64
+
let rec walk cur =
65
+
if Some cur = idom_y then () (* reached idom(Y): stop, don't add *)
66
+
else begin
67
+
df := IntMap.add cur (IntSet.add id (IntMap.find cur !df)) !df;
68
+
(match idom cfg doms cur with
69
+
| None -> ()
70
+
| Some n -> walk n)
71
+
end
72
+
in
73
+
walk p
74
+
) b.preds
75
+
) cfg.blocks;
76
+
!df
+27
lib/ir.ml
+27
lib/ir.ml
···
1
+
type value = int
2
+
3
+
type inst =
4
+
| Alloca of value
5
+
| Load of value * value
6
+
| Store of value * value
7
+
| Phi of value * (value * value) list
8
+
| Move of value * value
9
+
| BinOp of value * value * value
10
+
| Ret of value option
11
+
| Br of value
12
+
| CondBr of value * value * value
13
+
14
+
type block = {
15
+
id : value;
16
+
insts : inst list;
17
+
mutable preds : value list;
18
+
mutable succs : value list;
19
+
}
20
+
21
+
type func = {
22
+
blocks : block list;
23
+
entry : value;
24
+
}
25
+
26
+
let mk_block id insts =
27
+
{ id; insts; preds = []; succs = [] }
+171
lib/nemo.ml
+171
lib/nemo.ml
···
1
+
open Ir
2
+
open Cfg
3
+
open Dominance
4
+
5
+
module PairMap = Map.Make(struct
6
+
type t = int * int
7
+
let compare (a1, b1) (a2, b2) =
8
+
match Stdlib.compare a1 a2 with
9
+
| 0 -> Stdlib.compare b1 b2
10
+
| c -> c
11
+
end)
12
+
13
+
type alloca_info = {
14
+
aid : value;
15
+
mutable stores : (value * value) list;
16
+
mutable loads : (value * value) list;
17
+
}
18
+
19
+
let find_allocas cfg =
20
+
let info = ref IntMap.empty in
21
+
IntMap.iter (fun _ b ->
22
+
List.iter (function
23
+
| Alloca a -> info := IntMap.add a { aid = a; stores = []; loads = [] } !info
24
+
| _ -> ()
25
+
) b.insts
26
+
) cfg.blocks;
27
+
IntMap.iter (fun _ b ->
28
+
List.iter (function
29
+
| Store (v, p) ->
30
+
(try let ai = IntMap.find p !info in ai.stores <- (v, b.id) :: ai.stores with Not_found -> ())
31
+
| Load (r, p) ->
32
+
(try let ai = IntMap.find p !info in ai.loads <- (r, b.id) :: ai.loads with Not_found -> ())
33
+
| _ -> ()
34
+
) b.insts
35
+
) cfg.blocks;
36
+
IntMap.fold (fun _ ai acc -> ai :: acc) !info []
37
+
38
+
let place_phi cfg df allocas =
39
+
let phi_needed = ref PairMap.empty in
40
+
List.iter (fun ai ->
41
+
let blocks_with_store = ref (List.fold_left (fun s (_, b) -> IntSet.add b s) IntSet.empty ai.stores) in
42
+
let work = ref (IntSet.elements !blocks_with_store) in
43
+
while !work <> [] do
44
+
let b = List.hd !work in
45
+
work := List.tl !work;
46
+
(try
47
+
let frontiers = IntMap.find b df in
48
+
IntSet.iter (fun f ->
49
+
let key = (f, ai.aid) in
50
+
if not (PairMap.mem key !phi_needed) then begin
51
+
phi_needed := PairMap.add key true !phi_needed;
52
+
if not (IntSet.mem f !blocks_with_store) then begin
53
+
blocks_with_store := IntSet.add f !blocks_with_store;
54
+
work := f :: !work
55
+
end
56
+
end
57
+
) frontiers
58
+
with Not_found -> ())
59
+
done
60
+
) allocas;
61
+
!phi_needed
62
+
63
+
let rename cfg phi_needed allocas =
64
+
let stacks = ref IntMap.empty in
65
+
let push aid v =
66
+
let s = try IntMap.find aid !stacks with Not_found -> [] in
67
+
stacks := IntMap.add aid (v :: s) !stacks
68
+
in
69
+
let pop aid =
70
+
let s = IntMap.find aid !stacks in
71
+
if List.tl s = [] then stacks := IntMap.remove aid !stacks
72
+
else stacks := IntMap.add aid (List.tl s) !stacks
73
+
in
74
+
let top aid =
75
+
try Some (List.hd (IntMap.find aid !stacks))
76
+
with Not_found -> None
77
+
in
78
+
(* For each (block, alloca) needing a phi, accumulate operands as the DFS
79
+
visits each predecessor. Each entry maps a predecessor id to the SSA
80
+
version of the variable on that edge. These are read only AFTER the full
81
+
DFS completes, so every predecessor has been visited. *)
82
+
let phi_ops = ref (PairMap.map (fun _ -> ref []) phi_needed) in
83
+
let body_insts = ref IntMap.empty in
84
+
let process_block id b =
85
+
(* A block that needs a phi defines a new version of the variable (using
86
+
the alloca id as the SSA name). Push it so subsequent loads in this
87
+
block read the phi's result rather than a stale predecessor version. *)
88
+
let pushed = ref [] in
89
+
List.iter (fun ai ->
90
+
if PairMap.mem (id, ai.aid) phi_needed then (push ai.aid ai.aid; pushed := ai.aid :: !pushed)
91
+
) allocas;
92
+
let insts = ref [] in
93
+
List.iter (fun i ->
94
+
match i with
95
+
| Alloca a -> () (* promoted away *)
96
+
| Store (v, p) ->
97
+
(* A store defines a new version of the variable. Push it once. *)
98
+
push p v; pushed := p :: !pushed
99
+
| Load (r, p) ->
100
+
(match top p with
101
+
| Some v -> insts := Move (r, v) :: !insts
102
+
| None -> ())
103
+
| Phi _ -> () (* input phis are replaced by the inserted ones *)
104
+
| _ -> insts := i :: !insts
105
+
) b.insts;
106
+
(* Rev to restore original order (insts was built by consing). *)
107
+
body_insts := IntMap.add id (List.rev !insts) !body_insts;
108
+
!pushed
109
+
in
110
+
let add_phi_operand block aid pred v =
111
+
if PairMap.mem (block, aid) phi_needed then
112
+
let ops_ref = PairMap.find (block, aid) !phi_ops in
113
+
(* Store as (pred, value) so IntMap.of_list below builds a pred->value map. *)
114
+
ops_ref := (pred, v) :: !ops_ref
115
+
in
116
+
let rec dfs id visited =
117
+
if IntSet.mem id visited then visited
118
+
else
119
+
let visited = IntSet.add id visited in
120
+
let b = IntMap.find id cfg.blocks in
121
+
let pushed = process_block id b in
122
+
let visited = List.fold_left (fun vis s ->
123
+
(* Edge id -> s: the version of each variable live on entry to s is the
124
+
current top of its stack. Record it as s's phi operand from id,
125
+
BEFORE descending into s (so s's own phi push doesn't interfere). *)
126
+
List.iter (fun ai ->
127
+
match top ai.aid with
128
+
| Some v -> add_phi_operand s ai.aid id v
129
+
| None -> ()
130
+
) allocas;
131
+
dfs s vis
132
+
) visited b.succs in
133
+
(* Pop every version this block pushed before backtracking. *)
134
+
List.iter pop pushed;
135
+
visited
136
+
in
137
+
ignore (dfs cfg.entry IntSet.empty);
138
+
(* Now build phi instructions from the fully-collected operands. *)
139
+
let phi_insts = ref IntMap.empty in
140
+
PairMap.iter (fun (block, aid) ops_ref ->
141
+
let ops_map = IntMap.of_list !ops_ref in
142
+
let b = IntMap.find block cfg.blocks in
143
+
let phi_ops_list = List.map (fun p ->
144
+
match IntMap.find_opt p ops_map with
145
+
| Some v -> (v, p)
146
+
| None -> (0, p)
147
+
) b.preds in
148
+
let existing = try IntMap.find block !phi_insts with Not_found -> [] in
149
+
phi_insts := IntMap.add block (Phi (aid, phi_ops_list) :: existing) !phi_insts
150
+
) !phi_ops;
151
+
(* Final per-block instruction list: phis first, then the rewritten body. *)
152
+
let new_insts = ref IntMap.empty in
153
+
IntMap.iter (fun id body ->
154
+
let phis = try IntMap.find id !phi_insts with Not_found -> [] in
155
+
new_insts := IntMap.add id (phis @ body) !new_insts
156
+
) !body_insts;
157
+
!new_insts
158
+
159
+
let mem2reg f =
160
+
let cfg = build_cfg f in
161
+
let doms = dominators cfg in
162
+
let df = dominance_frontiers cfg doms in
163
+
let allocas = find_allocas cfg in
164
+
let phi_needed = place_phi cfg df allocas in
165
+
let new_insts = rename cfg phi_needed allocas in
166
+
let new_blocks = IntMap.mapi (fun id b ->
167
+
{ b with insts = try IntMap.find id new_insts with Not_found -> b.insts }
168
+
) cfg.blocks in
169
+
(* Preserve the original input block order for deterministic output. *)
170
+
let ordered = List.map (fun b -> IntMap.find b.id new_blocks) f.blocks in
171
+
{ f with blocks = ordered }
+91
lib/parser.ml
+91
lib/parser.ml
···
1
+
open Ir
2
+
3
+
let read_lines path =
4
+
let ic = open_in path in
5
+
let lines = ref [] in
6
+
(try
7
+
while true do lines := input_line ic :: !lines done
8
+
with End_of_file -> close_in ic);
9
+
List.rev !lines
10
+
11
+
let parse_func (lines : string list) : func =
12
+
let blocks = ref [] in
13
+
let current_id = ref 0 in
14
+
let current_insts = ref [] in
15
+
let finish_block () =
16
+
if !current_id <> 0 then
17
+
blocks := mk_block !current_id (List.rev !current_insts) :: !blocks;
18
+
current_id := 0;
19
+
current_insts := []
20
+
in
21
+
let parse_instr s =
22
+
(* Tokenize: split on whitespace and standalone commas, but keep a
23
+
parenthesized group (e.g. "(10,2)") as a single token. This lets
24
+
"load 200, 100" and "phi 7 (10,2) (20,3)" both parse cleanly. *)
25
+
let tokens = ref [] in
26
+
let buf = Buffer.create 16 in
27
+
let depth = ref 0 in
28
+
let flush () =
29
+
if Buffer.length buf > 0 then begin
30
+
tokens := String.trim (Buffer.contents buf) :: !tokens;
31
+
Buffer.clear buf
32
+
end
33
+
in
34
+
String.iter (fun c ->
35
+
match c with
36
+
| '(' -> incr depth; Buffer.add_char buf c
37
+
| ')' -> decr depth; Buffer.add_char buf c
38
+
| (' ' | '\t' | '\n' | ',') when !depth = 0 ->
39
+
flush ()
40
+
| _ -> Buffer.add_char buf c
41
+
) s;
42
+
flush ();
43
+
let parts = List.rev !tokens in
44
+
match parts with
45
+
| "alloca" :: id :: [] -> Alloca (int_of_string id)
46
+
| "load" :: r :: p :: [] -> Load (int_of_string r, int_of_string p)
47
+
| "store" :: v :: p :: [] -> Store (int_of_string v, int_of_string p)
48
+
| "move" :: r :: v :: [] -> Move (int_of_string r, int_of_string v)
49
+
| "binop" :: r :: v1 :: v2 :: [] -> BinOp (int_of_string r, int_of_string v1, int_of_string v2)
50
+
| "phi" :: r :: ops ->
51
+
(* Each op is a parenthesized "(v,b)"; strip parens and split on ','. *)
52
+
let parse_op op =
53
+
(* Drop the surrounding parentheses, then split on ','. *)
54
+
let inner = String.sub op 1 (String.length op - 2) in
55
+
match String.split_on_char ',' inner with
56
+
| [v; b] -> (int_of_string (String.trim v), int_of_string (String.trim b))
57
+
| _ -> failwith ("Bad phi operand: " ^ op)
58
+
in
59
+
Phi (int_of_string r, List.map parse_op ops)
60
+
| "ret" :: [] -> Ret None
61
+
| "ret" :: v :: [] -> Ret (Some (int_of_string v))
62
+
| "br" :: t :: [] -> Br (int_of_string t)
63
+
| "condbr" :: c :: t1 :: t2 :: [] -> CondBr (int_of_string c, int_of_string t1, int_of_string t2)
64
+
| _ -> failwith ("Unknown instr: " ^ s)
65
+
in
66
+
List.iter (fun line ->
67
+
let trimmed = String.trim line in
68
+
if trimmed = "" then ()
69
+
else if String.ends_with ~suffix:":" trimmed then
70
+
let label = String.sub trimmed 0 (String.length trimmed - 1) in
71
+
(* Accept either "1" or "block 1" style headers. *)
72
+
let id_str =
73
+
if String.starts_with ~prefix:"block " label
74
+
then String.sub label 6 (String.length label - 6)
75
+
else label
76
+
in
77
+
finish_block ();
78
+
current_id := int_of_string id_str
79
+
else
80
+
current_insts := parse_instr trimmed :: !current_insts
81
+
) lines;
82
+
finish_block ();
83
+
let blocks = List.rev !blocks in
84
+
let entry = match blocks with
85
+
| [] -> 0
86
+
| b :: _ -> b.id
87
+
in
88
+
{ blocks; entry }
89
+
90
+
let parse_func_from_file path : func =
91
+
parse_func (read_lines path)
+22
lib/printer.ml
+22
lib/printer.ml
···
1
+
open Ir
2
+
3
+
let string_of_inst = function
4
+
| Alloca a -> Printf.sprintf "alloca %d" a
5
+
| Load (r, p) -> Printf.sprintf "load %d, %d" r p
6
+
| Store (v, p) -> Printf.sprintf "store %d, %d" v p
7
+
| Phi (r, ops) ->
8
+
let ops_str = List.map (fun (v, b) -> Printf.sprintf "(%d, %d)" v b) ops |> String.concat ", " in
9
+
Printf.sprintf "phi %d [%s]" r ops_str
10
+
| Move (r, v) -> Printf.sprintf "move %d, %d" r v
11
+
| BinOp (r, v1, v2) -> Printf.sprintf "binop %d, %d, %d" r v1 v2
12
+
| Ret None -> "ret"
13
+
| Ret (Some v) -> Printf.sprintf "ret %d" v
14
+
| Br t -> Printf.sprintf "br %d" t
15
+
| CondBr (c, t1, t2) -> Printf.sprintf "condbr %d, %d, %d" c t1 t2
16
+
17
+
let print_block b =
18
+
Printf.printf "block %d:\n" b.id;
19
+
List.iter (fun i -> Printf.printf " %s\n" (string_of_inst i)) b.insts
20
+
21
+
let print_func f =
22
+
List.iter print_block f.blocks
+17
main.ml
+17
main.ml
···
1
+
open Ir
2
+
open Parser
3
+
open Nemo
4
+
open Printer
5
+
6
+
let () =
7
+
if Array.length Sys.argv < 2 then begin
8
+
Printf.eprintf "Usage: %s <file.ir>\n" Sys.argv.(0);
9
+
exit 1
10
+
end;
11
+
let path = Sys.argv.(1) in
12
+
let f = parse_func_from_file path in
13
+
let f' = mem2reg f in
14
+
Printf.printf "Before (mem2reg input):\n";
15
+
print_func f;
16
+
Printf.printf "\nAfter (SSA):\n";
17
+
print_func f'
+86
tests/test.ml
+86
tests/test.ml
···
1
+
open Ir
2
+
open Nemo
3
+
open Printer
4
+
open Parser
5
+
6
+
(* ---- Inline unit tests (no file I/O) ---- *)
7
+
8
+
let test_simple () =
9
+
Printf.printf "Test Simple\n";
10
+
let b1 = mk_block 1 [Alloca 100; Store (1, 100); Br 2] in
11
+
let b2 = mk_block 2 [Load (200, 100); Ret (Some 200)] in
12
+
let f = { blocks = [b1; b2]; entry = 1 } in
13
+
Printf.printf "Before:\n";
14
+
print_func f;
15
+
let f' = mem2reg f in
16
+
Printf.printf "\nAfter:\n";
17
+
print_func f'
18
+
19
+
let test_phi () =
20
+
Printf.printf "\nTest Phi\n";
21
+
let b1 = mk_block 1 [Alloca 100; Store (1, 100); Br 2] in
22
+
let b2 = mk_block 2 [Alloca 101; Store (2, 101); Br 3] in
23
+
let b3 = mk_block 3 [Load (200, 100); Load (201, 101); Ret (Some 200)] in
24
+
let f = { blocks = [b1; b2; b3]; entry = 1 } in
25
+
Printf.printf "Before:\n";
26
+
print_func f;
27
+
let f' = mem2reg f in
28
+
Printf.printf "\nAfter:\n";
29
+
print_func f'
30
+
31
+
let test_branch () =
32
+
Printf.printf "\nTest Branch with Phi\n";
33
+
let b1 = mk_block 1 [Alloca 100; CondBr (1, 2, 3)] in
34
+
let b2 = mk_block 2 [Store (10, 100); Br 4] in
35
+
let b3 = mk_block 3 [Store (20, 100); Br 4] in
36
+
let b4 = mk_block 4 [Load (200, 100); Ret (Some 200)] in
37
+
let f = { blocks = [b1; b2; b3; b4]; entry = 1 } in
38
+
Printf.printf "Before:\n";
39
+
print_func f;
40
+
let f' = mem2reg f in
41
+
Printf.printf "\nAfter:\n";
42
+
print_func f'
43
+
44
+
(* ---- File-based regression test ---- *)
45
+
46
+
(* Parse a .ir file, run mem2reg, and assert the printed SSA output equals the
47
+
expected string. Exits with status 1 on mismatch so `make test` fails. *)
48
+
let string_of_func f =
49
+
let buf = Buffer.create 256 in
50
+
List.iter (fun b ->
51
+
Buffer.add_string buf (Printf.sprintf "block %d:\n" b.id);
52
+
List.iter (fun i ->
53
+
Buffer.add_string buf (Printf.sprintf " %s\n" (Printer.string_of_inst i))
54
+
) b.insts
55
+
) f.blocks;
56
+
Buffer.contents buf
57
+
58
+
let test_file name expected () =
59
+
Printf.printf "\nTest file: %s\n" name;
60
+
let f = parse_func_from_file ("examples/" ^ name) in
61
+
let f' = mem2reg f in
62
+
let got = string_of_func f' in
63
+
if got <> expected then begin
64
+
Printf.eprintf "FAIL %s\n--- expected ---\n%s\n--- got ---\n%s\n" name expected got;
65
+
exit 1
66
+
end;
67
+
print_string got
68
+
69
+
let expected_simple =
70
+
"block 1:\n br 2\nblock 2:\n move 200, 1\n ret 200\n"
71
+
72
+
let expected_basic =
73
+
"block 1:\n br 2\nblock 2:\n move 200, 42\n ret 200\n"
74
+
75
+
(* Example exercising the phi-handling path: an input phi node is filled in
76
+
from the renamed value at each predecessor. *)
77
+
let expected_phi =
78
+
"block 1:\n condbr 1, 2, 3\nblock 2:\n br 4\nblock 3:\n br 4\nblock 4:\n phi 100 [(20, 3), (10, 2)]\n move 200, 100\n ret 200\n"
79
+
80
+
let () =
81
+
test_simple ();
82
+
test_phi ();
83
+
test_branch ();
84
+
test_file "simple.ir" expected_simple ();
85
+
test_file "basic.ir" expected_basic ();
86
+
test_file "phi.ir" expected_phi ()