(**
   A simple (and currently unused) set of tools for obtaining provenance 
   information about a given calculus expression and/or the domains of 
   variables in this expression.
*)

open Type
open Calculus

let (_,_,error,bug) = Debug.Logger.functions_for_module "CALCULUSPROVENANCE"

type lf_t = 
   | RelSource of string
   | InlineSource
   | NoSource
   | AnySource

module ProvenanceBase = struct
   type t = lf_t
   let zero = AnySource
   let one  = NoSource

   let is_zero (v: t) = (v = zero)

   let is_one (v: t) = (v = one)  
end

module ProvenanceRing = Ring.Make(ProvenanceBase)

type t = ProvenanceRing.expr_t

let union         : t list -> t = ProvenanceRing.mk_sum
let join          : t list -> t = ProvenanceRing.mk_prod
let rel_source rel: t           = ProvenanceRing.mk_val (RelSource(rel))
let inline_source : t           = ProvenanceRing.mk_val InlineSource
let no_source     : t           = ProvenanceRing.mk_val NoSource
let fold                        = ProvenanceRing.fold

(**
   Generate a string representation of a provenance expression
*)
let string_of_provenance: (t -> string) =
   fold (fun x -> "("^(ListExtras.string_of_list ~sep:" || " (fun a->a) x)^")")
        (fun x -> "("^(ListExtras.string_of_list ~sep:" && " (fun a->a) x)^")")
        (fun x -> "(NOT "^x^")")
        (function
            | RelSource(r) -> r
            | InlineSource -> "[inline expr]"
            | NoSource -> "[nothing]"
            | AnySource -> "*"
         )

(**
   Compute the full provenance information for all output variables, and the
   value generated by a specified expression
*)
let rec provenance_of_expr ?(in_scope=[]) ?(inline_vars=[]) (expr:expr_t): 
                           ((var_t * t) list * t) =
   let scope = ListAsSet.union in_scope (List.map (fun v -> (v,inline_source))
                                                  inline_vars) in
   let merge merge_op var_val_prov =
      let (var_prov, val_prov) = List.split var_val_prov in
         (  List.map (fun (a, a_prov) -> (a, merge_op a_prov))
                     (ListExtras.reduce_assoc (List.flatten var_prov)),
            merge_op val_prov
         )
   in
   let var_provenance vars =
         try 
            (* We use information flow to derive the provenance of variables in 
               a value or a comparison.  These variables don't introduce any new
               variables into scope, but the value of the expression itself 
               derives its value from these variables. The value doesn't 
               introduce any new variables into scope, but the value itself 
               derives its value from the variables in scope. *)
            join (List.map (ListAsFunction.apply_strict scope) vars)
         with ListAsFunction.NonFunctionalMappingException -> (
            error ~detail:(fun () ->
               "Expression:\n"^(CalculusPrinter.string_of_expr expr)^"\n"^
               "Vars:"^(ListExtras.ocaml_of_list string_of_var vars)^"\n"^
               "Scope: "^(ListExtras.ocaml_of_list (fun (vn,prov) ->
                     (string_of_var vn)^"->"^(string_of_provenance prov)
                  ) scope)
            ) "Computing provenance of an expression with an unbound variable"; 
            (no_source)
         )
   
   in
   let rcr scope e = provenance_of_expr ~in_scope:scope e in
   begin match expr with 
      | CalcRing.Sum(sl) -> merge union (List.map (rcr scope) sl)
      | CalcRing.Prod([]) -> ([], no_source)
      | CalcRing.Prod(head::rest) ->
         let (head_prov) = rcr scope head in
         let (new_scope,_) = merge join [scope,no_source; head_prov] in
            merge join [head_prov; rcr new_scope (CalcRing.mk_prod rest)]
      | CalcRing.Neg(x) -> rcr scope x

      | CalcRing.Val(Value(Arithmetic.ValueRing.Val(Arithmetic.AVar(v)))) ->
            ( [], var_provenance [v] )

      | CalcRing.Val(Value(v)) ->
            ( [], join [inline_source; 
                        var_provenance (Arithmetic.vars_of_value v)] )

      | CalcRing.Val(Cmp(_,v1,v2)) ->
            ( [], join [inline_source;
                        var_provenance (ListAsSet.union 
                                    (Arithmetic.vars_of_value v1) 
                                    (Arithmetic.vars_of_value v1))]
            )

      | CalcRing.Val(CmpOrList(v,_)) ->
            ( [], join [inline_source; var_provenance 
                                          (Arithmetic.vars_of_value v)]
            )
      
      | CalcRing.Val(Rel(rn,rv)) ->
            (  List.map (fun v -> (v, rel_source rn)) rv,
               rel_source rn
            )

      | CalcRing.Val(DeltaRel(rn,rv)) ->
            (  List.map (fun v -> (v, rel_source rn)) rv,
               rel_source rn
            )
            
      | CalcRing.Val(DomainDelta(subexp)) ->
         let (var_prov, val_prov) = rcr scope subexp in
            ( var_prov, inline_source )

      | CalcRing.Val(AggSum(gb_vars, subexp)) ->
         (* The value provenance is affected by any gb_vars that get projected
            away and aggregated over *)
         let (var_prov, val_prov) = rcr scope subexp in
            List.fold_right (fun (curr_var,curr_var_prov) 
                                 (old_var_prov, old_val_prov) ->
               if List.mem curr_var gb_vars then
                  (  (curr_var,curr_var_prov)::old_var_prov, 
                     old_val_prov
                  )
               else 
                  (  old_var_prov,
                     join [old_val_prov; curr_var_prov]
                  )
            ) var_prov ([], val_prov)
            
      | CalcRing.Val(Lift(lv,subexp)) ->
         let (var_prov, val_prov) = rcr scope subexp in
            if List.mem_assoc lv var_prov then (
               error ~detail:(fun () -> CalculusPrinter.string_of_expr expr)
                  ("Computing provenance of lift into variable bound in "^
                   "nested expression"); ([], no_source)
            ) else (
               ( (lv,val_prov)::var_prov, inline_source )
            )

