open Ast
open Util

(* Utility : find declaration of a const / a node *)

let find_const_decl p id =
  match List.find (function
                   | AST_const_decl (c, _) when c.c_name = id -> true
                   | _ -> false) 
                  p with
  | AST_const_decl d -> d
  | _ -> assert false

let find_node_decl p id =
  match List.find (function
                   | AST_node_decl (c, _) when c.n_name = id -> true
                   | _ -> false) 
                  p with
  | AST_node_decl d -> d
  | _ -> assert false

let extract_const_decls =
  List.fold_left
    (fun l d -> match d with
      | AST_const_decl d -> d::l
      | _ -> l)
    []


(* Utility : scopes *)


(* path to node * subscope prefix in node * equations in scope *)
type scope = id * id * eqn ext list

let get_root_scope p root =
    let (n, _) = find_node_decl p root in
    ("", "", n.body)


(* Utility : find instances declared in an expression *)

(* extract_instances :
      prog -> expr ext -> (id * eqs * (var_def * expr ext) list) list
*)
let rec extract_instances p e = match fst e with
  | AST_idconst _ | AST_identifier _
  | AST_int_const _ | AST_real_const _ | AST_bool_const _ -> []
  | AST_unary (_, e') | AST_pre (e', _) | AST_not e' -> extract_instances p e'
  | AST_binary(_, e1, e2) | AST_binary_rel (_, e1, e2) | AST_binary_bool (_, e1, e2)
  | AST_arrow(e1, e2) ->
    extract_instances p e1 @ extract_instances p e2
  | AST_if(e1, e2, e3) ->
    extract_instances p e1 @ extract_instances p e2 @ extract_instances p e3
  | AST_instance((f, _), args, id) ->
    let more = List.flatten (List.map (extract_instances p) args) in
    let (node, _) = find_node_decl p f in
    let args_x = List.map2 (fun id arg -> id, arg) node.args args in
    (id, node.body, args_x)::more

(* Utility : find pre declarations in an expression *)

(* extract_pre : expr ext -> (id * expr ext) list *)
let rec extract_pre e = match fst e with
  | AST_identifier _ | AST_idconst _
  | AST_int_const _ | AST_real_const _ | AST_bool_const _ -> []
  | AST_unary (_, e') | AST_not e' -> extract_pre e'
  | AST_binary(_, e1, e2) | AST_binary_rel (_, e1, e2) | AST_binary_bool (_, e1, e2)
  | AST_arrow(e1, e2) ->
    extract_pre e1 @ extract_pre e2
  | AST_if(e1, e2, e3) ->
    extract_pre e1 @ extract_pre e2 @ extract_pre e3
  | AST_instance((f, _), args, id) ->
    List.flatten (List.map extract_pre args)
  | AST_pre(e', n) ->
    (n, e')::(extract_pre e')

(* Some errors *)

let combinatorial_cycle v = error ("Combinatorial cycle with variable: " ^ v)
let no_variable e = error ("No such variable: " ^ e)
let type_error e = error ("Type error: " ^ e)
let not_implemented e = error ("Not implemented: " ^ e)
let invalid_arity e = error ("Invalid arity (" ^ e ^ ")")