open Support.Pervasive
open Support.Error
open Union_find
open Prim
open List

type 'a alist = (string * 'a) list

type mutability = Mutable | Constant

type value_kind = Lvalue | Rvalue


type where_clause = ((string * typ list) list (* concept constraints *)
		     * (typ * typ) list)      (* same type constraints *)

and exp_typ = typ * value_kind * mutability

(* The class_member type is for parsing only *)
and class_member =
    Field of info * string * typ
  | Constructor of 
      info
      * string list          (* template parameters *)
      * where_clause
      * param list           (* parameters *)
      * (string * exp list) list  (* member inits *)
      * stmt                 (* body *)
  | Destructor of info * stmt

(* The concept_member type is for parsing only *)
and concept_member =
   AssocC of info          (* require an associated type *)
      * string
  | RefinesC of info           (* refinement *)
      * string
      * typ list
  | RequiresC of info           (* requirement *)
      * string
      * typ list
  | SameTypeC of info          (* same type requirement *)
      * typ * typ
  | FunC of info
      * string               (* function name *)
      * string list          (* template parameters *)
      * where_clause
      * param list           (* parameters *)
      * exp_typ              (* return type *)
      * stmt option          (* body *)

(* The model_member type is for parsing only *)
and model_member =
    AssocM of info * string * typ
  | FunM of info
      * string               (* function name *)
      * bool                 (* allow recursion? *)
      * string list          (* template parameters *)
      * where_clause
      * param list           (* parameters *)
      * exp_typ              (* return type *)
      * stmt                 (* body *)

and param = string * exp_typ
      
and constructor_info =
    info
    * string list             (* template parameters *)
    * where_clause
    * param list            (* parameters *)
    * (string * exp list) list   (* member inits *)
    * stmt                  (* body *)

and function_info =
    info
    * bool                    (* allow recursion? *)
    * string list             (* template parameters *)
    * where_clause
    * param list          (* parameters *)
    * exp_typ             (* return type *)
    * stmt                (* body *)

and decl =
    LetD of info * typ * mutability * string * exp (* type, name, initializing expression *)
  | LetFwdD of info * typ * mutability * string
  | TypedefD of info * string * typ
  | FunD of info
      * string               (* function name *)
      * bool                 (* allow recursion? *)
      * string list          (* template parameters *)
      * where_clause
      * param list           (* parameters *)
      * exp_typ              (* return type *)
      * stmt                 (* body *)
  | ExternScopeD of info * string * decl list
  | FunFwdD of info
      * string               (* function name *)
      * string list          (* template parameters *)
      * where_clause
      * param list           (* parameters *)
      * exp_typ              (* return type *)
  | ExternFunFwdD of info
      * string               (* language: C, C++ *)
      * string               (* function name *)
      * string list          (* template parameters *)
      * where_clause
      * param list           (* parameters *)
      * exp_typ              (* return type *)
  | PrimD of info
      * primitive
      * string list
      * param list
      * exp_typ
  | OvldD of info
      * string               (* function name to overload *)
      * exp                  (* function to add to the overload set *)

  | ClassFwdD of info * string * string list * where_clause
  | StructFwdD of info * string * string list * where_clause
  | UnionFwdD of info * string * string list * where_clause
  | StructD of info
      * string               (* struct name *)
      * string list          (* template parameters *)
      * where_clause
      * (string * typ) list  (* field names and types *)
  | UnionD of info
      * string               (* union name *)
      * string list          (* template parameters *)
      * where_clause
      * (string * typ) list  (* field names and types *)
  | ClassD of info
      * string               (* class name *)
      * string list          (* type parameters *)
      * where_clause
      * (string * typ) list  (* fields *)
      * constructor_info list
      * stmt                 (* destructor body *)
  | ConceptD of info
      * string               (* concept name *)
      * string list          (* type parameters *)
      * string list          (* associated types *)
      * (string * typ list) list (* refinements *)
      * (string * typ list) list (* requires *)
      * (string               (* function name *)
	 * (typ                  (* function type *)
	      * string list      (* param names *)
	      * stmt option)) list    (* default body *)
      * (typ * typ) list     (* same-types *)
  | ModelD of info
      * bool                 (* open into the enclosing scope? *)
      * string list          (* type parameters *)
      * where_clause
      * string               (* concept name *)
      * typ list             (* type arguments *)
      * (string * typ) list  (* associated type definitions *)
      * (string * function_info) list

  | ModelFwdD of info
      * bool
      * string list          (* type parameters *)
      * where_clause
      * string               (* concept name *)
      * typ list             (* type arguments *)
      * (string * typ) list  (* associated type definitions *)

  | ModuleD of info
      * string               (* module name *)
      * string list          (* type parameters *)
      * where_clause
      * decl list            (* body *)
  | PrivateD of info 
  | PublicD of info
  | ImportD of info
      * import_exp           (* what to import *)
      * (info * string * typ list) list (* module path from which to import *)
      * string               (* rename *)
  | IncludeD of info         (* Include a C header *)
      * string               (* file name *)
  | UseD of info            (* Include a G file *)
      * string               (* file name *)
  | EmptyD

