(** This module describes the values manipulated by the plug-in. *)

exception Unknown_cost of string
exception Unknown_prototype of string


let string_of_mset to_list sep f mset =
  let filter (_, occ) = occ <> 0 in
  let f' (elt, occ) =
    if occ = 1 then (f elt)
    else Printf.sprintf "%d*%s" occ (f elt) in
  Misc.List.to_string sep f' (List.filter filter (to_list mset))

type prototypes = string list Misc.String.Map.t


module type S = sig

  type relation
  val is_large : relation -> bool
  val has_lower_type : relation -> bool

  type t

  val top : t
  val of_int : int -> t
  val of_var : string -> t
  val add    : t -> t -> t
  val minus  : t -> t -> t
  val mul    : t -> t -> t
  val div    : t -> t -> t
  val max    : t -> t -> t
  val cond   : t -> relation -> t -> t -> t -> t
  val join   : t -> t -> t
  val widen  : t -> t -> t
  val narrow : t -> t -> t

  val le : t -> t -> bool

  val replace_vars : t Misc.String.Map.t -> t -> t

  val to_string : t -> string
  val string_of_relation : relation -> string

  val compare : t -> t -> int

end


module Make (S : S) = struct

  let s_add_list = function
    | [] -> S.of_int 0
    | e :: l -> List.fold_left S.add e l

  module Args = struct

    type t = S.t list

    let compare = Misc.List.compare S.compare

    let le args1 args2 =
      if List.length args1 <> List.length args2 then false
      else
	let f res arg1 arg2 = res && (S.le arg1 arg2) in
	List.fold_left2 f true args1 args2

    let replace_vars replacements = List.map (S.replace_vars replacements)

    let to_string = Misc.List.to_string ", " S.to_string

  end


  module Externs = struct

    module M = Misc.String.MSet
    include M

    let add = union

    let le = subset

    let replace_vars _ externs = externs

    let to_string = string_of_mset to_list " + " (fun x -> x)

    let to_ext externs =
      let f x occ ext = S.add (S.mul (S.of_int occ) (S.of_var x)) ext in
      fold f externs (S.of_int 0)

  end


  module FunCall = struct

    type t =
	{ caller : string ;
	  id : int ;
	  callee : string ;
	  args : Args.t }
    let compare = Pervasives.compare

    let caller fun_call = fun_call.caller
    let id fun_call = fun_call.id
    let callee fun_call = fun_call.callee
    let args fun_call = fun_call.args

    let make caller id callee args = { caller ; id ; callee ; args }

    let apply f f_caller f_id f_callee f_args fun_call =
      let caller_res = f_caller (caller fun_call) in
      let id_res = f_id (id fun_call) in
      let callee_res = f_callee (callee fun_call) in
      let args_res = f_args (args fun_call) in
      f caller_res id_res callee_res args_res

    let apply2 f f_caller f_id f_callee f_args fun_call1 fun_call2 =
      let caller_res = f_caller (caller fun_call1) (caller fun_call2) in
      let id_res = f_id (id fun_call1) (id fun_call2) in
      let callee_res = f_callee (callee fun_call1) (callee fun_call2) in
      let args_res = f_args (args fun_call1) (args fun_call2) in
      f caller_res id_res callee_res args_res

    let le =
      let f b1 b2 b3 b4 = b1 && b2 && b3 && b4 in
      apply2 f (=) (=) (=) Args.le

    let replace_vars replacement =
      apply make Misc.id Misc.id Misc.id (Args.replace_vars replacement)

    let reduce
	to_list is_solved replace_vars of_fun_call
	prototypes costs fun_call =
      let callee = callee fun_call in
      let args = args fun_call in
      if Misc.String.Map.mem callee prototypes then
	let formals = Misc.String.Map.find callee prototypes in
	if List.length formals = List.length args then
	  let replacements =
	    Misc.String.Map.of_list (List.combine formals args) in
	  if Misc.String.Map.mem callee costs then
	    let cost' = Misc.String.Map.find callee costs in
	    if is_solved cost' then to_list (replace_vars replacements cost')
	    else [of_fun_call fun_call]
	  else raise (Unknown_cost callee)
	else
	  raise
	    (Failure ("FunCall.reduce: formals and actuals for " ^
			 "function " ^ callee ^ " have different sizes."))
      else raise (Unknown_prototype callee)

    let to_string fun_call =
      Printf.sprintf "%s@[%s,%d](%s)"
	(callee fun_call) (caller fun_call) (id fun_call)
	(Args.to_string (args fun_call))

  end

  module FunCalls = struct

    module M = Multiset.Make (FunCall)
    include M

    let singleton_ fun_call = M.add fun_call empty

    let singleton caller id callee args =
      singleton (FunCall.make caller id callee args)

    let add = union

    let le1 fun_call occ fun_calls =
      let f fun_call' occ' = (FunCall.le fun_call fun_call') && (occ <= occ') in
      exists f fun_calls

    let le fun_calls1 fun_calls2 =
      let f fun_call occ = le1 fun_call occ fun_calls2 in
      for_all f fun_calls1

    let called_funs fun_calls =
      let f fun_call _ called_funs =
	Misc.String.Set.add (FunCall.callee fun_call) called_funs in
      fold f fun_calls Misc.String.Set.empty

    let replace_vars replacements fun_calls =
      let f fun_call _ fun_calls =
	let fun_call = FunCall.replace_vars replacements fun_call in
	add (singleton_ fun_call) fun_calls in
      fold f fun_calls empty

    let to_string = string_of_mset to_list " + " FunCall.to_string

  end


  module rec LoopCost : sig
    type t
    val compare : t -> t -> int
    val make : string -> int -> S.relation -> S.t -> S.t -> S.t -> Cost.t -> t
    val le : t -> t -> bool
    val called_funs : t -> Misc.String.Set.t
    val replace_vars : S.t Misc.String.Map.t -> t -> t
    val reduce : prototypes -> Cost.t Misc.String.Map.t -> t -> t
    val to_string : t -> string
    val to_ext : t -> S.t
  end = struct

    type t =
	{ fun_name : string ;
	  id : int ;
	  relation : S.relation ;
	  init_value : S.t ;
	  exit_value : S.t ;
	  increment : S.t ;
	  body_cost : Cost.t }

    let make fun_name id relation init_value exit_value increment body_cost =
      { fun_name ; id ; relation ; init_value ; exit_value ; increment ;
	body_cost }

    let fun_name loop_cost = loop_cost.fun_name
    let id loop_cost = loop_cost.id
    let relation loop_cost = loop_cost.relation
    let init_value loop_cost = loop_cost.init_value
    let exit_value loop_cost = loop_cost.exit_value
    let increment loop_cost = loop_cost.increment
    let body_cost loop_cost = loop_cost.body_cost

    let compare = Pervasives.compare

    let apply f
	f_fun_name f_id f_relation f_init_value f_exit_value f_increment
	f_body_cost loop_cost =
      let fun_name_res = f_fun_name (fun_name loop_cost) in
      let id_res = f_id (id loop_cost) in
      let relation_res = f_relation (relation loop_cost) in
      let init_value_res = f_init_value (init_value loop_cost) in
      let exit_value_res = f_exit_value (exit_value loop_cost) in
      let increment_res = f_increment (increment loop_cost) in
      let body_cost_res = f_body_cost (body_cost loop_cost) in
      f
	fun_name_res id_res relation_res init_value_res exit_value_res
	increment_res body_cost_res

    let apply2 f
	f_fun_name f_id f_relation f_init_value f_exit_value f_increment
	f_body_cost loop_cost1 loop_cost2 =
      let fun_name_res =
	f_fun_name (fun_name loop_cost1) (fun_name loop_cost2) in
      let id_res = f_id (id loop_cost1) (id loop_cost2) in
      let relation_res =
	f_relation (relation loop_cost1) (relation loop_cost2) in
      let init_value_res =
	f_init_value (init_value loop_cost1) (init_value loop_cost2) in
      let exit_value_res =
	f_exit_value (exit_value loop_cost1) (exit_value loop_cost2) in
      let increment_res =
	f_increment (increment loop_cost1) (increment loop_cost2) in
      let body_cost_res =
	f_body_cost (body_cost loop_cost1) (body_cost loop_cost2) in
      f
	fun_name_res id_res relation_res init_value_res exit_value_res
	increment_res body_cost_res

    let le =
      let f b1 b2 b3 b4 b5 b6 b7 = b1 && b2 && b3 && b4 && b5 && b6 && b7 in
      apply2 f (=) (=) (=) S.le S.le S.le Cost.le

    let called_funs loop_cost = Cost.called_funs (body_cost loop_cost)

    let replace_vars replacements =
      let arg_replace_vars = S.replace_vars replacements in
      apply make Misc.id Misc.id Misc.id arg_replace_vars arg_replace_vars
	arg_replace_vars (Cost.replace_vars replacements)

    let reduce prototypes costs =
      apply make Misc.id Misc.id Misc.id Misc.id Misc.id Misc.id
	(Cost.reduce prototypes costs)

    let to_string loop_cost =
      Printf.sprintf "%s@%d(%s %s %s %s (%s))"
	(fun_name loop_cost)
	(id loop_cost)
	(S.string_of_relation (relation loop_cost))
	(S.to_string (init_value loop_cost))
	(S.to_string (exit_value loop_cost))
	(S.to_string (increment loop_cost))
	(Cost.to_string (body_cost loop_cost))

    let to_ext loop_cost =
      let rel = relation loop_cost in
      let init_value = init_value loop_cost in
      let exit_value = exit_value loop_cost in
      let increment = increment loop_cost in
      let body_cost = body_cost loop_cost in
      let rel_op = if S.has_lower_type rel then S.minus else S.add in
      let rel_added = S.of_int (if S.is_large rel then 0 else 1) in
      let iteration_nb = rel_op increment rel_added in
      let iteration_nb = S.minus iteration_nb init_value in
      let iteration_nb = S.add exit_value iteration_nb in
      let iteration_nb = S.div iteration_nb increment in
      let body_cost = Cost.to_ext body_cost in
      let cond_body_cost =
	S.cond init_value rel exit_value body_cost (S.of_int 0) in
      S.mul iteration_nb cond_body_cost

  end


  and LoopCosts : sig
    type t
    val empty : t
    val singleton :
      string -> int -> S.relation -> S.t -> S.t -> S.t -> Cost.t -> t
    val add : t -> t -> t
    val replace_vars : S.t Misc.String.Map.t -> t -> t
    val called_funs : t -> Misc.String.Set.t
    val reduce : prototypes -> Cost.t Misc.String.Map.t -> t -> t
    val to_string : t -> string
    val le : t -> t -> bool
    val to_ext : t -> S.t
  end = struct

    module M = Multiset.Make (LoopCost)
    include M

    let singleton_ loop_cost = M.add loop_cost empty

    let singleton
	fun_name id relation init_value exit_value increment body_cost =
      let loop_cost =
	LoopCost.make
	  fun_name id relation init_value exit_value increment body_cost in
      singleton_ loop_cost

    let add = union

    let le1 loop_cost occ loop_costs =
      let f loop_cost' occ' res =
	res || ((LoopCost.le loop_cost loop_cost') && (occ <= occ')) in
      fold f loop_costs false

    let le loop_costs1 loop_costs2 =
      let f loop_cost occ res = res && le1 loop_cost occ loop_costs2 in
      fold f loop_costs1 true

    let called_funs loop_costs =
      let f loop_cost _ called_funs =
	Misc.String.Set.union (LoopCost.called_funs loop_cost) called_funs in
      fold f loop_costs Misc.String.Set.empty

    let replace_vars replacements loop_costs =
      let f loop_cost _ loop_costs =
	let loop_cost = LoopCost.replace_vars replacements loop_cost in
	add (singleton_ loop_cost) loop_costs in
      fold f loop_costs empty

    let reduce prototypes replacements loop_costs =
      let f loop_cost occ loop_costs =
	add_occ (LoopCost.reduce prototypes replacements loop_cost) occ
	  loop_costs in
      fold f loop_costs empty

    let to_string = string_of_mset to_list " + " LoopCost.to_string

    let to_ext loop_costs =
      let f loop_cost occ ext =
	S.add (S.mul (S.of_int occ) (LoopCost.to_ext loop_cost)) ext in
      fold f loop_costs (S.of_int 0)

  end


  and Base : sig
    type t
    val compare : t -> t -> int
    val of_int : int -> t
    val of_extern : string -> t
    val of_fun_call_ : FunCall.t -> t
    val of_fun_call : string -> int -> string -> Args.t -> t
    val of_loop_cost :
      string -> int -> S.relation -> S.t -> S.t -> S.t -> Cost.t -> t
    val add : t -> t -> t
    val called_funs : t -> Misc.String.Set.t
    val replace_vars : S.t Misc.String.Map.t -> t -> t
    val reduce : prototypes -> Cost.t Misc.String.Map.t -> t -> Base.t list
    val le : t -> t -> bool
    val to_string : t -> string
    val to_ext : t -> S.t
  end = struct

    type t =
	{ constant : int ;
	  externs : Externs.t ;
	  fun_calls : FunCalls.t ;
	  loop_costs : LoopCosts.t }

    let make constant externs fun_calls loop_costs =
      { constant ; externs ; fun_calls ; loop_costs }

    let compare = Pervasives.compare

    let constant base = base.constant
    let externs base = base.externs
    let fun_calls base = base.fun_calls
    let loop_costs base = base.loop_costs

    let set_fun_calls fun_calls base = { base with fun_calls }
    let set_loop_costs loop_costs base = { base with loop_costs }

    let to_string base =
      Printf.sprintf "%d + (%s) + (%s) + (%s)"
	(constant base)
	(Externs.to_string (externs base))
	(FunCalls.to_string (fun_calls base))
	(LoopCosts.to_string (loop_costs base))

    let of_int i =  make i Externs.empty FunCalls.empty LoopCosts.empty

    let of_extern x =
      make 0 (Externs.singleton x) FunCalls.empty LoopCosts.empty

    let of_fun_call_ fun_call =
      make 0 Externs.empty (FunCalls.singleton_ fun_call) LoopCosts.empty

    let of_fun_call caller id callee args =
      let fun_call = FunCall.make caller id callee args in
      of_fun_call_ fun_call

    let of_loop_cost
	fun_name id relation init_value exit_value increment body_cost =
      let loop_costs =
	LoopCosts.singleton
	  fun_name id relation init_value exit_value increment body_cost in
      make 0 Externs.empty FunCalls.empty loop_costs

    let apply f f_constant f_externs f_fun_calls f_loop_costs base =
      let constant_res = f_constant (constant base) in
      let externs_res = f_externs (externs base) in
      let fun_calls_res = f_fun_calls (fun_calls base) in
      let loop_costs_res = f_loop_costs (loop_costs base) in
      f constant_res externs_res fun_calls_res loop_costs_res

    let apply2 f f_constant f_externs f_fun_calls f_loop_costs base1 base2 =
      let constant_res = f_constant (constant base1) (constant base2) in
      let externs_res = f_externs (externs base1) (externs base2) in
      let fun_calls_res = f_fun_calls (fun_calls base1) (fun_calls base2) in
      let loop_costs_res = f_loop_costs (loop_costs base1) (loop_costs base2) in
      f constant_res externs_res fun_calls_res loop_costs_res

    let add = apply2 make (+) Externs.add FunCalls.add LoopCosts.add

    let le =
      let f b1 b2 b3 b4 = b1 && b2 && b3 && b4 in
      apply2 f (<=) Externs.le FunCalls.le LoopCosts.le

    let replace_vars replacements =
      apply make Misc.id
	(Externs.replace_vars replacements)
	(FunCalls.replace_vars replacements)
	(LoopCosts.replace_vars replacements)

    let called_funs base =
      Misc.String.Set.union
	(FunCalls.called_funs (fun_calls base))
	(LoopCosts.called_funs (loop_costs base))

    let reduce prototypes costs base =
      let f fun_call occ base_list =
	let added_bases =
	  FunCall.reduce
	    Cost.to_list Cost.is_solved Cost.replace_vars of_fun_call_
	    prototypes costs fun_call in
	let added_bases =
	  if added_bases = [] then [of_int 0] else added_bases in
	let added_bases =
	  let f base = Misc.repeat occ (add base) (of_int 0) in
	  List.map f added_bases in
	let f_base_list res added_base =
	  res @ (List.map (add added_base) base_list) in
	List.fold_left f_base_list [] added_bases in
      let loop_costs = LoopCosts.reduce prototypes costs (loop_costs base) in
      let base = set_loop_costs loop_costs base in
      let base' = set_fun_calls FunCalls.empty base in
      FunCalls.fold f (fun_calls base) [base']

    let to_ext base =
      let f_fun_calls fun_calls =
	if not (FunCalls.is_empty fun_calls) then
	  raise (Failure "Base.to_ext: function calls")
	else S.of_int 0 in
      let f ext1 ext2 ext3 ext4 = s_add_list [ext1 ; ext2 ; ext3 ; ext4] in
      apply f S.of_int Externs.to_ext f_fun_calls LoopCosts.to_ext base

  end


  and Cost : sig
    type t
    val of_int : int -> t
    val of_extern : string -> t
    val of_fun_call : string -> int -> string -> Args.t -> t
    val of_loop_cost :
      string -> int -> S.relation -> S.t -> S.t -> S.t -> Cost.t -> t
    val of_base : Base.t -> t
    val empty : t
    val add : t -> t -> t
    val join : t -> t -> t
    val widen : t -> t -> t
    val narrow : t -> t -> t
    val called_funs : t -> Misc.String.Set.t
    val has_fun_calls : t -> bool
    val replace_vars : S.t Misc.String.Map.t -> t -> t
    val reduce : prototypes -> Cost.t Misc.String.Map.t -> t -> t
    val is_solved : t -> bool
    val to_list : t -> Base.t list
    val to_string : t -> string
    val le : t -> t -> bool
    val to_ext : t -> S.t
  end = struct

    module M = Eset.Make (Base)
    include M

    let to_string cost =
      if is_empty cost then "0"
      else Misc.List.to_string " max " Base.to_string (to_list cost)

    let of_base base = singleton base

    let of_int i = of_base (Base.of_int i)

    let of_extern x = of_base (Base.of_extern x)

    let of_fun_call caller id callee args =
      of_base (Base.of_fun_call caller id callee args)

    let of_loop_cost
	fun_name loop_id relation init_value exit_value increment body_cost =
      of_base
	(Base.of_loop_cost
	   fun_name loop_id relation init_value exit_value increment body_cost)

    let join1 base cost =
      let f_exists base' = Base.le base base' in
      if exists f_exists cost then cost
      else
	let f_absorb base' = Base.le base' base in
	M.add base (M.diff cost (M.filter f_absorb cost))

    let add cost1 cost2 =
      if is_empty cost1 then cost2
      else
	if is_empty cost2 then cost1
	else
	  let f2 base1 base2 = join1 (Base.add base1 base2) in
	  let f1 base1 = fold (f2 base1) cost2 in
	  fold f1 cost1 empty

    let join cost1 cost2 =
      if is_empty cost1 then cost2
      else
	if is_empty cost2 then cost1
	else fold join1 cost1 cost2

    let widen = join

    let narrow = join (* TODO: improve *)

    let mem base cost =
      let f base' res = res || (Base.le base base') in
      fold f cost false

    let le cost1 cost2 =
      let f base res = res && (mem base cost2) in
      fold f cost1 true


    let called_funs cost =
      let f base called_funs =
	Misc.String.Set.union (Base.called_funs base) called_funs in
      fold f cost Misc.String.Set.empty

    let has_fun_calls cost = not (Misc.String.Set.is_empty (called_funs cost))

    let replace_vars replacements cost =
      let f base cost = join1 (Base.replace_vars replacements base) cost in
      fold f cost empty

    let reduce prototypes costs cost =
      let f base cost =
	let base_list = Base.reduce prototypes costs base in
	let f_join cost base = join1 base cost in
	List.fold_left f_join cost base_list in
      fold f cost empty

    let is_solved cost = not (has_fun_calls cost)

    let to_ext cost =
      if is_empty cost then S.of_int 0
      else
	let f base ext = S.max (Base.to_ext base) ext in
	let base = choose cost in
	let cost = remove base cost in
	fold f cost (Base.to_ext base)

  end


  type t = Top | C of Cost.t


  let to_string = function
    | Top -> "top"
    | C cost -> Cost.to_string cost


  let of_int i = C (Cost.of_int i)

  let of_extern fun_name = C (Cost.of_extern fun_name)

  let of_fun_call caller id callee args =
    C (Cost.of_fun_call caller id callee args)

  let of_loop_cost
      fun_name loop_id relation init_value exit_value increment body_cost =
    C (Cost.of_loop_cost
	 fun_name loop_id relation init_value exit_value increment body_cost)


  let is_top = function Top -> true | _ -> false

  let extract = function
    | Top -> raise (Failure "Cost_value.extract")
    | C cost -> cost

  let top = Top

  let bot = of_int 0


  let top_absorbs f = function
    | Top -> Top
    | C cost -> C (f cost)

  let top_absorbs2 f cost1 cost2 = match cost1, cost2 with
    | Top, _ | _, Top -> Top
    | C cost1, C cost2 -> C (f cost1 cost2)


  let add = top_absorbs2 Cost.add

  let join = top_absorbs2 Cost.join

  let widen = top_absorbs2 Cost.widen

  let narrow cost1 cost2 = match cost1, cost2 with
    | cost, Top | Top, cost -> cost
    | C cost1, C cost2 -> C (Cost.narrow cost1 cost2)

  let le cost1 cost2 = match cost1, cost2 with
    | _, Top -> true
    | Top, _ -> false
    | C cost1, C cost2 -> Cost.le cost1 cost2


  let reduce_ prototypes costs cost =
    let called_funs = Cost.called_funs cost in
    let costs =
      let f fun_name _ = Misc.String.Set.mem fun_name called_funs in
      Misc.String.Map.filter f costs in
    let f fun_name cost costs = match cost, costs with
      | _, None -> None
      | Top, _ -> None
      | C cost, Some costs -> Some (Misc.String.Map.add fun_name cost costs) in
    match Misc.String.Map.fold f costs (Some Misc.String.Map.empty) with
      | None -> Top
      | Some costs -> C (Cost.reduce prototypes costs cost)

  let reduce prototypes costs = function
    | Top -> Top
    | C cost -> reduce_ prototypes costs cost

  let reduces prototypes costs =
    Misc.String.Map.map (reduce prototypes costs) costs


  let replace_vars replacements = top_absorbs (Cost.replace_vars replacements)


  let has_fun_calls = function
    | Top -> false
    | C cost -> Cost.has_fun_calls cost


  let is_concrete cost = (not (is_top cost)) && (not (has_fun_calls cost))

  let to_ext = function
    | Top -> S.top
    | C cost -> Cost.to_ext cost


end