Source file unique_binding.ml

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

module PredTable =
  DatalogLib.Datalog_AbstractSyntax.AbstractSyntax.Predicate.PredIdTable

module PredIds =
  DatalogLib.Datalog_AbstractSyntax.AbstractSyntax.Predicate.PredIds

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

module RgNodeMap = Map.Make (Rgg.Rg_graph.V)

let new_id_for_adorned_pred ~intensional:is_intensional (pred, adornment)
    old_table table =
  let pred_sym =
    ASPred.PredIdTable.find_sym_from_id pred.ASPred.p_id old_table
  in
  let () = Log.debug (fun m -> m "Dealing with pred '%s'" pred_sym) in
  let ad_pred =
    if is_intensional then
      Printf.sprintf "%s_%s" pred_sym
        (Adornment.to_string adornment)
    else pred_sym
  in
  ASPred.PredIdTable.add_sym ad_pred table

type ad_pred_key = {
  key_rule : int;
  key_head_ad : Adornment.status list;
  key_pred : ASPred.pred_id; (* Probably useless *)
  key_pos : int;
}

module AdPredMap = Map.Make (struct
  type t = ad_pred_key

  let compare k1 k2 =
    match k1.key_rule - k2.key_rule with
    | 0 -> (
        match Adornment.compare k1.key_head_ad k2.key_head_ad with
        | 0 -> k1.key_pos - k2.key_pos
        | r -> r)
    | r -> r
end)

