open List;;
type prog = Prog of stmt
and stmt = CompStmt of (stmt * stmt) |
            AssignStmt of (string * exp) |
            PrintStmt of exp |
            IfStmt of (exp * stmt * stmt) |
            WhileStmt of (exp * stmt)
and exp = Number of int |
            Id of string |
            PlusOp of (exp * exp) |
            MinusOp of (exp * exp) |
            MulOp of (exp * exp) |
            DivOp of (exp * exp)
and node = Node of stmt
and pred = node
and succ = node
and edge = pred * succ * int
and nodeList = node list
and edgeList = edge list
and cfg = nodeList * edgeList * node * nodeList;; (* node list, edge list, first node, last nodes*)

let ast =
CompStmt(AssignStmt("a",PlusOp(Id("a"),Id("b"))),
    CompStmt(AssignStmt("b",MinusOp(Id("a"),Id("b"))),
        CompStmt(AssignStmt("a",MinusOp(Id("a"),Id("b"))), WhileStmt(Id("c"),CompStmt(AssignStmt("x",MulOp(Id("x"),Id("y"))),
                                                                                AssignStmt("y",DivOp(Id("y"),Number(2))))))));;
let rec addToNodeList (listIn: node list) (el: node)
= match listIn with
| [] -> [el]
| h::t -> h::(addToNodeList t el);;

let rec addToEdgeList (listIn: edge list) (el: edge)
= match listIn with
| [] -> [el]
| h::t -> h::(addToEdgeList t el);;

let rec addEdgeListToEdgeList (edgeList : edge list) (newEdgeList : edge list) =
    match edgeList with
    | [] -> newEdgeList
    | h::t -> h::(addEdgeListToEdgeList t newEdgeList)
(* let rec getLastSucNode (edgeList: edge list) : node list =
    match edgeList with
    | [pred, succ, int] -> succ
    | h::t -> getLastSucNode t;; *)

let rec getDefList (nodeList: node list) : string list =
begin
    match nodeList with
    | [] -> ["_"]
    | Node(h)::t ->
          begin
              match h with
              | CompStmt(st1, st2) -> ["ERROR"]
              | AssignStmt(var, exp) -> var::(getDefList t)
              | PrintStmt(exp) -> "_"::(getDefList t)
              | IfStmt(exp, stmt1, strm2) -> "_"::(getDefList t)
              | WhileStmt(exp, stmt) -> "_"::(getDefList t)
          end
end

let rec getSuccForNode (edgeList: edge list) (node: node) : node list =
begin
    match edgeList with
    | [] -> []
    | h::t ->
    begin
        match h with
        | (pred, succ, _) -> if pred == node then succ::(getSuccForNode t node)
                                else (getSuccForNode t node)
    end
end

let rec genNodes (transfAst: stmt) (nlist: node list) : node list =
begin
    match transfAst with
    | CompStmt (st1, st2) -> genNodes st2 (genNodes st1 nlist)
    | AssignStmt (var, exp) -> addToNodeList nlist (Node(AssignStmt(var, exp)))
    | PrintStmt(exp) -> addToNodeList nlist (Node (PrintStmt (exp)))
    | IfStmt (exp, stmt1, stmt2) -> genNodes stmt2 (genNodes stmt1 (addToNodeList nlist (Node (IfStmt(exp,stmt1,stmt2)))))
    | WhileStmt (exp, stmt) -> genNodes stmt (addToNodeList nlist (Node (WhileStmt(exp, stmt))))
end

let rec getLastCompNode (transfAst: stmt) : node =
begin
    match transfAst with
    | CompStmt(_,st2) -> getLastCompNode (st2)
    | AssignStmt (var,exp) -> Node(AssignStmt(var,exp))
    | PrintStmt(exp) -> Node(PrintStmt(exp))
    | IfStmt(exp,stmt1,stmt2) -> Node(IfStmt(exp,stmt1,stmt2))
    | WhileStmt(exp,stmt) -> Node(WhileStmt(exp,stmt))
end

let rec getFirstCompNode (transfAst: stmt) : node =
begin
    match transfAst with
    | CompStmt(st1,_) -> getLastCompNode (st1)
    | AssignStmt (var,exp) -> Node(AssignStmt(var,exp))
    | PrintStmt(exp) -> Node(PrintStmt(exp))
    | IfStmt(exp,stmt1,stmt2) -> Node(IfStmt(exp,stmt1,stmt2))
    | WhileStmt(exp,stmt) -> Node(WhileStmt(exp,stmt))
end

