(* Scope analyzer *)

open Ast
open Util
open Ast_util

type renaming = (id * bool) VarMap.t

(* consts_renaming : prog -> renaming *)
let consts_renaming p =
  let cdecl = extract_const_decls p in
  List.fold_left
    (fun rn (d,_) -> VarMap.add d.c_name (d.c_name, true) rn)
    VarMap.empty
    cdecl

(* add_var_defs : string -> renaming -> var_def list -> renaming *)
let add_var_defs prefix =
  List.fold_left
    (fun rn (_,n,_) -> VarMap.add n (prefix^n, false) rn)

(* rename_expr : expr ext -> renaming -> expr ext *)
let rec rename_expr exp rn =
  let sub e = rename_expr e rn in
  let a = match fst exp with
    | AST_identifier(id, l) ->
      begin try let (nn, const) = VarMap.find id rn in
          if const then AST_idconst(nn, l) else AST_identifier(nn, l)
      with Not_found -> loc_error l no_variable id
      end
    | AST_idconst(id, l) -> assert false (* expression not already renamed *)
    | AST_unary(op, e) -> AST_unary(op, sub e)
    | AST_binary(op, a, b) -> AST_binary(op, sub a, sub b)
    | AST_binary_rel(op, a, b) -> AST_binary_rel(op, sub a, sub b)
    | AST_binary_bool(op, a, b) -> AST_binary_bool(op, sub a, sub b)
    | AST_not(e) -> AST_not(sub e)
    | AST_pre (e, x) -> AST_pre(sub e, x)
    | AST_arrow(a, b) -> AST_arrow(sub a, sub b)
    | AST_if(a, b, c) -> AST_if(sub a, sub b, sub c)
    | AST_instance(i, l, s) -> AST_instance(i, List.map sub l, s)
    | x -> x  (* other cases : constant literals *)
  in a, snd exp

(* rename_var : renaming -> id -> id *)
let rename_var rn id =
  try let (nn, const)= VarMap.find id rn in
      if const then error ("Constant "^id^" cannot be defined in node.");
      nn
  with Not_found -> no_variable id

let rename_var_ext rn (id, loc) = (rename_var rn id, loc)

let rename_var_def rn (p, id, t) = (p, rename_var rn id, t)

(* rename_eqn : renaming -> eqn ext -> eqn ext *)
let rec rename_eqn rn e =
  let a = match fst e with
    | AST_assign(s, e) -> AST_assign (List.map (rename_var_ext rn) s, rename_expr e rn)
    | AST_guarantee(g, e) -> AST_guarantee(g, rename_expr e rn)
    | AST_assume(a, e) -> AST_assume(a, rename_expr e rn)
    | AST_automaton(aid, states, ret) ->
      let rename_state ((st:Ast.state),l) =
        let rn = add_var_defs (aid^"."^st.st_name^".") rn st.st_locals in
        { st with
          st_locals = List.map (rename_var_def rn) st.st_locals;
          body = List.map (rename_eqn rn) st.body;
          until = List.map (fun (c, n, r) -> (rename_expr c rn, n, r)) st.until;
        }, l
      in 
      AST_automaton(aid, List.map rename_state states, List.map (rename_var rn) ret)
    | AST_activate (c, ret) ->
      let rec rename_activate_if = function
        | AST_activate_if(c, a, b) ->
          AST_activate_if(rename_expr c rn, rename_activate_if a, rename_activate_if b)
        | AST_activate_body b ->
          let rn = add_var_defs (b.act_id^".") rn b.act_locals in
          AST_activate_body{ b with
            act_locals = List.map (rename_var_def rn) b.act_locals;
            body = List.map (rename_eqn rn) b.body;
          }
      in
      AST_activate(rename_activate_if c, List.map (rename_var rn) ret)
  in a, snd e

(* rename_node : renaming -> node_decl -> node_decl *)
let rename_node const_rn node =
  let rn = add_var_defs "" const_rn node.args in
  let rn = add_var_defs "" rn node.ret in
  let rn = add_var_defs "" rn node.var in
  { node with
    body = List.map (rename_eqn rn) node.body }

(* rename_prog : prog -> prog *)
let rename_prog p =
  let const_rn = consts_renaming p in
  let rn_toplevel = function
  | AST_node_decl (n, l) -> AST_node_decl (rename_node const_rn n, l)
  | x -> x
  in
  List.map rn_toplevel p