Source file datalog_AbstractSyntax.ml

open UtilsLib.IdGenerator
open UtilsLib.Utils
module RuleIdMap = IntMap

module Var = struct
  type t = Var of int

  let compare (Var i) (Var j) = i - j
  let succ (Var i) = Var (i + 1)
  let start = Var 0

  let to_string (Var i) =
    let dec = decompose ~input:i ~base:26 in
    List.fold_left
      (fun acc i -> Printf.sprintf "%s%c" acc (char_of_int (97 + i)))
      "" dec

  let pp m (Var i) =
    let dec = decompose ~input:i ~base:26 in
    UtilsLib__Utils.pp_list ~sep:""
      (fun m i -> Format.fprintf m "%c" (char_of_int (97 + i)))
      m dec
end

module VarGen = IdGen (Var)

module Const = struct
  type t = Const of int

  let compare (Const i) (Const j) = i - j
  let start = Const 0
  let succ (Const i) = Const (i + 1)
  let to_string (Const i) = string_of_int i
  let pp m (Const i) = Format.pp_print_int m i
end

module ConstGen = IdGen (Const)

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

module AbstractSyntax = struct
  (** These modules are the abstract syntactic representations of the
      predicates, of the rules, and of the programs *)
  module Predicate = struct
    type term = Var of VarGen.id | Const of ConstGen.id

    module VarMap = Map.Make (Var)

    let map_content_compare (k1, map1) (k2, map2) =
      try
        let val1 = VarMap.find k1 map1 in
        try val1 - VarMap.find k2 map2 with Not_found -> 1
      with Not_found -> (
        try
          let _ = VarMap.find k2 map2 in
          -1
        with Not_found -> 0)

    let term_compare l1 l2 =
      let rec term_compare_aux l1 l2 (l1_vars, l2_vars) pos =
        match (l1, l2) with
        | [], [] -> 0
        | [], _ -> -1
        | _, [] -> 1
        | Const _ :: _, Var _ :: _ -> 1
        | Var _ :: _, Const _ :: _ -> -1
        | Const a1 :: tl1, Const a2 :: tl2 ->
            let res = ConstGen.compare a1 a2 in
            if ConstGen.compare a1 a2 <> 0 then res
            else term_compare_aux tl1 tl2 (l1_vars, l2_vars) (pos + 1)
        | Var a1 :: tl1, Var a2 :: tl2 ->
            let res = map_content_compare (a1, l1_vars) (a2, l2_vars) in
            if VarGen.compare a1 a2 <> 0 then res
            else
              term_compare_aux tl1 tl2
                (VarMap.add a1 pos l1_vars, VarMap.add a2 pos l2_vars)
                (pos + 1)
      in
      term_compare_aux l1 l2 (VarMap.empty, VarMap.empty) 0

    let pp_term cst_id_table fmt = function
      | Var v -> Var.pp fmt v
      | Const c ->
          Format.fprintf fmt "%s"
            (ConstGen.Table.find_sym_from_id c cst_id_table)

    type pred_id = int

    module PredIdMap = IntIdGen.IdMap
    module PredIdTable = IntIdGen.Table

    type predicate = {
      p_id : pred_id;
      arity : int;
      arguments : term list;
          (** It is assumed that the size of the list is the
		       arity *)
    }

    let pp_pred_id = Format.pp_print_int