let rec genEdges (transfAst: stmt) (n: int) : edge list =
  begin
    match transfAst with
        | CompStmt (st1, st2) -> let firstList = (genEdges st1 (n+1)) in begin
                    (getLastCompNode(st1), getFirstCompNode(st2), n)::List.append firstList (genEdges st2 (n+(List.length firstList)+1))
                end
        | AssignStmt (var, exp) -> []
        | PrintStmt (exp) -> []
        | IfStmt (exp, st1, st2) -> let firstList = (genEdges st1 (n+2)) in begin
                                            (Node(IfStmt(exp,st1,st2)),getFirstCompNode(st1),n) ::
                                                ((Node(IfStmt(exp,st1,st2)),getFirstCompNode(st2),n+1) ::
                                                List.append firstList (genEdges st2 (n+2+(List.length firstList))))
                                    end
        | WhileStmt (exp, st) -> (Node(WhileStmt(exp,st)),getFirstCompNode(st),n) ::
                                    ((getLastCompNode(st),Node(WhileStmt(exp,st)),n+1) ::
                                        (genEdges st (n+2)))
  end

let rec print_exp (expr: exp) =
begin
  match expr with
    | Number(nr) -> Printf.printf "%d" nr
    | Id (id) -> Printf.printf "%s" id
    | PlusOp (exp1, exp2) -> print_exp exp1; Printf.printf " + "; print_exp exp2
    | MinusOp (exp1, exp2) -> print_exp exp1; Printf.printf " - "; print_exp exp2
    | MulOp (exp1, exp2) -> print_exp exp1; Printf.printf " * "; print_exp exp2
    | DivOp (exp1, exp2) -> print_exp exp1; Printf.printf " / "; print_exp exp2
end

let print_Stmt (stmt: stmt) =
begin
  match stmt with
    | CompStmt (st1, st2) -> Printf.printf "comp print error\n"
    | AssignStmt (var1, exp1) ->  Printf.printf "%s = " var1; print_exp exp1; Printf.printf "\n"
    | PrintStmt (exp1) ->  Printf.printf "Print stmt: "; print_exp exp1; Printf.printf "\n"
    | IfStmt (exp1, stm1, stm2) -> Printf.printf "it is if\n"
    | WhileStmt (exp1, stm1) ->  Printf.printf "it is while\n"
end

let rec print_NodeList = function
[] -> print_string "\n"
| (Node (stmt))::l -> print_Stmt stmt; print_NodeList l;;


let rec print_StringList = function
[] -> print_string "\n"
| str::l -> print_string str; print_StringList   l;;

let rec print_EdgeList = function
[] -> print_string "\n"
| (Node(pred), Node(succ), n)::l -> Printf.printf "Edge nr = %d\n" n; Printf.printf "Pred list:\n"; print_Stmt pred;
  Printf.printf "Succ list:\n"; print_Stmt succ; print_EdgeList l;;

let rec getFirstNode (transfAst: stmt) =
  begin
    match transfAst with
        | CompStmt (st1, _) -> getFirstNode (st1)
        | AssignStmt (var, exp) -> Node (AssignStmt (var, exp))
        | PrintStmt (exp) -> Node (PrintStmt (exp))
        | IfStmt (exp, st1, st2) -> Node (IfStmt (exp, st1, st2))
        | WhileStmt (exp, st) -> Node (WhileStmt (exp, st))
  end

let rec genLastNodes (transfAst: stmt) (lastNodes: node list)  =
begin
  match transfAst with
  | CompStmt (_, st) -> genLastNodes st lastNodes
  | AssignStmt (var, exp) -> addToNodeList lastNodes (Node (AssignStmt (var, exp)))
  | PrintStmt (exp) -> addToNodeList lastNodes (Node (PrintStmt (exp)))
  | IfStmt (_, st1, st2) -> genLastNodes st2 (genLastNodes st1 lastNodes)
  | WhileStmt (exp, st) -> addToNodeList [getLastCompNode(st)] (Node (WhileStmt (exp, st)))
end

let print_controlFlowG (cfg: cfg) =
  match cfg with
  | (nlist, elist, (Node (stmt)), lastn) -> Printf.printf "First node: ";print_Stmt (stmt); Printf.printf "Node list:\n";
   print_NodeList nlist; Printf.printf "Edge list:\n"; print_EdgeList elist; Printf.printf "Last nodes list:\n";
   print_NodeList lastn;;

let transformAst (transfAst: stmt) =
  (genNodes transfAst [], genEdges transfAst 1, getFirstNode transfAst, genLastNodes transfAst [])

let rez = transformAst ast;;
  (* print_controlFlowG rez;; *)