module AdornmentTrie = struct
  exception Not_found
  exception Conflict

  type 'a t = Node of ('a option * 'a alternative)
  and 'a alternative = 'a t option * 'a t option
  (* Free on the left, and Bound on the right *)

  type key = Adornment.status list

  let empty = Node (None, (None, None))

  let rec add ?(overwrite = false) k v t =
    match (k, t) with
    | [], Node (None, alt) -> Node (Some v, alt)
    | [], Node (Some _, alt) when overwrite -> Node (Some v, alt)
    | [], Node (Some _, _) -> raise Conflict
    | Adornment.Free :: tl, Node (n, (None, right)) ->
        Node (n, (Some (add ~overwrite tl v empty), right))
    | Adornment.Free :: tl, Node (n, (Some child, right)) ->
        Node (n, (Some (add ~overwrite tl v child), right))
    | Adornment.Bound :: tl, Node (n, (left, None)) ->
        Node (n, (left, Some (add ~overwrite tl v empty)))
    | Adornment.Bound :: tl, Node (n, (left, Some child)) ->
        Node (n, (left, Some (add ~overwrite tl v child)))

  let rec find key t =
    match (key, t) with
    | [], Node (None, _) -> raise Not_found
    | [], Node (Some v, _) -> v
    | Adornment.Free :: _, Node (_, (None, _)) -> raise Not_found
    | Adornment.Free :: tl, Node (_, (Some t', _)) -> find tl t'
    | Adornment.Bound :: _, Node (_, (_, None)) -> raise Not_found
    | Adornment.Bound :: tl, Node (_, (_, Some t')) -> find tl t'

  let rec fold_aux f key acc t =
    match t with
    | Node (None, (None, None)) -> acc
    | Node (Some v, (None, None)) -> f (List.rev key) v acc
    | Node (None, (Some left, None)) ->
        fold_aux f (Adornment.Free :: key) acc left
    | Node (None, (None, Some right)) ->
        fold_aux f (Adornment.Bound :: key) acc right
    | Node (None, (Some left, Some right)) ->
        fold_aux f
          (Adornment.Bound :: key)
          (fold_aux f (Adornment.Free :: key) acc left)
          right
    | Node (Some v, (Some left, None)) ->
        fold_aux f (Adornment.Free :: key) (f (List.rev key) v acc) left
    | Node (Some v, (None, Some right)) ->
        fold_aux f (Adornment.Bound :: key) (f (List.rev key) v acc) right
    | Node (Some v, (Some left, Some right)) ->
        fold_aux f
          (Adornment.Bound :: key)
          (fold_aux f (Adornment.Free :: key) (f (List.rev key) v acc) left)
          right

  let fold f acc t = fold_aux f [] acc t

  let iter f t = fold (fun k v () -> f k v) () t

  let pp ?sep:_ _pp_a _fmt _t = failwith "Not implemented"
end

type data = {
  adorned_prog : ASProg.program;
  unadorned_map : ASPred.pred_id ASPred.PredIdMap.t;
  (* map from the id of the adorned predicate (the predicate of the
     unique binding program) to the id of the original predicate it
     replaces *)
  rule_to_rule_map : (int * int) Datalog_AbstractSyntax.RuleIdMap.t;
  (* map from the id of the new rule to the pair
     consisting of the id of the rule it replaces and the
     i_rhs_num *)
  adornments : Adornment.status list AdPredMap.t;
  (* a map to keep track of the adornments of the
     predicates in a given rule with a given adornment
     (avoids parsing the predicate symbol) *)
  pred_map : ASPred.pred_id AdornmentTrie.t ASPred.PredIdMap.t;
      (* map from original predicate ids to adornments to id of the
         corresponding unique binding predicate *)
}

let find_goal Rgg.{ rule; position; _ } =
  match
    List.find_opt (fun (_, pos) -> pos = position + 1) rule.ASRule.e_rhs
  with
  | Some res -> res
  | None -> List.find (fun (_, pos) -> pos = position + 1) rule.ASRule.i_rhs

let extend_map_to_trie p_id (ad, new_p_id) m =
  match ASPred.PredIdMap.find_opt p_id m with
  | None -> ASPred.PredIdMap.add p_id AdornmentTrie.(add ad new_p_id empty) m
  | Some t ->
      ASPred.PredIdMap.add p_id
        AdornmentTrie.(add ~overwrite:true ad new_p_id t)
        m

let rec get_rule_rhs_aux ~from:r_n node graph prog
    ((e_body, i_body, data_acc) as acc) =
  match node with
  | Rgg.Goal (goal_pred, goal_ad) ->
      (* The current node is a [Goal] node, hence it will generate a
         predicate at position [r_n.position+1] in the rule *)
      (* Find an id (possibly create one) for this goal *)
      let is_intensional =
        ASPred.(PredIds.mem goal_pred.p_id prog.ASProg.i_preds)
      in
      let new_goal_id, new_table =
        new_id_for_adorned_pred ~intensional:is_intensional (goal_pred, goal_ad)
          prog.ASProg.pred_table data_acc.adorned_prog.ASProg.pred_table
      in
      let () = Log.debug (fun m -> m "Intensional? %b" is_intensional) in
      let new_pred_map =
        extend_map_to_trie goal_pred.ASPred.p_id (goal_ad, new_goal_id)
          data_acc.pred_map
      in
      (* We need to recover the argument characteristics from the rule
         being processed *)
      let goal_in_rule, goal_position = find_goal r_n in
      let () = Log.debug (fun m -> m "At position '%d'" goal_position) in
      let () = assert (ASPred.(goal_pred.p_id = goal_in_rule.p_id)) in
      let new_goal_pred = ASPred.{ goal_in_rule with p_id = new_goal_id } in
      (* Record the mapping from [new_body_pred] to [body_pred] *)
      let new_unadorned_map =
        ASPred.(PredIdMap.add new_goal_id goal_pred.p_id data_acc.unadorned_map)
      in
      (* First update the program with the new predicate table *)
      let new_prog =
        { data_acc.adorned_prog with ASProg.pred_table = new_table }
      in
      if is_intensional then
        (* This new adorned predicate has to be added to the
           intensional predicates of the new program *)
        let new_prog' =
          {
            new_prog with
            ASProg.i_preds =
              ASPred.PredIds.add new_goal_id new_prog.ASProg.i_preds;
          }
        in
        ( e_body,
          (new_goal_pred, goal_position) :: i_body,
          {
            data_acc with
            adorned_prog = new_prog';
            unadorned_map = new_unadorned_map;
            pred_map = new_pred_map;
          } )
      else
        ( (new_goal_pred, goal_position) :: e_body,
          i_body,
          {
            data_acc with
            adorned_prog = new_prog;
            unadorned_map = new_unadorned_map;
            pred_map = new_pred_map;
          } )
  | Rgg.Rule new_r_n when new_r_n.Rgg.position = r_n.Rgg.position + 1 ->
      (* This rule node is going to be processed recursively on each of its successor nodes *)
      (* TODO: tail recursion is possible because a rule node has 0 or 2 successors *)
      Rgg.Rg_graph.fold_succ
        (fun n l_acc -> get_rule_rhs_aux ~from:new_r_n n graph prog l_acc)
        graph node acc
  | Rgg.Rule _ -> failwith "Bug: wrong from_position for rule"

let get_rule_rhs ~from:r_n graph prog acc =
  Rgg.Rg_graph.fold_succ
    (fun n l_acc -> get_rule_rhs_aux ~from:r_n n graph prog l_acc)
    graph (Rgg.Rule r_n) acc

let derive_program graph prog =
  let () =
    Log.debug (fun m -> m "Going to derive program for:@,@[<v>  @[%a@]@]" 
                  (ASProg.pp ~with_position:false ~with_id:true) prog)
  in
  let new_prog =
    ASProg.
      {
        rules = ASRule.Rules.empty;
        pred_table = ASPred.PredIdTable.empty;
        const_table = prog.const_table;
        i_preds = ASPred.PredIds.empty;
        rule_id_gen = UtilsLib.IdGenerator.IntIdGen.init ();
        head_to_rules = ASPred.PredIdMap.empty;
        e_pred_to_rules = ASPred.PredIdMap.empty;
      }
  in
  Rgg.Rg_graph.fold_vertex
    (fun node l_data ->
      match node with
      | Rgg.Rule _ ->
          (* If it's a rule node, there is nothing to do. We only
             consider goals *)
          l_data
      | Rgg.Goal (head_pred_sym, _)
        when not ASPred.(PredIds.mem head_pred_sym.p_id prog.ASProg.i_preds) ->
          (* If the goal node is extensional, there is nothing to do
             either. It will be dealt with when constructing the rules
             for intensional predicates *)
          l_data
      | Rgg.Goal (head_pred, head_adornment) ->
          (* We are dealing with an intensional predicate.*)

          (* We can generate a new id and a new symbol table for the
             adorned head *)
          let head_id, new_table =
            new_id_for_adorned_pred ~intensional:true
              (head_pred, head_adornment)
              prog.ASProg.pred_table l_data.adorned_prog.ASProg.pred_table
          in
          let new_pred_map =
            extend_map_to_trie head_pred.ASPred.p_id (head_adornment, head_id)
              l_data.pred_map
          in
          let new_unadorned_map =
            ASPred.PredIdMap.add head_id head_pred.ASPred.p_id
              l_data.unadorned_map
          in

          (* We then check all successors of this node in order to find
             those corresponding to rules it is a head of *)
          Rgg.Rg_graph.fold_succ
            (fun vertex l_acc ->
              match vertex with
              | Rgg.Goal _ ->
                  failwith
                    "A goal node should not be the successor of a goal node"
              | Rgg.Rule r_n when r_n.Rgg.position <> 0 ->
                  (* If the successor is a rule node with a position
                     different from 0, then the goal g is not the head
                     of the rule r_n, so it is bypassed *)
                  l_acc
              | Rgg.Rule r_n ->
                  (* The head predicate is assigned the parameters of the r_n to be processed *)
                  let () = assert (head_adornment = r_n.Rgg.adorned_head) in

                  let () =
                    Log.debug (fun m ->
                        m "Extracting rule:@,@[  @[%a@]@]"
                          (ASRule.pp ~with_position:true
                             ~with_id:true prog.ASProg.pred_table
                             prog.ASProg.const_table) r_n.Rgg.rule)
                  in
                  let head_pred =
                    ASPred.{ r_n.Rgg.rule.ASRule.lhs with p_id = head_id }
                  in

                  (* The rule node describes a rule the goal g is a head
                     of. We fetch the rhs following this node *)
                  let e_rhs, i_rhs, l_acc' =
                    get_rule_rhs ~from:r_n graph prog ([], [], l_acc)
                  in

                  (* Generate a new id for the rule *)
                  let new_rule_id, new_rule_id_gen =
                    UtilsLib.IdGenerator.IntIdGen.get_fresh_id
                      l_acc'.adorned_prog.ASProg.rule_id_gen
                  in
                  let new_prog =
                    {
                      l_acc'.adorned_prog with
                      ASProg.rule_id_gen = new_rule_id_gen;
                    }
                  in
                  (* We sort the (pred,position) lists *)
                  let sorted_e_rhs =
                    List.sort (fun (_, pos1) (_, pos2) -> pos1 - pos2) e_rhs
                  in
                  let sorted_i_rhs =
                    List.sort (fun (_, pos1) (_, pos2) -> pos1 - pos2) i_rhs
                  in
                  (* TODO: Est-ce vraiment nécessaire *)
                  let i_length = List.length sorted_i_rhs in
                  let () =
                    Log.debug (fun m ->
                        m "i_length = %d, i_rhs_num = %d" i_length
                          r_n.Rgg.rule.ASRule.i_rhs_num)
                  in
                  let () = assert (i_length = r_n.Rgg.rule.ASRule.i_rhs_num) in

                  let new_rule =
                    ASRule.
                      {
                        id = new_rule_id;
                        lhs = head_pred;
                        e_rhs = sorted_e_rhs;
                        i_rhs = sorted_i_rhs;
                        i_rhs_num = i_length;
                        rhs_num = r_n.Rgg.rule.rhs_num;
                      }
                  in

                  (* We add the rule to the program *)
                  let new_prog' =
                    ASProg.
                      {
                        new_prog with
                        rules = ASRule.Rules.add new_rule new_prog.rules;
                        head_to_rules =
                          ASRule.extend_head_id_map_to_rules head_id new_rule
                            new_prog.head_to_rules;
                      }
                  in
                  {
                    l_acc' with
                    adorned_prog = new_prog';
                    rule_to_rule_map =
                      Datalog_AbstractSyntax.RuleIdMap.add new_rule_id
                        (r_n.Rgg.rule.ASRule.id, i_length)
                        l_acc'.rule_to_rule_map;
                  })
            graph node
            {
              l_data with
              adorned_prog =
                {
                  l_data.adorned_prog with
                  ASProg.pred_table = new_table;
                  ASProg.i_preds =
                    ASPred.PredIds.add head_id
                      l_data.adorned_prog.ASProg.i_preds;
                };
              pred_map = new_pred_map;
              unadorned_map = new_unadorned_map;
            })
    graph
    {
      adorned_prog = new_prog;
      unadorned_map = ASPred.PredIdMap.empty;
      rule_to_rule_map = Datalog_AbstractSyntax.RuleIdMap.empty;
      adornments = AdPredMap.empty;
      pred_map = ASPred.PredIdMap.empty;
    }