and import_exp =
    FunImportName of info * string
  | FunImport of info
	* string               (* function name *)
	* string list          (* template parameters *)
        * where_clause
	* param list           (* parameters *)
	* typ                  (* return type *)
  | ImportConcept of info * string
  | ImportType of info * string
  | ImportModel of info * string * typ list
  | ImportNamedModel of info * string

and req =                      (* in where clauses *)
    ModelR of info
        * string               (* concept name *)
        * typ list             (* type arguments *)
  | SameTypeR of info          (* same type requirement *)
	* typ * typ

and stmt =
    LetS of info * typ * string * exp (* type, name, initializing expression *)
  | TypedefS of info * string * typ
  | ExprS of info * exp
  | ReturnS of info * exp option
  | IfS of info * exp * stmt * stmt
  | WhileS of info * exp * stmt
  | CompoundS of info * stmt list
  | SwitchS of info * exp * (string * stmt) list
  | ImportS of info * import_exp * (info * string * typ list) list * string

and exp =
    (* constants *)
    NullE of info     (* null pointer *)
  | IntE of info * int
  | FloatE of info * float
  | DoubleE of info * float
  | BoolE of info * bool
  | StringE of info * string
  | CharE of info * char

  | VarE of info * string
  | FunE of info
	* string list          (* template parameters *)
	* param list           (* parameters *)
	* exp_typ              (* return type *)
        * (string * (exp * typ)) list  (* inits *)
	* stmt                  (* body *)
  | ApplyE of info * exp * exp list
  | InstE of info * exp * typ list
  | PrimE of info * primitive

  (* generic high-level allocation *)
  | NewE of info * alloc * typ * exp list
  | PlacementNewE of info * exp * typ * exp list
  | NewArrayE of info * alloc * typ * exp
  | NewArrayNoInitE of info * alloc * typ * exp
  | DeleteE of info * exp
  | DestroyE of info * exp

  | SizeofE of info * typ

  | StructE of info * alloc * string * typ list * (string * exp) list
  | ThisE of info

  | MemE of info * exp * string

  | ModelMemE of info * string * typ list * string (* This needs updating! -Jeremy *)
  | IfE of info * exp * exp * exp
  | SeqE of info * exp list

and typ =
    VarT of info * string
  | VertexT of int
  | ClassT of info * string * typ list
  | IntT of info
  | LongT of info
  | LongLongT of info
  | ShortT of info
  | UnsignedT of info * typ
  | SignedT of info * typ
  | ConstT of info * typ
  | FloatT of info
  | DoubleT of info
  | LongDoubleT of info
  | BoolT of info
  | StringT of info
  | CharT of info
  | WCharT of info
  | VoidT of info
  | NullT of info
  | PtrT of info * typ
  | FunT of info 
	* string list
	* where_clause
	* exp_typ list
	* exp_typ
  | OvldT of info * (Cpp_ast.exp * typ) list
  | AssocT of info * string * typ list * string (* access associated type *)
  | ElipsesT of info

let num_tparams t =
  (match t with
    FunT (_,ts,_,_,_) -> length ts
  | _ -> error UNKNOWN "num_tparams, expected function")

let rec subst_typ (sub : (string * typ) list) (t : typ) : typ =
  (match t with
      VarT (i, n) -> (try assoc n sub with Not_found -> VarT (i,n))
    | VertexT i -> VertexT i (* This is problematic! -Jeremy *)
    | ClassT (i, cn, ts) ->
	let ts = map (subst_typ sub) ts in
	ClassT (i, cn, ts)
    | PtrT (i, t) -> PtrT (i, subst_typ sub t)
    | FunT (i, ts, (rs,ss), ps, (rt,rv,rm)) -> 
	let nts = map (fun t -> uniquify_name t) ts in
	let sub2 = map2 (fun t nt -> (t, VarT (i,nt))) ts nts in
	let subty = (fun t -> subst_typ sub (subst_typ sub2 t)) in
	let rs = map (fun (n,targs) -> (n, map subty targs)) rs in
	let ss = map (fun (s,t) -> (subty s, subty t)) ss in
	let ps = map (fun (t,v,m) -> (subty t,v,m)) ps in
	let rt = subty rt in
	  FunT (i, nts, (rs,ss), ps, (rt,rv,rm))
    | OvldT (i, ts) -> error i "subst_typ: unhandled ovld type"
    | AssocT (i, cn, targs, tn) -> 
	let targs = map (subst_typ sub) targs in
	AssocT (i, cn, targs, tn)
    | ElipsesT i -> ElipsesT i
    | NullT i -> NullT i
    | VoidT i -> VoidT i
    | CharT i -> CharT i
    | WCharT i -> WCharT i
    | StringT i -> StringT i
    | BoolT i -> BoolT i
    | LongDoubleT i -> LongDoubleT i
    | DoubleT i -> DoubleT i
    | FloatT i -> FloatT i
    | ConstT (i,t) -> ConstT (i, subst_typ sub t)
    | SignedT (i,t) -> SignedT (i, subst_typ sub t)
    | UnsignedT (i,t) -> UnsignedT (i, subst_typ sub t)
    | ShortT i -> ShortT i
    | LongLongT i -> LongLongT i
    | LongT i -> LongT i
    | IntT i -> IntT i
    )

