pivot.ml

(*************************************************************************)
(*  This file is part of Colibri2.                                       *)
(*                                                                       *)
(*  Copyright (C) 2014-2021                                              *)
(*    CEA   (Commissariat à l'énergie atomique et aux énergies           *)
(*           alternatives)                                               *)
(*                                                                       *)
(*  you can redistribute it and/or modify it under the terms of the GNU  *)
(*  Lesser General Public License as published by the Free Software      *)
(*  Foundation, version 2.1.                                             *)
(*                                                                       *)
(*  It is distributed in the hope that it will be useful,                *)
(*  but WITHOUT ANY WARRANTY; without even the implied warranty of       *)
(*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the        *)
(*  GNU Lesser General Public License for more details.                  *)
(*                                                                       *)
(*  See the GNU Lesser General Public License version 2.1                *)
(*  for more details (enclosed in the file licenses/LGPLv2.1).           *)
(*************************************************************************)

let debug =
  Debug.register_info_flag ~desc:"for the normalization by pivoting" "LRA.pivot"

module WithUnsolved (P : sig
  type t

  val name : string

  include Colibri2_popop_lib.Popop_stdlib.Datatype with type t := t

  val of_one_node : Node.t -> t

  val is_one_node : t -> Node.t option

  val subst : t -> Node.t -> t -> t option

  val normalize : t -> f:(Node.t -> t) -> t

  type data

  val nodes : t -> data Node.M.t

  type info

  val info : _ Egraph.t -> t -> info

  val attach_info_change :
    Egraph.wt -> (Egraph.rt -> Node.t -> Events.enqueue) -> unit

  val solve : info -> info -> (Node.t * t) option

  val set : Egraph.wt -> Node.t -> t -> unit
end) : sig
  val assume_equality : Egraph.wt -> Node.t -> P.t -> unit

  val init : Egraph.wt -> unit

  val get_repr : _ Egraph.t -> Node.t -> P.t option

  val iter_eqs : _ Egraph.t -> Node.t -> f:(P.t -> unit) -> unit

  val attach_repr_change :
    _ Egraph.t -> ?node:Node.t -> (Egraph.wt -> Node.t -> unit) -> unit

  val attach_eqs_change :
    _ Egraph.t -> ?node:Node.t -> (Egraph.wt -> Node.t -> unit) -> unit
end = struct
  type t = {
    repr : P.t;
        (** When a pivot is not possible because the equality can be null, the other
            product are waiting on the side, they are also normalized *)
    eqs : P.S.t;
  }
  [@@deriving ord, eq]

  let pp fmt t = Fmt.pf fmt "%a,%a" P.pp t.repr P.S.pp t.eqs

  let dom =
    Dom.Kind.create
      (module struct
        type nonrec t = t

        let name = P.name
      end)

  let used_in_poly : Node.S.t Node.HC.t =
    Node.HC.create Node.S.pp (P.name ^ "_used_in_poly")

  let set_poly d cl p chg =
    Egraph.set_dom d dom cl p;
    List.iter (fun p -> P.set d cl p) chg

  let add_used_product d cl' new_cls =
    Node.M.iter
      (fun used _ ->
        Node.HC.change
          (function
            | Some b -> Some (Node.S.add cl' b)
            | None -> (
                match Egraph.get_dom d dom used with
                | None ->
                    (* If a used node have no polynome associated, we set it to
                        itself. This allows to be warned when this node is merged. It
                        is the reason why this module doesn't specifically wait for
                        representative change *)
                    Egraph.set_dom d dom used
                      { repr = P.of_one_node used; eqs = P.S.empty };
                    Some (Node.S.of_list [ cl'; used ])
                | Some p ->
                    assert (P.equal (P.of_one_node used) p.repr);
                    assert false))
          used_in_poly d used)
      new_cls

  let add_used_t d cl' t =
    add_used_product d cl' (P.nodes t.repr);
    P.S.iter (fun p -> add_used_product d cl' (P.nodes p)) t.eqs

  let subst_doms d cl p =
    let b =
      match Node.HC.find used_in_poly d cl with
      | exception Not_found -> Node.S.empty
      | b -> b
    in
    Node.S.iter
      (fun cl' ->
        match Egraph.get_dom d dom cl' with
        | None -> assert false (* absurd: can't be used and absent *)
        | Some q ->
            let fold (new_cl, acc, chg) (q : P.t) =
              let new_cl =
                Node.M.set_union new_cl
                  (Node.M.set_diff (P.nodes p) (P.nodes q))
              in
              match P.subst q cl p with
              | None -> (new_cl, P.S.add q acc, chg)
              | Some q -> (new_cl, P.S.add q acc, q :: chg)
            in
            let new_cl, acc, chg = fold (Node.M.empty, P.S.empty, []) q.repr in
            let repr = P.S.choose acc in
            let new_cl, acc, chg =
              P.S.fold_left fold (new_cl, acc, chg) q.eqs
            in
            let eqs = P.S.remove repr acc in
            add_used_product d cl' new_cl;
            set_poly d cl' { repr; eqs } chg)
      b

  let norm_product d p =
    P.normalize p ~f:(fun cl ->
        let cl = Egraph.find d cl in
        match Egraph.get_dom d dom cl with
        | None -> P.of_one_node cl
        | Some p -> p.repr)

  module Th = struct
    let merged v1 v2 =
      Base.phys_equal v1 v2
      ||
      match (v1, v2) with
      | None, None -> true
      | Some v', Some v -> equal v' v
      | _ -> false

    let norm_dom cl = function
      | None ->
          let r = P.of_one_node cl in
          { repr = r; eqs = P.S.empty }
      | Some p -> p

    let add_itself d cl norm =
      add_used_t d cl norm;
      Egraph.set_dom d dom cl norm

    let rec merge d (_, cl1) (_, cl2) _inv =
      let cl1 = Egraph.find d cl1 in
      let cl2 = Egraph.find d cl2 in
      assert (not (Egraph.is_equal d cl1 cl2));
      merge_aux d cl1 cl2

    and merge_aux d cl1 cl2 =
      let p1o = Egraph.get_dom d dom cl1 in
      let p2o = Egraph.get_dom d dom cl2 in
      assert (not (Option.is_none p1o && Option.is_none p2o));
      match (p1o, p2o) with
      | None, None -> assert false (* absurd: no need to merge *)
      | Some p, None ->
          assert (Option.is_none (Node.HC.find_opt used_in_poly d cl2));
          add_itself d cl2 p
      | None, Some p ->
          assert (Option.is_none (Node.HC.find_opt used_in_poly d cl1));
          add_itself d cl1 p
      | Some p1, Some p2 -> (
          match
            solve d
              (Base.List.cartesian_product
                 (part d (p1.repr :: P.S.elements p1.eqs))
                 (part d (p2.repr :: P.S.elements p2.eqs)))
          with
          | `Solved ->
              (* The domains have been subsituted, and possibly recursively *)
              merge_aux d cl1 cl2
          | `Not_solved ->
              (* nothing to solve *)
              let repr =
                match (P.is_one_node p1.repr, P.is_one_node p2.repr) with
                | None, None -> p1.repr (* arbitrary *)
                | Some _, None -> p1.repr
                | None, Some _ -> p2.repr
                | Some cl1', Some cl2' ->
                    assert (Node.equal cl1' cl2');
                    p1.repr
              in
              let eqs =
                p1.eqs |> P.S.add p1.repr |> P.S.union p2.eqs |> P.S.add p2.repr
                |> P.S.remove repr
              in
              let p = { repr; eqs } in
              Egraph.set_dom d dom cl1 p;
              Egraph.set_dom d dom cl2 p)

    and merge_one_new_eq d cl eq =
      let po = Egraph.get_dom d dom cl in
      match po with
      | None ->
          add_used_product d cl (P.nodes eq);
          set_poly d cl { repr = eq; eqs = P.S.empty } [ eq ]
      | Some p -> (
          if (not (P.S.mem eq p.eqs)) && not (P.equal eq p.repr) then
            match
              solve d
                (Base.List.cartesian_product
                   (part d (p.repr :: P.S.elements p.eqs))
                   (part d [ eq ]))
            with
            | `Solved ->
                (* The domains have been substituted, and possibly recursively *)
                let eq = norm_product d eq in
                merge_one_new_eq d cl eq
            | `Not_solved ->
                (* nothing to solve *)
                let repr = p.repr in
                let eqs = p.eqs |> P.S.add eq |> P.S.remove repr in
                let p = { repr; eqs } in
                add_used_product d cl (P.nodes eq);
                set_poly d cl p [ eq ])

    and subst d cl eq =
      Debug.dprintf4 debug "[Pivot] subst %a with %a" Node.pp cl P.pp eq;
      let po = Egraph.get_dom d dom cl in
      match po with
      | None ->
          let p = { repr = eq; eqs = P.S.empty } in
          add_used_product d cl (P.nodes eq);
          set_poly d cl p [ eq ]
      | Some p ->
          assert (P.equal p.repr (P.of_one_node cl));
          subst_doms d cl eq;
          assert (
            P.equal eq (Base.Option.value_exn (Egraph.get_dom d dom cl)).repr);
          merge_one_new_eq d cl eq

    and part d l = List.map (fun p -> P.info d p) l

    and solve d l =
      let exception Solved of Node.t * P.t in
      let criteria i1 i2 =
        let aux i1 i2 =
          match P.solve i1 i2 with
          | None -> ()
          | Some (n, p) -> raise (Solved (n, p))
        in
        aux i1 i2;
        aux i2 i1
      in
      match List.iter (fun (a, b) -> criteria a b) l with
      | exception Solved (n, p) ->
          subst d n p;
          `Solved
      | () -> `Not_solved

    let key = dom

    type nonrec t = t

    let pp = pp
  end

  let () = Dom.register (module Th)

  let get_repr d n =
    let open CCOpt in
    let+ p = Egraph.get_dom d dom n in
    p.repr

  let iter_eqs d n ~f =
    match Egraph.get_dom d dom n with
    | None -> ()
    | Some p ->
        f p.repr;
        P.S.iter f p.eqs

  let assume_equality d n (p : P.t) =
    let n = Egraph.find d n in
    let p = norm_product d p in
    Th.merge_one_new_eq d n p

  module ChangePos = struct
    type runable = Node.S.t

    let print_runable = Node.S.pp

    let run d ns =
      Node.S.iter
        (fun n ->
          let p = Base.Option.value_exn (Egraph.get_dom d dom n) in
          let l = Th.part d (p.repr :: P.S.elements p.eqs) in
          let l = Base.List.cartesian_product l l in
          ignore (Th.solve d l))
        ns

    let delay = Events.Delayed_by 10

    let key =
      Events.Dem.create
        (module struct
          type t = Node.S.t

          let name = "Dom_product.ChangePos"
        end)

    let init d =
      P.attach_info_change d (fun d n ->
          match Node.HC.find_opt used_in_poly d n with
          | Some ns -> Events.EnqRun (key, ns, None)
          | None -> Events.EnqAlready)
  end

  let () = Events.register (module ChangePos)

  let init d = ChangePos.init d

  let attach_eqs_change d ?node f =
    match node with
    | Some x -> Daemon.attach_dom d x dom f
    | None -> Daemon.attach_any_dom d dom f

  let attach_repr_change = attach_eqs_change
end