fourier.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).           *)
(*************************************************************************)

open Colibri2_popop_lib
open Colibri2_core

let comparisons = Datastructure.Push.create Ground.pp "LRA.fourier.comparisons"

let scheduled =
  Datastructure.Ref.create Base.Bool.pp "LRA.fourier.scheduled" false

let debug =
  Debug.register_info_flag ~desc:"Reasoning about <= < in LRA" "fourier"

let stats_run = Debug.register_stats_int ~name:"Fourier.run" ~init:0

type eq = { p : Polynome.t; bound : Bound.t; origins : Ground.S.t }
[@@deriving show]
(** p <= 0 or p < 0 *)

let divide d (p : Polynome.t) =
  try
    (match Polynome.is_cst p with None -> () | Some _ -> raise Exit);
    if Q.sign p.cst <> 0 then raise Exit;
    let l = Node.M.bindings p.poly in
    let l =
      List.fold_left
        (fun acc (e, q) ->
          match Dom_product.get_repr d e with
          | None when Egraph.is_equal d RealValue.zero e -> acc
          | None -> raise Exit
          | Some p -> (p, q) :: acc)
        [] l
    in
    Debug.dprintf4 debug "@[eq:%a@ %a@]" Polynome.pp p
      Fmt.(list ~sep:(any "+") (using CCPair.swap (pair Q.pp Product.pp)))
      l;
    match l with
    | [] -> raise Exit
    | (hd, _) :: _ ->
        let common =
          List.fold_left
            (fun acc (p, _) ->
              Node.M.inter
                (fun _ a b -> if Q.equal a b then Some a else None)
                acc p.Product.poly)
            hd.Product.poly l
        in
        let common, pos =
          Node.M.fold_left
            (fun (acc, pos) n v ->
              if Dom_interval.is_strictly_positive d n then
                (Node.M.add n v acc, pos)
              else if Dom_interval.is_strictly_negative d n then
                (Node.M.add n v acc, not pos)
              else (acc, pos))
            (Node.M.empty, true) common
        in
        if Node.M.is_empty common then raise Exit;
        Debug.dprintf4 debug "[Fourier] found possible division: %a / %a"
          Polynome.pp p (Node.M.pp Q.pp) common;
        let cst, l =
          List.fold_left
            (fun (cst, acc) (p, v) ->
              let p = Product.of_map (Node.M.set_diff p.Product.poly common) in

              let v = if pos then v else Q.neg v in
              match Product.classify p with
              | ONE -> (Q.add v cst, acc)
              | NODE n -> (cst, (n, v) :: acc)
              | PRODUCT ->
                  let n = Dom_product.node_of_product p in
                  (Q.add Q.one cst, (n, v) :: acc))
            (Q.zero, []) l
        in
        Some (Polynome.of_list cst l, common, pos)
  with Exit -> None

let delta d eq =
  (* add info to delta *)
  if Node.M.is_num_elt 2 eq.p.poly then
    match Node.M.bindings eq.p.poly with
    | [ (src, a); (dst, b) ] when Q.equal Q.one a && Q.equal Q.minus_one b ->
        Delta.add d src (Q.neg eq.p.cst) eq.bound dst
    | [ (dst, b); (src, a) ] when Q.equal Q.one a && Q.equal Q.minus_one b ->
        Delta.add d src (Q.neg eq.p.cst) eq.bound dst
    | _ -> ()

let apply d ~f truth g =
  let aux d bound a b =
    let bound =
      if truth then bound
      else match bound with Bound.Large -> Strict | Strict -> Large
    in
    let a, b = if truth then (a, b) else (b, a) in
    f d bound a b
  in
  match Ground.sem g with
  | { app = { builtin = Expr.Lt }; tyargs = []; args; _ } ->
      let a, b = IArray.extract2_exn args in
      aux d Bound.Strict a b
  | { app = { builtin = Expr.Leq }; tyargs = []; args; _ } ->
      let a, b = IArray.extract2_exn args in
      aux d Large a b
  | { app = { builtin = Expr.Geq }; tyargs = []; args; _ } ->
      let a, b = IArray.extract2_exn args in
      aux d Large b a
  | { app = { builtin = Expr.Gt }; tyargs = []; args; _ } ->
      let a, b = IArray.extract2_exn args in
      aux d Strict b a
  | _ -> assert false

let add_eq d eqs p bound origins =
  match (Polynome.is_cst p, bound) with
  | None, _ ->
      let eq = { p; bound; origins } in
      delta d eq;
      eq :: eqs
  | Some p, Bound.Large when Q.sign p = 0 ->
      Ground.S.iter
        (fun g ->
          let truth = Base.Option.value_exn (Boolean.is d (Ground.node g)) in
          apply d truth g ~f:(fun d bound a b ->
              match bound with
              | Strict
                when Dom_interval.is_integer d a && Dom_interval.is_integer d b
                ->
                  Dom_polynome.assume_equality d a
                    (Polynome.of_list Q.minus_one [ (b, Q.one) ])
              | Large ->
                  Dom_polynome.assume_equality d a (Polynome.monome Q.one b)
              | Strict -> assert false))
        origins;
      eqs
  | Some p, _ when Q.sign p < 0 ->
      Debug.dprintf2 debug "[Fourier] discard %a" Ground.S.pp origins;
      eqs
  | Some p, bound ->
      Debug.dprintf4 Debug.contradiction "[LRA/Fourier] Found %a%a0" Q.pp p
        Bound.pp bound;
      Egraph.contradiction d