let asd = getDefList (genNodes ast []);;
print_StringList asd;;

(* let rez = genEdges ast 1;;
print_EdgeList rez;; *)



(* let rec genEdges (transfAst: stmt) (elist: edge list) (pred: node list) (succ: node list) (n: int) =
  begin
    match transfAst with
        | CompStmt (st1, st2) -> genEdges st2 (genEdges st1 elist pred (getCompSucc transfAst [] 0) n) (getCompPred transfAst [] 0) succ (n+1)
        | AssignStmt (_, _) -> addToEdgeList elist (pred, succ, n)
        | PrintStmt (_) -> addToEdgeList elist (pred, succ, n)
        | IfStmt (_, st1, st2) -> genEdges st2 (genEdges st1 (addToEdgeList elist (pred, succ, n)) pred succ (n+1)) pred succ n
        | WhileStmt (exp, st) -> genEdges st (addToEdgeList elist ((getPredForWhile st pred), (getSuccForWhile st succ), n)) [(Node (WhileStmt (exp, st)))] [(Node (WhileStmt (exp, st)))] (n+1)
  end *)
(* let rec getCompPred (transfAst: stmt) (predList: node list) (firstCall: int) =
begin
    match transfAst with
    | CompStmt (st1, st2) ->
        if firstCall = 0 then
          (match st1 with
          | CompStmt (_, _) -> getCompPred st2 predList 0
          | AssignStmt (_, _) -> getCompPred st1 predList 1
          | PrintStmt (_) -> getCompPred st1 predList 1
          | IfStmt (_, st3, st4) -> getCompPred st4 (getCompPred st3 predList 1) 1
          | WhileStmt (_, st3) -> getCompPred st3 predList 1
        ) else (getCompPred st1 predList 1)
    | AssignStmt (var, exp) -> addToNodeList predList (Node (AssignStmt (var, exp)))
    | PrintStmt (exp) -> addToNodeList predList (Node (PrintStmt (exp)))
    | IfStmt (exp, st1, st2) -> addToNodeList predList (Node (IfStmt (exp, st1, st2)))
    | WhileStmt (exp, st) -> addToNodeList predList (Node (WhileStmt (exp, st)))
end


let rec getCompSucc (transfAst: stmt) (succList: node list) (firstCall: int) =
begin
    match transfAst with
    | CompStmt (st1, st2) ->
        if firstCall = 0 then (
          match st1 with
          | CompStmt (_, _) -> getCompSucc st1 succList 0
          | AssignStmt (_, _) -> getCompSucc st2 succList 1
          | PrintStmt (_) -> getCompSucc st2 succList 1
          | IfStmt (_, st3, st4) -> getCompSucc st4 (getCompSucc st3 succList 1) 1
          | WhileStmt (_, st3) -> getCompSucc st3 (getCompSucc st2 succList 1) 1
        ) else (getCompSucc st1 succList 1)
    | AssignStmt (var, exp) -> addToNodeList succList (Node (AssignStmt (var, exp)))
    | PrintStmt (exp) -> addToNodeList succList (Node (PrintStmt (exp)))
    | IfStmt (exp, st1, st2) -> addToNodeList succList (Node (IfStmt (exp, st1, st2)))
    | WhileStmt (exp, st) -> addToNodeList succList (Node (WhileStmt (exp, st)))
end *)

(* let rec getPredForWhile (transfAst: stmt) (predList: node list) =
begin
    match transfAst with
    | CompStmt (_, st2) -> getPredForWhile st2 predList
    | AssignStmt (var, exp) -> addToNodeList predList (Node (AssignStmt (var, exp)))
    | PrintStmt (exp1) -> addToNodeList predList (Node (PrintStmt (exp1)))
    | IfStmt (_, st1, st2) -> getPredForWhile st1 (getPredForWhile st2 predList)
    | WhileStmt (exp, st) -> addToNodeList predList (Node (WhileStmt (exp, st)))
end

let rec getSuccForWhile (transfAst: stmt) (succList: node list) =
begin
  if succList = [] then
    (match transfAst with
        | CompStmt (st1, _) -> getSuccForWhile st1 succList
        | AssignStmt (var, exp) -> addToNodeList succList (Node (AssignStmt (var, exp)))
        | PrintStmt (exp) -> addToNodeList succList (Node (PrintStmt (exp)))
        | IfStmt (exp, st1, st2) -> addToNodeList succList (Node (IfStmt (exp, st1, st2)))
        | WhileStmt (exp, st) -> addToNodeList succList (Node (WhileStmt (exp, st)))
    )
  else succList
end *)