(***** BEGIN EXISTS HACK *****)
      | CalcRing.Val(Exists(subexp)) ->
         let (var_prov, val_prov) = rcr scope subexp in
            ( var_prov, inline_source )
(***** END EXISTS HACK *****)

      | CalcRing.Val(External(e)) ->
            bug ~detail:(fun () -> Calculus.string_of_leaf (External(e)))
                "Provenance for a materialized expression"; ([], no_source)
   end

(**
   Compute the full provenance information for the value generated by a
   specified expression.
*)
let rec provenance_of_val ?(inline_vars=[]) (expr:expr_t): t =
   snd (provenance_of_expr ~inline_vars:inline_vars expr)

(**
   Compute the full provenance information for the domain of a specified 
   variable in a specified expression
*)
let rec provenance_of_var ?(inline_vars=[]) (expr:expr_t) (var:var_t): t =
   List.assoc var (fst (provenance_of_expr ~inline_vars:inline_vars expr))

*
   A simple (and currently unused) set of tools for obtaining provenance 
   information about a given calculus expression and/or the domains of 
   variables in this expression.

*
   Generate a string representation of a provenance expression

*
   Compute the full provenance information for all output variables, and the
   value generated by a specified expression

 We use information flow to derive the provenance of variables in 
               a value or a comparison.  These variables don't introduce any new
               variables into scope, but the value of the expression itself 
               derives its value from these variables. The value doesn't 
               introduce any new variables into scope, but the value itself 
               derives its value from the variables in scope. 
 The value provenance is affected by any gb_vars that get projected
            away and aggregated over 
**** BEGIN EXISTS HACK ****
**** END EXISTS HACK ****
*
   Compute the full provenance information for the value generated by a
   specified expression.

*
   Compute the full provenance information for the domain of a specified 
   variable in a specified expression


open Type
open Calculus

let (_,_,error,bug) = Debug.Logger.functions_for_module "CALCULUSPROVENANCE"

type lf_t = 
   | RelSource of string
   | InlineSource
   | NoSource
   | AnySource

module ProvenanceBase = struct
   type t = lf_t
   let zero = AnySource
   let one  = NoSource

   let is_zero (v: t) = (v = zero)

   let is_one (v: t) = (v = one)  
end

module ProvenanceRing = Ring.Make(ProvenanceBase)

type t = ProvenanceRing.expr_t

let union         : t list -> t = ProvenanceRing.mk_sum
let join          : t list -> t = ProvenanceRing.mk_prod
let rel_source rel: t           = ProvenanceRing.mk_val (RelSource(rel))
let inline_source : t           = ProvenanceRing.mk_val InlineSource
let no_source     : t           = ProvenanceRing.mk_val NoSource
let fold                        = ProvenanceRing.fold

let string_of_provenance: (t -> string) =
   fold (fun x -> "("^(ListExtras.string_of_list ~sep:" || " (fun a->a) x)^")")
        (fun x -> "("^(ListExtras.string_of_list ~sep:" && " (fun a->a) x)^")")
        (fun x -> "(NOT "^x^")")
        (function
            | RelSource(r) -> r
            | InlineSource -> "[inline expr]"
            | NoSource -> "[nothing]"
            | AnySource -> "*"
         )