let mk_eq d bound a b =
  let ( ! ) n =
    match Dom_polynome.get_repr d n with
    | None -> Polynome.monome Q.one (Egraph.find_def d n)
    | Some p -> p
  in
  let p, bound' =
    match bound with
    | Bound.Strict
      when Dom_interval.is_integer d a && Dom_interval.is_integer d b ->
        (Polynome.add_cst (Polynome.sub !a !b) Q.one, Bound.Large)
    | _ -> (Polynome.sub !a !b, bound)
  in
  ( p,
    bound',
    Polynome.sub (Polynome.monome Q.one a) (Polynome.monome Q.one b),
    bound )

let make_equations d (eqs, vars) g =
  Debug.dprintf2 debug "[Fourier] %a" Ground.pp g;
  match Boolean.is d (Ground.node g) with
  | None -> (eqs, vars)
  | Some truth -> (
      let p, bound, p_non_norm, bound_non_norm = apply d ~f:mk_eq truth g in
      Debug.dprintf6 debug "[Fourier] %a(%a)%a0" Polynome.pp p Polynome.pp
        p_non_norm Bound.pp bound;
      let eqs, vars =
        ( add_eq d eqs p bound (Ground.S.singleton g),
          Node.M.union_merge (fun _ _ _ -> Some ()) vars p.Polynome.poly )
      in
      match divide d p_non_norm with
      | Some (p', _, _) ->
          let p' = Dom_polynome.normalize d p' in
          ( add_eq d eqs p' bound_non_norm Ground.S.empty,
            Node.M.union_merge (fun _ _ _ -> Some ()) vars p'.Polynome.poly )
      | None -> (eqs, vars))

let fm (d : Egraph.wt) =
  Debug.dprintf0 debug "[Fourier]";
  Debug.incr stats_run;
  Datastructure.Ref.set scheduled d false;
  let eqs, vars =
    Datastructure.Push.fold comparisons d ~f:(make_equations d)
      ~init:([], Node.S.empty)
  in
  let rec split n positive negative absent = function
    | [] -> (positive, negative, absent)
    | eq :: l -> (
        match Node.M.find_opt n eq.p.poly with
        | None -> split n positive negative (eq :: absent) l
        | Some q ->
            if Q.sign q > 0 then split n ((q, eq) :: positive) negative absent l
            else split n positive ((q, eq) :: negative) absent l)
  in
  let rec aux eqs vars =
    match Node.S.choose vars with
    | exception Not_found -> ()
    | n -> (
        let vars = Node.S.remove n vars in
        let positive, negative, absent = split n [] [] [] eqs in
        match (positive, negative) with
        | [], _ | _, [] -> aux absent vars
        | _, _ ->
            let eqs =
              Lists.fold_product
                (fun eqs (pq, peq) (nq, neq) ->
                  let q = Q.div pq (Q.neg nq) in
                  let p = Polynome.cx_p_cy Q.one peq.p q neq.p in
                  let bound =
                    match (peq.bound, neq.bound) with
                    | Large, Large -> Bound.Large
                    | _ -> Bound.Strict
                  in
                  add_eq d eqs p bound (Ground.S.union peq.origins neq.origins))
                absent positive negative
            in
            aux eqs vars)
  in
  aux eqs vars

module Daemon = struct
  let key =
    Events.Dem.create
      (module struct
        type t = unit

        let name = "LRA.fourier"
      end)

  let enqueue d _ =
    if Datastructure.Ref.get scheduled d then Events.EnqAlready
    else (
      Datastructure.Ref.set scheduled d true;
      Events.EnqRun (key, (), None))

  let delay = Events.Delayed_by 64

  type runable = unit

  let print_runable = Unit.pp

  let run d () = fm d
end

let () = Events.register (module Daemon)

(** {2 Initialization} *)
let converter d (f : Ground.t) =
  match Ground.sem f with
  | {
   app = { builtin = Expr.Lt | Expr.Leq | Expr.Geq | Expr.Gt };
   tyargs = [];
   args;
  } ->
      let attach n =
        Dom_polynome.events_repr_change d ~node:n Daemon.enqueue;
        Events.attach_repr d n Daemon.enqueue
      in
      IArray.iter ~f:attach args;
      Events.attach_value d (Ground.node f) Boolean.dom (fun d n _ ->
          Daemon.enqueue d n);
      Datastructure.Push.push comparisons d f;
      Events.new_pending_daemon d Daemon.key ();
      Choice.register d (Boolean.chobool (Ground.node f))
  | _ -> ()

let init env = Ground.register_converter env converter