Source file rewriting.ml

open UtilsLib
open DatalogLib
module ASProg = Datalog_AbstractSyntax.AbstractSyntax.Program
module ASPred = Datalog_AbstractSyntax.AbstractSyntax.Predicate
module ASRule = Datalog_AbstractSyntax.AbstractSyntax.Rule
open Datalog

module Log = Xlog.Make (struct
  let name = "Rewriting"
end)

module PredicateSet = Set.Make (struct
  type t = ASPred.predicate

  let compare pred pred2 = ASPred.compare ~with_arguments:false pred pred2
end)

module QueryMap = Map.Make (struct
  type t = ASPred.pred_id * Adornment.status list

  let compare = Stdlib.compare
end)

module Elt = struct
  type t = Adornment.status

  let to_string = Adornment.status_to_string
end

module Ads = Product.Make (Elt)

let generate_permutations len =
  let values = Ads.list_to_tuple [ Adornment.Bound ] in 
  Ads.product len values

type magic_context = {
  predicates_from_magic : Magic.extra_pred_type ASPred.PredIdMap.t;
  (* sets for the supp_zero, supp_k and magic added predicates *)
  unique_to_original_rules_ids_map : (int * int) Datalog_AbstractSyntax.RuleIdMap.t;
  magic_to_unique_rule_ids_map : int Datalog_AbstractSyntax.RuleIdMap.t;
  adorn_to_unadorm_pred_ids_map :
    Unique_binding.ASPred.pred_id Unique_binding.ASPred.PredIdMap.t;
  unique_binding_program : ASProg.program;
}

exception Break of ASRule.rule

let find_in_set f s =
  try
    ASRule.Rules.fold
      (fun r acc ->
        match acc with
        | None -> if f r then raise (Break r) else acc
        | Some _ -> acc)
      s None
  with Break r -> Some r

let pp_find_rule msg ?(with_id = false) prog fmt id =
  match find_in_set (fun r -> r.ASRule.id = id) prog.ASProg.rules with
  | Some r ->
    ASRule.pp ~with_id prog.ASProg.pred_table prog.ASProg.const_table fmt r
  | None ->
    Format.fprintf fmt "Could not find any rule (%s) with id %d" msg id

let pp_mapping ?(with_id = false) f new_prog old_prog fmt map =
  let () =
    Format.fprintf
      fmt
      "[log_mapping] Old prog is:@,@[<v>  @[%a@]@]@,[log_mapping] New prog is:@,@[<v> @[%a@]@]"
      (ASProg.pp ~with_position:false ~with_id)
      old_prog
      (ASProg.pp ~with_position:false ~with_id)
      new_prog in
  Datalog_AbstractSyntax.RuleIdMap.iter
    (fun k v ->
       let () = Format.fprintf fmt "Lookup for the %d -->> %d mapping:@," k (f v) in
       Format.fprintf
         fmt
         "%a   -->> %a@,"
         (pp_find_rule ~with_id "in new_prog" new_prog)
         k
         (pp_find_rule ~with_id "in old_prog" old_prog)
         (f v))
    map

