(*
    use "set.ml";;
    use "imp-syntax";;
*)

type blabel = string;;
type block = BAssignVar of ide * exp
    | BAssignArray of ide * exp * exp
    | BSkip
    | BGuard of exp;;

type cfg = (blabel * block) set        (* nodes *)
         * (blabel * blabel) set;;     (* edges *)

type s =
    SSkip of blabel
    | SAssignVar of blabel * ide * exp
    | SAssignArray of blabel * ide * exp * exp
    | SCseq of s * s
    | SIf of (blabel * exp) * s * s
    | SWhile of (blabel * exp) * s;;

let rec get_last_label n : blabel = match n with
    SSkip l -> l
    | SAssignVar(l, _, _) -> l
    | SAssignArray(l, _, _, _) -> l
    | SCseq(_, s) -> get_last_label s
    | SIf(_, _, s) -> get_last_label s
    | SWhile(_, s) -> get_last_label s;;

let labeledprog imp : s = match imp with
    ImpProg(_, c) ->
        let rec labeledprog' c i = match c with
            Skip -> SSkip(string_of_int i)
            | AssignVar(n, e) -> SAssignVar(string_of_int i, n, e)
            | AssignArray(n, p, e) -> SAssignArray(string_of_int i, n, p, e)
            | Cseq(c1, c2) ->
                let l' = (labeledprog' c1 i) in
                let i' = int_of_string (get_last_label l') in
                SCseq(l', (labeledprog' c2 (i' + 1)))
            | If(e, c1, c2) ->
                let l' = (labeledprog' c1 (i + 1)) in
                let i' = int_of_string (get_last_label l') in
                SIf((string_of_int i, e), l', (labeledprog' c2 (i' + 1)))
            | While(e, c) -> SWhile((string_of_int i, e), (labeledprog' c (i + 1)))
    in labeledprog' c 1;;

let rec init s : blabel = match s with
        SSkip l -> l
        | SAssignVar(l, _, _) -> l
        | SAssignArray(l, _, _, _) -> l
        | SCseq(s1, _) -> init s1
        | SIf((l, _), _, _) -> l
        | SWhile((l, _), _) -> l;;

let rec final s : blabel set = match s with
    SSkip l -> [l]
    | SAssignVar(l, _, _) -> [l]
    | SAssignArray(l, _, _, _) -> [l]
    | SCseq(_, s2) -> final s2
    | SIf(_, s1, s2) -> final s1 @ final s2
    | SWhile((l, _), _) -> [l];;

let rec labels s : blabel set = match s with
    SSkip l -> [l]
    | SAssignVar(l, _, _) -> [l]
    | SAssignArray(l, _, _, _) -> [l]
    | SCseq(s1, s2) -> labels s1 @ labels s2
    | SIf((l, _), s1, s2) -> [l] @ labels s1 @ labels s2
    | SWhile((l, _), s) -> [l] @ labels s;;

let rec flow s : (blabel * blabel) set = match s with
    SSkip l -> []
    | SAssignVar _ -> []
    | SAssignArray _ -> []
    | SCseq(s1, s2) ->
        flow s1 @ flow s2 @ [(List.hd (final s1), init s2)]
    | SIf((l, _), s1, s2) ->
        flow s1 @ flow s2 @ [init s1, init s2]
    | SWhile((l, _), s) ->
        flow s @ [(l, init s)] @ [List.hd (final s), l];;

let flow_rev s : (blabel * blabel) set =
    let rec flow_rev' f = match f with
        [] -> []
        | (l, l')::f -> [(l', l)] @ flow_rev' f
    in flow_rev' (flow s);;

let rec blocks s : (blabel * block) set = match s with
    SSkip l -> [l, BSkip]
    | SAssignVar(l, n, e) -> [l, BAssignVar(n, e)]
    | SAssignArray(l, n, p, e) -> [(l, BAssignArray(n, p, e))]
    | SCseq(s1, s2) -> blocks s1 @ blocks s2
    | SIf((l, e), s1, s2) -> [l, BGuard e] @ blocks s1 @ blocks s2
    | SWhile((l, e), s) -> [l, BGuard e] @ blocks s;;

let cfg_of_imp imp : cfg =
    let s = labeledprog imp in
    (blocks s, flow s);;