let subst_subst sub1 sub2 =
  map (fun (n,t) -> (n, subst_typ sub1 t)) sub2


let mut2str m =
  (match m with
    Constant -> "c"
  | Mutable -> "m")

let vk2str v =
  (match v with
    Lvalue -> "l"
  | Rvalue -> "r")


let rec exp_typ2mangled_name et =
  let (t,v,m) = et in
    Printf.sprintf "%s%s%s" (typ2mangled_name t) (vk2str v) (mut2str m)

and typ2mangled_name t =
  (match t with
    VarT (i, n) -> n
  | VertexT i -> Printf.sprintf "<vertex %d>" i
  | ClassT (i, n, ts) -> Printf.sprintf "class_%s" n
  | IntT i -> "i"
  | LongT i -> "l"
  | LongLongT i -> "ll"
  | ShortT i -> "s"
  | FloatT i -> "f"
  | DoubleT i -> "d"
  | LongDoubleT i -> "ld"
  | BoolT i -> "b"
  | StringT i -> "str"
  | CharT i -> "chr"
  | WCharT i -> "w"
  | VoidT i -> "v"
  | NullT i -> "n"
  | ElipsesT i -> "e"
  | PtrT (i, pt) -> 
      Printf.sprintf "p%s" (typ2mangled_name pt)
  | FunT (i, ts, (rs,ss), ps, (rt,rv,rm)) -> 
      let sub = (map2 (fun t n -> (t,VarT (i, Printf.sprintf "T%d" n))) ts (iota (length ts))) in
      let ps = map (fun (t,v,m) -> (subst_typ sub t, v, m)) ps in
      let rs = map (fun (cn,targs) -> (cn, map (subst_typ sub) targs)) rs in
      let rs = map (fun (cn,targs) -> 
			   (Printf.sprintf "%s_%s" cn
			      (String.concat "_" (map typ2mangled_name targs))))
		 rs in
      let ss = map (fun (s,t) -> (subst_typ sub s, subst_typ sub t)) ss in
      let ss = map (fun (s,t) ->
			   Printf.sprintf "%s_eq_%s" (typ2mangled_name s) (typ2mangled_name t)) ss in
      let rt = subst_typ sub rt in
      Printf.sprintf "fn_where_%s_%s__%s_%s" 
	(String.concat "_" rs) 
	(String.concat "_" ss)
	(String.concat "_" (map exp_typ2mangled_name ps)) 
	(exp_typ2mangled_name (rt,rv,rm))
  | UnsignedT (i, t) -> Printf.sprintf "u%s" (typ2mangled_name t)
  | SignedT (i, t) -> Printf.sprintf "s%s" (typ2mangled_name t)
  | ConstT (i, t) -> Printf.sprintf "%sc" (typ2mangled_name t)
  | OvldT (i, ts) -> error i "typ2mangled_name: unhandled ovld type"
  | AssocT (i, cn, targs, tn) -> 
      (*
	Printf.sprintf "%s_%s_%s" cn (String.concat "_" (map typ2mangled_name targs)) tn
      *)
      tn
  )
  
let pod (t : typ) : bool =
  (match t with
  | IntT _ | FloatT _ | DoubleT _ | BoolT _ | StringT _ | CharT _ 
  | NullT _ ->
      true
  | PtrT (_, VarT (_,_)) ->
      false
  | PtrT (_, _) ->
      true
  | _ -> false)

let empty_where = ([],[])


let combine_value_kind v1 v2 =
  (match (v1,v2) with
    (Lvalue,Lvalue) -> Lvalue
  | _ -> Rvalue)

let combine_mutability m1 m2 =
  (match (m1,m2) with
    (Mutable,Mutable) -> Mutable
  | _ -> Constant)