let rewrite_program ~msg prog query =
  let () = Log.debug (fun m -> m "Starting rewrite for %s" msg) in
  let pred, ad = query in
  let pred_name =
    ASPred.PredIdTable.find_sym_from_id pred.ASPred.p_id prog.ASProg.pred_table
  in
  let c = Mtime_clock.counter () in
  let () =
    Log.info (fun m ->
        let adornment = Adornment.to_string ad in
        m "%sMagic program for query: %s_%s/%d" msg pred_name adornment
          pred.ASPred.arity ~tags:(Xlog.stamp c))
  in
  let () =
    Log.debug (fun m ->
        m "@[<v>Rewriting the following program:@,@[<v>%a@]@]"
          (fun fmt rules ->
             ASRule.Rules.iter
               (fun r ->
                  Format.fprintf
                    fmt
                    "Head: %a ->@,Rule: @[<hv>%a@]@,"
                    (ASPred.pp ~with_arity:true 
                       prog.ASProg.pred_table prog.ASProg.const_table)
                    r.ASRule.lhs
                    (ASRule.pp ~with_position:true ~with_id:true
                     ~with_arity:true prog.ASProg.pred_table
                     prog.ASProg.const_table)
                    r)
               rules)
          prog.ASProg.rules)  in
  let t = Timer.top () in
  let rgg = Rgg.build_rgg prog query in
  let () =
    Timer.debug (fun m ->
        let adornment = Adornment.to_string ad in
        m "@[For query: %s_%s/%d, rgg@ built@ in: %a@]" pred_name adornment
          pred.ASPred.arity Timer.elapsed t)  in
  let t = Timer.top () in
  (* ATTENTION : map du programme unique binding (pas magic !) vers program original *)
  let unique_binding = Unique_binding.derive_program rgg prog in
  let () = Log.info (fun m -> m "Done.") in
  let unique_prog, unadorned_map, unique_prog_rules_to_old_rule_id_map
      =
    Unique_binding.
      ( unique_binding.adorned_prog,
        unique_binding.unadorned_map,
        unique_binding.rule_to_rule_map )
  in
  let () =
    Timer.debug (fun m ->
        m "Unique binding computed in: %a" Timer.elapsed t)
  in
  let t = Timer.top () in
  let magic_prog, new_preds, magic_rules_to_unique_rule_id_map =
    Magic.make_magic unique_prog
  in
  let () =
    Timer.debug (fun m ->
        m "Running magic lasted: %a" Timer.elapsed t)
  in
  let () =
    Log.debug
      (fun m ->
         m
           "(SP) Mapping of unique binding rules to original program rules:@,@[<v>  @[%a@]@]"
           (pp_mapping
              ~with_id:true
              (fun (id, _) -> id)
              unique_prog
              prog)
           unique_prog_rules_to_old_rule_id_map)
  in
  let () =
    Log.debug (fun m ->
          m "(SP) Mapping of magic rules to unique binding program rules:@,@[<v>  @[%a@]@]"
            (pp_mapping
               ~with_id:true
               (fun id -> id)
               magic_prog
               unique_prog)
            magic_rules_to_unique_rule_id_map)
  in
  
  let () =
    Log.debug (fun m ->
        m
          "New Magic prog:@,@[<v>@[%a@]@]"
          (ASProg.pp ~with_position:false ~with_id:true)
          magic_prog)
  in
  let magic_context =
    {
      predicates_from_magic = new_preds;
      adorn_to_unadorm_pred_ids_map = unadorned_map;
      unique_to_original_rules_ids_map = unique_prog_rules_to_old_rule_id_map;
      magic_to_unique_rule_ids_map = magic_rules_to_unique_rule_id_map;
      unique_binding_program = unique_prog;
    }
  in
  let res = (magic_prog, magic_context) in
  res


let rewrite_programs ?(msg = "") prog =
  let msg = if msg = "" then msg else Printf.sprintf "%s. " msg in
  let possible_adornments = Utils.IntMap.empty in
  let all_queries_to_prog, _ =
    ASRule.Rules.fold
      (fun rule ((queries_to_prog,all_adornments) as acc) ->
         let head = rule.ASRule.lhs in
         (* Check if the predicate corresponding to the head has already
            been processed *)
         if QueryMap.exists (fun (p_id, _) _ -> p_id = head.ASPred.p_id) queries_to_prog then
           (* if yes, then do nothing *)
           let () = Log.debug (fun m -> m "Intensional predicate %a already processed. Skipping it" (ASPred.pp prog.ASProg.pred_table prog.ASProg.const_table) head) in
           acc
         else
           (* generates all possible adornments for this predicate *)
           let adornments,new_all_adornments =
             match Utils.IntMap.find_opt head.ASPred.arity all_adornments with
             | None ->
               let ads = generate_permutations head.ASPred.arity in
               ads, Utils.IntMap.add head.ASPred.arity ads all_adornments
             | Some ads -> ads, all_adornments in
           (Ads.fold
              (fun l_acc ad ->
                 let adornment = Ads.tuple_to_list ad in
                 let () = Log.info (fun m -> m "Processing Intensional predicate %a with adornment: %s" (ASPred.pp prog.ASProg.pred_table prog.ASProg.const_table) head (Adornment.to_string adornment)) in
                 QueryMap.add (head.ASPred.p_id, adornment)
                   (rewrite_program ~msg prog (head, adornment))
                   l_acc)
              queries_to_prog
              adornments,
            new_all_adornments))
      prog.ASProg.rules
      (QueryMap.empty,possible_adornments) in
  let () = Log.info (fun m -> m "rewrite_programms done.") in
  all_queries_to_prog
  


    