let pp ?position ?(with_id = false) ?(with_arity = false) pred_id_table cst_id_table fmt predicate =
  let pp_id fmt id =
    match with_id with
    | false -> ()
    | true -> Format.fprintf fmt "[p_id = %d]" id in
  let pp_position fmt position =
    match position with
    | None -> ()
    | Some i -> Format.fprintf fmt "(* position: %d*)" i
  in
  let pp_arity fmt arity =
    match with_arity with
    | false -> ()
    | true -> Format.fprintf fmt "/%d" arity
  in
      Format.fprintf fmt "@[<h>@[%s%a%a(@[%a@])%a@]@]"
        (PredIdTable.find_sym_from_id predicate.p_id pred_id_table)
        pp_arity
        predicate.arity
        pp_id
        predicate.p_id
        (pp_list ~sep:"," (pp_term cst_id_table))
        predicate.arguments
        pp_position
        position

    let compare ?(with_arguments = true)
        ({ p_id = id1; arity = a1; arguments = l1 } : predicate)
        ({ p_id = id2; arity = a2; arguments = l2 } : predicate) =
      let res = id1 - id2 in
      if res <> 0 then res
      else
        let res = a1 - a2 in
        if res <> 0 then res
        else if with_arguments then term_compare l1 l2
        else res

    let fake_pred_id = -1

    module PredIds = IntSet

    module TermSet = Set.Make (struct
      type t = term

      let compare term1 term2 =
        match (term1, term2) with
        | Var id1, Var id2 -> VarGen.compare id1 id2
        | Const id1, Const id2 -> ConstGen.compare id1 id2
        | Var _, Const _ -> 1
        | Const _, Var _ -> -1
    end)

    let pp_terms cst_id_table fmt set =
      let first = ref true in
      TermSet.iter
        (fun term ->
           let () =
             if !first then
               pp_term cst_id_table fmt term
             else
               Format.fprintf fmt ",%a" (pp_term cst_id_table) term
           in first := false)
        set

    let get_variables pred =
      let is_variable arg = match arg with Var _ -> true | _ -> false in
      TermSet.of_list (List.filter is_variable pred.arguments)

    (** [variables_of_pred_aux acc pred] returns [acc'] where [acc']
       is the union of the set of arguments (terms) of [pred] with
       [acc] *)
    let free_variables_of_pred_aux acc pred =
      List.fold_left
        (fun acc' term ->
          match term with Var _ -> TermSet.add term acc' | _ -> acc')
        acc pred.arguments

    let free_variables_of_preds_aux acc preds =
      List.fold_left free_variables_of_pred_aux acc preds

    let get_variables_of_preds preds =
      let new_free_variable_set =
        free_variables_of_preds_aux TermSet.empty preds
      in
      new_free_variable_set

    let copy_predicate ~new_id:id pred = { pred with p_id = id }

    let pp_predids pred_table fmt p_ids =
      PredIds.iter
        (fun elt ->
          Format.fprintf fmt "%s@,"
            (PredIdTable.find_sym_from_id elt pred_table))
        p_ids
  end

  module Proto_Rule = struct
    type t = {
      proto_id : int;
      proto_lhs : Predicate.predicate;
      proto_rhs : Predicate.predicate list;
          (** represents the predicates of the rule *)
    }

    let pp pred_id_table cst_id_table fmt r =
      let pp_head fmt lhs = Predicate.pp pred_id_table cst_id_table fmt lhs in
      let pp_tail fmt rhs =
        match rhs with
        | [] -> Format.fprintf fmt "."
        | _ ->
            Format.fprintf fmt ":- %a."
              (pp_list ~sep:"," (Predicate.pp pred_id_table cst_id_table))
              rhs
      in
      Format.fprintf fmt "%a%a" pp_head r.proto_lhs pp_tail r.proto_rhs
  end

  module Rule = struct
    type rule = {
      id : int;
      lhs : Predicate.predicate;
      e_rhs : (Predicate.predicate * int) list;
      i_rhs : (Predicate.predicate * int) list;
      i_rhs_num : int;
      rhs_num : int;
    }

    let pp ?(with_position = false) ?(with_id = false)
        ?(with_arity = false) pred_id_table cst_id_table fmt r =
      let pp_head = Predicate.pp ~with_id ~with_arity pred_id_table cst_id_table in
      let vdash, e_i_sep =
        match (r.e_rhs, r.i_rhs) with
        | [], [] -> (".", format_of_string "")
        | [], _ -> (":-", format_of_string "")
        | _, [] -> (":-", format_of_string "")
        | _, _ -> (":-", format_of_string" ,@,")
      in
      let pp_pred fmt (pred, pos) =
        match with_position with
        | false -> Predicate.pp  ~with_id ~with_arity pred_id_table cst_id_table fmt pred
        | true -> Predicate.pp ~position:pos ~with_id ~with_arity pred_id_table cst_id_table fmt pred
      in
      let pp_rule_id fmt id = if with_id then Format.fprintf fmt "(id: %6d)@;" id else () in
      Format.fprintf
        fmt
        ("@[%a@]@;%a %s@[ @[<hv>%a"^^e_i_sep^^"%a@]@]")
        pp_rule_id
        r.id
        pp_head
        r.lhs
        vdash
        (pp_list ~sep:",@," pp_pred)
        r.e_rhs
        (pp_list ~sep:",@," pp_pred)
        r.i_rhs

    module Rules = Set.Make (struct
      type t = rule

      let compare { id = i; _ } { id = j; _ } = i - j
    end)

    module RuleMap = Map.Make (struct
      type t = rule

      let compare { id = i; _ } { id = j; _ } = i - j
    end)

    let extend_head_id_map_to_rules id r m =
      match Predicate.PredIdMap.find_opt id m with
      | None -> Predicate.PredIdMap.add id Rules.(singleton r) m
      | Some rules -> Predicate.PredIdMap.add id Rules.(add r rules) m

    let ids_to_rules ids id_to_rule_map =
      IntSet.fold
        (fun e acc -> Rules.add (IntMap.find e id_to_rule_map) acc)
        ids Rules.empty

    let pp_rules ?(with_position = false) ?(with_id = false) pred_id_table cst_id_table fmt rules =
      let pp_rule = pp ~with_position ~with_id pred_id_table cst_id_table in
      Rules.iter (fun r -> Format.fprintf fmt "%a@," pp_rule r) rules

    let init_split_rhs proto_preds intensional_pred =
      let i_num, i_p, e_p, length =
        List.fold_left
          (fun (i_num, i_preds, e_preds, i) ({ Predicate.p_id; _ } as pred) ->
            if Predicate.PredIds.mem p_id intensional_pred then
              (i_num + 1, (pred, i) :: i_preds, e_preds, i + 1)
            else (i_num, i_preds, (pred, i) :: e_preds, i + 1))
          (0, [], [], 1) proto_preds
      in
      (i_num, i_p, e_p, length - 1)

    let update_split_rhs init proto_preds intensional_pred =
      List.fold_left
        (fun (i_preds, e_preds) (({ Predicate.p_id; _ }, _) as pred) ->
          if Predicate.PredIds.mem p_id intensional_pred then
            (pred :: i_preds, e_preds)
          else (i_preds, pred :: e_preds))
        init proto_preds

    let extend_map_to_set k v map_to_set =
      let current_set =
        try Predicate.PredIdMap.find k map_to_set
        with Not_found -> IntSet.empty
      in
      Predicate.PredIdMap.add k (IntSet.add v current_set) map_to_set

    let proto_rule_to_rule proto_rule intensional_pred =
      let i_num, i_preds, e_preds, length =
        init_split_rhs proto_rule.Proto_Rule.proto_rhs intensional_pred
      in
      {
        id = proto_rule.Proto_Rule.proto_id;
        lhs = proto_rule.Proto_Rule.proto_lhs;
        e_rhs = List.rev e_preds;
        i_rhs = List.rev i_preds;
        i_rhs_num = i_num;
        rhs_num = length;
      }

    let update rule intensional_pred =
      let i_preds, e_preds =
        update_split_rhs (rule.i_rhs, []) rule.e_rhs intensional_pred
      in
      { rule with e_rhs = e_preds; i_rhs = i_preds }

    let set_new_id id rule = { rule with id }
    (* Replace
     **********
                                { lhs = rule.lhs
                                ; e_rhs = rule.e_rhs      
                                ; i_rhs = rule.i_rhs
                                ; i_rhs_num = rule.i_rhs_num
                                ; id
                                }
     *)

    let set_new_id_from_gen rule gen =
      let new_id, new_gen = IntIdGen.get_fresh_id gen in
      (set_new_id new_id rule, new_gen)

    let get_variables_in_rule rule =
      let head_result = Predicate.get_variables rule.lhs in
      let variables =
        List.fold_left
          (fun acc (pred, _) -> Predicate.free_variables_of_pred_aux acc pred)
          head_result rule.i_rhs
      in
      let new_result =
        List.fold_left
          (fun acc (pred, _) -> Predicate.free_variables_of_pred_aux acc pred)
          variables rule.e_rhs
      in
      new_result
    (* Replace
     **********************
       let i_rhs_result =
       Predicate.get_variables_of_preds (List.map fst rule.i_rhs)
       in
       let e_rhs_result =
       Predicate.get_variables_of_preds (List.map fst rule.e_rhs)
       in
       Predicate.TermSet.union head_result (Predicate.TermSet.union i_rhs_result e_rhs_result)
     *)

    let get_subgoal rule i =
      let new_result =
        match
          List.find_opt (fun (_pred, position) -> position = i + 1) rule.e_rhs
        with
        | Some s -> s
        | None -> (
            match
              List.find_opt
                (fun (_pred, position) -> position = i + 1)
                rule.i_rhs
            with
            | Some s -> s
            | None -> failwith (Printf.sprintf "Bug: No subgoal: %d" i))
      in
      new_result
    (* Replace *)
    (*****************)
    (*
             let rhs = rule.e_rhs @ rule.i_rhs in
             match List.find_opt (fun (_pred, position) -> position == i + 1) rhs with
             | Some s -> s
             | None -> failwith (Printf.sprintf "Error getting the subgoal : %d" i)
            *)
  end

  module Proto_Program = struct
    type t = {
      rules : Proto_Rule.t list;
      pred_table : Predicate.PredIdTable.table;
      const_table : ConstGen.Table.table;
      i_preds : Predicate.PredIds.t;
      rule_id_gen : IntIdGen.t;
      pred_to_rules : IntSet.t Predicate.PredIdMap.t;
    }

    type tables =
      Predicate.PredIdTable.table * (VarGen.Table.table * ConstGen.Table.table)

    let empty =
      {
        rules = [];
        pred_table = Predicate.PredIdTable.empty;
        const_table = ConstGen.Table.empty;
        i_preds = Predicate.PredIds.empty;
        rule_id_gen = IntIdGen.init ();
        pred_to_rules = Predicate.PredIdMap.empty;
      }

    let extension pred_table const_table rule_id_gen =
      {
        rules = [];
        pred_table;
        const_table;
        i_preds = Predicate.PredIds.empty;
        rule_id_gen;
        pred_to_rules = Predicate.PredIdMap.empty;
      }

    let add_proto_rule (f_lhs, f_rhs) prog =
      let rule_id, new_rule_id_gen = IntIdGen.get_fresh_id prog.rule_id_gen in
      let lhs, (new_pred_id_table, new_tables) =
        f_lhs (prog.pred_table, (VarGen.Table.empty, prog.const_table))
      in
      let rhs, (new_pred_id_table', (_, new_const_table)) =
        f_rhs (new_pred_id_table, new_tables)
      in
      let new_i_preds =
        match rhs with
        | [] -> prog.i_preds
        | _ -> Predicate.PredIds.add lhs.Predicate.p_id prog.i_preds
      in
      let new_rule =
        {
          Proto_Rule.proto_id = rule_id;
          Proto_Rule.proto_lhs = lhs;
          Proto_Rule.proto_rhs = rhs;
        }
      in
      {
        rules = new_rule :: prog.rules;
        pred_table = new_pred_id_table';
        const_table = new_const_table;
        i_preds = new_i_preds;
        rule_id_gen = new_rule_id_gen;
        pred_to_rules =
          List.fold_left
            (fun acc p -> Rule.extend_map_to_set p.Predicate.p_id rule_id acc)
            prog.pred_to_rules rhs;
      }
  end

  module Program = struct
    type program = {
      rules : Rule.Rules.t;
      pred_table : Predicate.PredIdTable.table;
      const_table : ConstGen.Table.table;
      i_preds : Predicate.PredIds.t;
      rule_id_gen : IntIdGen.t;
      head_to_rules : Rule.Rules.t Predicate.PredIdMap.t;
      e_pred_to_rules : Rule.Rules.t Predicate.PredIdMap.t;
    }

    type modifier = {
      modified_rules : Rule.Rules.t;
      new_pred_table : Predicate.PredIdTable.table;
      new_const_table : ConstGen.Table.table;
      new_i_preds : Predicate.PredIds.t;
      new_e_preds : Predicate.PredIds.t;
      new_rule_id_gen : IntIdGen.t;
    }

    let make_program
        {
          Proto_Program.rules;
          Proto_Program.pred_table;
          Proto_Program.const_table;
          Proto_Program.i_preds;
          Proto_Program.rule_id_gen;
          Proto_Program.pred_to_rules;
        } =
      let actual_rules, ids_to_rule_map, head_id_to_rules_map =
        List.fold_left
          (fun (acc, ids_to_rule_map, head_id_to_rules_map) p_rule ->
            let rule = Rule.proto_rule_to_rule p_rule i_preds in
            ( Rule.Rules.add rule acc,
              IntMap.add p_rule.Proto_Rule.proto_id rule ids_to_rule_map,
              Rule.extend_head_id_map_to_rules rule.Rule.lhs.Predicate.p_id rule
                head_id_to_rules_map ))
          (Rule.Rules.empty, IntMap.empty, Predicate.PredIdMap.empty)
          rules
      in
      {
        rules = actual_rules;
        pred_table;
        const_table;
        i_preds;
        rule_id_gen;
        head_to_rules = head_id_to_rules_map;
        e_pred_to_rules =
          Predicate.PredIdMap.fold
            (fun p rule_ids acc ->
              if Predicate.PredIds.mem p i_preds then
                Predicate.PredIdMap.remove p acc
              else
                Predicate.PredIdMap.add p
                  (Rule.ids_to_rules rule_ids ids_to_rule_map)
                  acc)
            pred_to_rules Predicate.PredIdMap.empty;
      }

    let extend prog
        {
          Proto_Program.rules;
          Proto_Program.pred_table;
          Proto_Program.const_table;
          Proto_Program.i_preds;
          Proto_Program.rule_id_gen;
          Proto_Program.pred_to_rules;
        } =
      let new_i_preds = Predicate.PredIds.union prog.i_preds i_preds in
      let updated_e_pred_map_to_r, updated_rules =
        (* all the rules that were pointed to by an extensional
           predicate that has to be turned into an intensional
           predicate because of the program extension have to be
           updated *)
        (* We check if some the new intensional predicates also are
           keys of the e_pred_to_rules map of the previous program *)
        Predicate.PredIds.fold
          (fun p_id ((e_p_to_r, rules_acc) as acc) ->
            try
              (* First we check wether this predicate was considered
                 as an extensional one *)
              let to_be_modified_rules =
                Predicate.PredIdMap.find p_id prog.e_pred_to_rules
              in
              (* If yes, we nee to remove it from the map *)
              ( Predicate.PredIdMap.remove p_id e_p_to_r,
                (* And to modify the rules it pointed to *)
                Rule.Rules.fold
                  (fun r acc ->
                    Rule.Rules.add
                      (Rule.update r new_i_preds)
                      (Rule.Rules.remove r acc))
                  to_be_modified_rules rules_acc )
            with (* If no, don't do anything *)
            | Not_found -> acc)
          i_preds
          (prog.e_pred_to_rules, prog.rules)
      in
      let new_rules, id_to_rule_map, new_head_to_rules =
        List.fold_left
          (fun (acc, id_to_rule_map, h_t_r) p_rule ->
            let rule = Rule.proto_rule_to_rule p_rule new_i_preds in
            ( Rule.Rules.add rule acc,
              IntMap.add p_rule.Proto_Rule.proto_id rule id_to_rule_map,
              Rule.extend_head_id_map_to_rules rule.Rule.lhs.Predicate.p_id rule
                h_t_r ))
          (updated_rules, IntMap.empty, prog.head_to_rules)
          rules
      in
      {
        rules = new_rules;
        pred_table;
        const_table;
        i_preds = new_i_preds;
        rule_id_gen;
        head_to_rules = new_head_to_rules;
        e_pred_to_rules =
          Predicate.PredIdMap.merge
            (fun _ opt_rule_ids opt_rules ->
              match (opt_rule_ids, opt_rules) with
              | None, None -> None
              | None, _ -> opt_rules
              | Some ids, None -> Some (Rule.ids_to_rules ids id_to_rule_map)
              | Some ids, Some rules ->
                  Some
                    (Rule.Rules.union rules
                       (Rule.ids_to_rules ids id_to_rule_map)))
            pred_to_rules updated_e_pred_map_to_r;
      }

    let is_in_idb pred prog = Predicate.(PredIds.mem pred.p_id prog.i_preds)

    let is_head (pred : Predicate.predicate) (rule : Rule.rule) =
      let head = rule.Rule.lhs in
      let () =
        Log.info (fun m ->
            m
              "testing pred_id/arity %d/%d agains rule %d with head \
               pred_id/arity %d/%d: %b and %b"
              pred.Predicate.p_id pred.Predicate.arity rule.Rule.id
              head.Predicate.p_id head.Predicate.arity
              Predicate.(head.p_id = pred.p_id)
              Predicate.(head.arity = pred.arity))
      in
      Predicate.(head.p_id = pred.p_id && head.arity = pred.arity)

    let match_rules (predicate : Predicate.predicate) program =
      (*
      let () = Log.info
          (fun m ->
             let () = m "Rules by head:%!" in
             Predicate.PredIdMap.iter
               (fun pred_id l_rules ->
                  Rule.Rules.iter
                    (fun r -> Log.info (fun m -> m "head: %d -> Rule: %s" pred_id (Rule.to_string ~with_position:true ~with_id:true r program.pred_table program.const_table)))
                    l_rules)
               program.head_to_rules) in
      let () = Log.info (fun m -> m "+++++++++++++++++++++++++++++++++++++++++++++++++++++++") in
      let () = Log.info
          (fun m ->
             let () = m "Program rules:%!" in
               Rule.Rules.iter
                 (fun r -> Log.info
                     (fun m ->
                        m
                          "head: %d -> Rule: %s"
                          r.Rule.lhs.Predicate.p_id
                          (Rule.to_string ~with_position:true ~with_id:true r program.pred_table program.const_table)))
                 program.rules) in *)
      let new_res =
        match
          Predicate.PredIdMap.find_opt predicate.Predicate.p_id
            program.head_to_rules
        with
        | None -> Rule.Rules.empty
        | Some set -> set
      in
      let new_res_ids =
        List.map (fun r -> r.Rule.id) (Rule.Rules.elements new_res)
      in
      let () =
        Log.info (fun m ->
            m "New matching rules are: %s"
              (UtilsLib.Utils.string_of_list ", "
                 (fun r -> string_of_int r.Rule.id)
                 (Rule.Rules.elements new_res)))
      in
      let () =
        Log.info (fun m ->
            let () = m "Matching new rules:" in
            Rule.Rules.iter
              (fun r ->
                Log.info (fun m ->
                    m "head: %d -> Rule: %a" r.Rule.lhs.Predicate.p_id
                      (Rule.pp ~with_position:true ~with_id:true 
                         program.pred_table program.const_table) r))
              new_res)
      in

      (* Replace *)
      let old_res = Rule.Rules.filter (is_head predicate) program.rules in
      let old_res_ids =
        List.map (fun r -> r.Rule.id) (Rule.Rules.elements old_res)
      in
      let () =
        Log.info (fun m ->
            m "Old matching rules are: %s"
              (UtilsLib.Utils.string_of_list ", "
                 (fun r -> string_of_int r.Rule.id)
                 (Rule.Rules.elements old_res)))
      in
      let () =
        assert (List.sort ( - ) new_res_ids = List.sort ( - ) old_res_ids)
      in
      let () = assert (Rule.Rules.equal new_res old_res) in
      (*new_res *)
      old_res

    let get_rule_by_id program id =
      let rules = program.rules in
      Rule.(Rules.find_first (fun x -> x.id = id) rules)

    let pp ?(with_position = false) ?(with_id = false) fmt prog =
      Format.fprintf fmt
        "@[<v>%a@]@,Intensional predicates are:@,@[<v>  @[<v>%a@]@]"
        (Rule.pp_rules ~with_position ~with_id prog.pred_table prog.const_table)
        prog.rules
        (Predicate.pp_predids prog.pred_table)
        prog.i_preds
    
  end
end