type decl_kind = Defn | FwdDecl

type binding = 
    VarB of Cpp_ast.exp * typ * mutability * decl_kind (* normal binding: new name, type, and mutability *)
  | OvldB of (Cpp_ast.exp * typ * decl_kind) list (* overloaded name *)

type var_env = (string * binding) list

type class_info = 
    string list (* type parameters *)
    * ((string * typ list) list * (typ * typ) list) (* where clause *)
    * (string * typ) list (* data members *)
    * typ list (* types of the constructors *)
type class_env = (string * class_info) list

type concept_info = 
    string list          (* type parameters *)
    * string list          (* associated types *)
    * (string * typ list) list (* refinements *)
    * (string * typ list) list (* requires *)
    * (string                (* function name *)
       * (string              (* mangled function name *)
	  * typ                 (* funcion type *)
	  * Cpp_ast.typ  (* C++ type *)
	  * string list
	  * (Cpp_ast.stmt option))) list   (* param names and default body *)
    * (typ * typ) list     (* same-types *)

type model_info =
    decl_kind
    * string list          (* type parameters *)
    * ((string * typ list) list * (typ * typ) list) (* where clause *)
    * typ list             (* type arguments *)
    * (string * typ) list  (* associated type definitions *)
    * string               (* dictionary name *)
    * string list          (* dictionary path *)

type concept_env = (string * concept_info) list    (* maps concept name to concept info *)
type model_env = (string * (model_info list)) list (* maps concept name to list of models *)

type vertex_info = { ty: typ; 
		     mutable pred: int list }
type graph = vertex_info array

type environment = { classes : class_env; 
		     unions : class_env; 
		     tparams : (string * typ) list;
		     vars : var_env; (* globals *)
		     local_vars : var_env;
		     concepts : concept_env;
		     models : model_env;
		     extern : bool;
		     type_graph : graph;
		     type_reps : union_find;
		     prims : (primitive * binding) list
		   }

let init_env = { unions = []; classes = []; tparams = []; 
		 vars = []; local_vars = []; 
		 concepts = []; models = []; extern = false;
		 type_graph = Array.make 1 { ty=IntT UNKNOWN; pred=[] };
		 type_reps = Array.make 1 { parent=0; rank=0 };
		 prims = []
	       }

let copy_env env =
  { env with type_graph = Array.copy env.type_graph;
    type_reps = Array.copy env.type_reps
  }

let combine_tps ts nts =
  map2 (fun t n -> (t, VarT (UNKNOWN, n))) ts nts

let env_extend_type_params ts env =
  let nts = map (fun t -> uniquify_name t) ts in
  ({ env with tparams = (combine_tps ts nts) @ (combine_tps nts nts) @ env.tparams },
   nts)


let rec uniquify_typ rename (t : typ) : typ =
  (match t with
    VarT (i, n) -> 
      (try let m = assoc n rename in VarT (i,m)
      with Not_found ->
	VarT (i, n))
  | VertexT v -> VertexT v
  | ClassT (i, n, ts) ->
      let ts = map (uniquify_typ rename) ts in ClassT (i, n, ts)
  | UnsignedT (i,t) -> UnsignedT (i, uniquify_typ rename t)
  | SignedT (i,t) -> SignedT (i, uniquify_typ rename t)
  | PtrT (i, t) -> PtrT (i, uniquify_typ rename t)
  | FunT (i, ts, (rs,ss), ps, (rt,rv,rm)) -> 
      let nts = map
	  (fun t -> uniquify_name t) ts in
      let rename2 = combine ts nts in
      let uniquify = (fun t -> uniquify_typ rename (uniquify_typ rename2 t)) in
      let rs = map
	  (fun (cn,targs) -> 
	    let targs = map uniquify targs in
	    (cn,targs)) rs in
      let ss = map
	  (fun (s,t) ->
	    (uniquify_typ rename s, uniquify t))
	  ss in
      let ps = map (fun (t,v,m) -> (uniquify t, v,m)) ps in
      let rt = uniquify rt in
      FunT (i, nts, (rs,ss), ps, (rt,rv,rm))
  | OvldT (i, ts) -> OvldT (i, map (fun (e,t) -> (e, uniquify_typ rename t)) ts)
  | AssocT (i, cn, ts, tn) -> 
      AssocT (i, cn, map (uniquify_typ rename) ts, tn)
  | _ -> t)
    
    
let poly (t : typ) : bool =
  (match t with
    VarT (_,_) 
  | AssocT (_,_,_,_) 
    -> true
  |  _ -> false)

let poly_ptr (t : typ) : bool =
  (match t with
    PtrT (_, ConstT (_, t)) -> poly t
  | PtrT (_, t) -> poly t
  | _ -> false)