let get_program query query_map =
  (*We build the adornment of term*)
  let bfs, _ = Adornment.adornment ~bound_variables:ASPred.TermSet.empty query in
  (*We get the progran in the map *)
  let the_magic_program, magic_context =
    QueryMap.find (query.ASPred.p_id, bfs) query_map
  in
  (*We build the magic rule*)
  (Magic.query_to_seed query the_magic_program, magic_context)


let rec unmagic_premises premises pred_from_magic derivations =
  match premises with
  | [supp_k; pred] when ASPred.(PredIdMap.find_opt supp_k.p_id pred_from_magic) = Some Magic.Supp_k &&
                        ASPred.(PredIdMap.find_opt pred.p_id pred_from_magic) = None ->
    let supp_k_premises_set = Datalog.Predicate.PredicateMap.find supp_k derivations in
    Datalog.Predicate.PremiseSet.fold
      (fun (l_premises, _, _) acc ->
         (List.fold_left
            (fun l_acc elt -> (pred::elt)::l_acc)
            acc
            (unmagic_premises l_premises pred_from_magic derivations)))
      supp_k_premises_set
      []
  | [pred; supp_k] when ASPred.(PredIdMap.find_opt supp_k.p_id pred_from_magic) = Some Magic.Supp_k &&
                        ASPred.(PredIdMap.find_opt pred.p_id pred_from_magic) = None ->
    let supp_k_premises_set = Datalog.Predicate.PredicateMap.find supp_k derivations in
    Datalog.Predicate.PremiseSet.fold
      (fun (l_premises, _, _) acc ->
         (List.fold_left
            (fun l_acc elt -> (pred::elt)::l_acc)
            acc
            (unmagic_premises l_premises pred_from_magic derivations)))
      supp_k_premises_set
      []
  | [supp_zero; pred] when ASPred.(PredIdMap.find_opt supp_zero.p_id pred_from_magic) = Some Magic.Supp_zero &&
                           ASPred.(PredIdMap.find_opt pred.p_id pred_from_magic) = None ->
    [[pred]]
  | [pred; supp_zero] when ASPred.(PredIdMap.find_opt supp_zero.p_id pred_from_magic) = Some Magic.Supp_zero &&
                           ASPred.(PredIdMap.find_opt pred.p_id pred_from_magic) = None ->
    [[pred]]
  | [supp_zero] when ASPred.(PredIdMap.find_opt supp_zero.p_id pred_from_magic) = Some Magic.Supp_zero ->
    []
  | _ -> failwith "Bug: Cannot be premises of a supp_k predicate"

let rewrite_derivations derivations
    {
      predicates_from_magic;
      adorn_to_unadorm_pred_ids_map;
      magic_to_unique_rule_ids_map;
      unique_to_original_rules_ids_map;
      unique_binding_program = _;
    } =
  let rename_predicate p =
    match ASPred.(PredIdMap.find_opt p.p_id adorn_to_unadorm_pred_ids_map) with
    | None -> failwith "Bug: only non-magic and non-supp predicates can be renamed"
    | Some id -> ASPred.copy_predicate ~new_id:id p in
  let find_original_rule_characteristics r_id =
    Datalog_AbstractSyntax.RuleIdMap.(find (find r_id magic_to_unique_rule_ids_map) unique_to_original_rules_ids_map) in    
  Datalog.Predicate.PredicateMap.fold
    (fun pred premises_set l_derivations ->
       match ASPred.(PredIdMap.find_opt pred.p_id predicates_from_magic) with
       | Some _ ->
         (* just remove derivation trees for supp_0, supp_k and magic
            predicates *)
         l_derivations
       | None ->
         (* otherwise, proceed *)
         let new_premises =
           Datalog.Predicate.PremiseSet.fold
             (fun (preds, r_id, _) ll_premise_set ->
                let original_rule_id, original_num = find_original_rule_characteristics r_id in
                match unmagic_premises preds predicates_from_magic derivations with
                | [] -> 
                  Datalog.Predicate.PremiseSet.add
                    ([], original_rule_id, original_num)
                    ll_premise_set
                | l_preds ->
                  List.fold_left
                    (fun acc l_premises ->
                       Datalog.Predicate.PremiseSet.add (List.rev_map rename_predicate l_premises, original_rule_id, original_num) acc)
                    ll_premise_set
                    l_preds)
             premises_set
             Datalog.Predicate.PremiseSet.empty in
         Datalog.Predicate.PredicateMap.add (rename_predicate pred) new_premises l_derivations)
    derivations
    Datalog.Predicate.PredicateMap.empty