let rec provenance_of_expr ?(in_scope=[]) ?(inline_vars=[]) (expr:expr_t): 
                           ((var_t * t) list * t) =
   let scope = ListAsSet.union in_scope (List.map (fun v -> (v,inline_source))
                                                  inline_vars) in
   let merge merge_op var_val_prov =
      let (var_prov, val_prov) = List.split var_val_prov in
         (  List.map (fun (a, a_prov) -> (a, merge_op a_prov))
                     (ListExtras.reduce_assoc (List.flatten var_prov)),
            merge_op val_prov
         )
   in
   let var_provenance vars =
         try 
            join (List.map (ListAsFunction.apply_strict scope) vars)
         with ListAsFunction.NonFunctionalMappingException -> (
            error ~detail:(fun () ->
               "Expression:\n"^(CalculusPrinter.string_of_expr expr)^"\n"^
               "Vars:"^(ListExtras.ocaml_of_list string_of_var vars)^"\n"^
               "Scope: "^(ListExtras.ocaml_of_list (fun (vn,prov) ->
                     (string_of_var vn)^"->"^(string_of_provenance prov)
                  ) scope)
            ) "Computing provenance of an expression with an unbound variable"; 
            (no_source)
         )
   
   in
   let rcr scope e = provenance_of_expr ~in_scope:scope e in
   begin match expr with 
      | CalcRing.Sum(sl) -> merge union (List.map (rcr scope) sl)
      | CalcRing.Prod([]) -> ([], no_source)
      | CalcRing.Prod(head::rest) ->
         let (head_prov) = rcr scope head in
         let (new_scope,_) = merge join [scope,no_source; head_prov] in
            merge join [head_prov; rcr new_scope (CalcRing.mk_prod rest)]
      | CalcRing.Neg(x) -> rcr scope x

      | CalcRing.Val(Value(Arithmetic.ValueRing.Val(Arithmetic.AVar(v)))) ->
            ( [], var_provenance [v] )

      | CalcRing.Val(Value(v)) ->
            ( [], join [inline_source; 
                        var_provenance (Arithmetic.vars_of_value v)] )

      | CalcRing.Val(Cmp(_,v1,v2)) ->
            ( [], join [inline_source;
                        var_provenance (ListAsSet.union 
                                    (Arithmetic.vars_of_value v1) 
                                    (Arithmetic.vars_of_value v1))]
            )

      | CalcRing.Val(CmpOrList(v,_)) ->
            ( [], join [inline_source; var_provenance 
                                          (Arithmetic.vars_of_value v)]
            )
      
      | CalcRing.Val(Rel(rn,rv)) ->
            (  List.map (fun v -> (v, rel_source rn)) rv,
               rel_source rn
            )

      | CalcRing.Val(DeltaRel(rn,rv)) ->
            (  List.map (fun v -> (v, rel_source rn)) rv,
               rel_source rn
            )
            
      | CalcRing.Val(DomainDelta(subexp)) ->
         let (var_prov, val_prov) = rcr scope subexp in
            ( var_prov, inline_source )

      | CalcRing.Val(AggSum(gb_vars, subexp)) ->
         let (var_prov, val_prov) = rcr scope subexp in
            List.fold_right (fun (curr_var,curr_var_prov) 
                                 (old_var_prov, old_val_prov) ->
               if List.mem curr_var gb_vars then
                  (  (curr_var,curr_var_prov)::old_var_prov, 
                     old_val_prov
                  )
               else 
                  (  old_var_prov,
                     join [old_val_prov; curr_var_prov]
                  )
            ) var_prov ([], val_prov)
            
      | CalcRing.Val(Lift(lv,subexp)) ->
         let (var_prov, val_prov) = rcr scope subexp in
            if List.mem_assoc lv var_prov then (
               error ~detail:(fun () -> CalculusPrinter.string_of_expr expr)
                  ("Computing provenance of lift into variable bound in "^
                   "nested expression"); ([], no_source)
            ) else (
               ( (lv,val_prov)::var_prov, inline_source )
            )

      | CalcRing.Val(Exists(subexp)) ->
         let (var_prov, val_prov) = rcr scope subexp in
            ( var_prov, inline_source )

      | CalcRing.Val(External(e)) ->
            bug ~detail:(fun () -> Calculus.string_of_leaf (External(e)))
                "Provenance for a materialized expression"; ([], no_source)
   end

let rec provenance_of_val ?(inline_vars=[]) (expr:expr_t): t =
   snd (provenance_of_expr ~inline_vars:inline_vars expr)

let rec provenance_of_var ?(inline_vars=[]) (expr:expr_t) (var:var_t): t =
   List.assoc var (fst (provenance_of_expr ~inline_vars:inline_vars expr))
