From 367c1945efafb570096a8558215f300f18c0716e Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Wed, 24 Mar 2021 15:30:01 -0400
Subject: [PATCH] feat(main): handle `check-sat-assuming` statement

---
 src/arith/base-term/Base_types.ml | 11 ++++++++---
 src/main/main.ml                  |  2 +-
 src/smtlib/Process.ml             |  7 +++++--
 src/smtlib/Typecheck.ml           | 15 ++++++++++++---
 4 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/src/arith/base-term/Base_types.ml b/src/arith/base-term/Base_types.ml
index 6895a84c..5c85ce9d 100644
--- a/src/arith/base-term/Base_types.ml
+++ b/src/arith/base-term/Base_types.ml
@@ -146,7 +146,7 @@ type statement =
   | Stmt_define of definition list
   | Stmt_assert of term
   | Stmt_assert_clause of term list
-  | Stmt_check_sat
+  | Stmt_check_sat of (bool * term) list
   | Stmt_exit
 
 let[@inline] term_equal_ (a:term) b = a==b
@@ -1110,14 +1110,19 @@ module Statement = struct
     | Stmt_define of definition list
     | Stmt_assert of term
     | Stmt_assert_clause of term list
-    | Stmt_check_sat
+    | Stmt_check_sat of (bool * term) list
     | Stmt_exit
 
   let pp out = function
     | Stmt_set_logic s -> Fmt.fprintf out "(set-logic %s)" s
     | Stmt_set_option l -> Fmt.fprintf out "(@[set-logic@ %a@])" (Util.pp_list Fmt.string) l
     | Stmt_set_info (a,b) -> Fmt.fprintf out "(@[set-info@ %s@ %s@])" a b
-    | Stmt_check_sat -> Fmt.string out "(check-sat)"
+    | Stmt_check_sat [] -> Fmt.string out "(check-sat)"
+    | Stmt_check_sat l ->
+      let pp_pair out (b,t) =
+        if b then pp_term out t else Fmt.fprintf out "(@[not %a@])" pp_term t in
+      Fmt.fprintf out "(@[check-sat-assuming@ (@[%a@])@])"
+        (Fmt.list pp_pair) l
     | Stmt_ty_decl (s,n) -> Fmt.fprintf out "(@[declare-sort@ %a %d@])" ID.pp s n
     | Stmt_decl (id,args,ret) ->
       Fmt.fprintf out "(@[<1>declare-fun@ %a (@[%a@])@ %a@])"
diff --git a/src/main/main.ml b/src/main/main.ml
index 1eaf4a6c..fa77de89 100644
--- a/src/main/main.ml
+++ b/src/main/main.ml
@@ -111,7 +111,7 @@ module Dimacs = struct
              (fun c ->
                 let lits = List.rev_map get_lit c in
                 stmts := Statement.Stmt_assert_clause lits :: !stmts);
-           stmts := Statement.Stmt_check_sat :: !stmts;
+           stmts := Statement.Stmt_check_sat [] :: !stmts;
            Ok (List.rev !stmts))
     with e ->
       E.of_exn_trace e
diff --git a/src/smtlib/Process.ml b/src/smtlib/Process.ml
index dbd6f529..bf23d59e 100644
--- a/src/smtlib/Process.ml
+++ b/src/smtlib/Process.ml
@@ -235,10 +235,13 @@ let process_stmt
     | Statement.Stmt_exit ->
       Log.debug 1 "exit";
       raise Exit
-    | Statement.Stmt_check_sat ->
+    | Statement.Stmt_check_sat l ->
+      let assumptions =
+        List.map (fun (sign,t) -> Solver.mk_atom_t solver ~sign t) l
+      in
       solve
         ?gc ?restarts ?dot_proof ~check ?pp_model ?time ?memory ?progress
-        ~assumptions:[] ?hyps
+        ~assumptions ?hyps
         solver;
       E.return()
     | Statement.Stmt_ty_decl (id,n) ->
diff --git a/src/smtlib/Typecheck.ml b/src/smtlib/Typecheck.ml
index 4a9b6c74..1199255a 100644
--- a/src/smtlib/Typecheck.ml
+++ b/src/smtlib/Typecheck.ml
@@ -513,9 +513,18 @@ and conv_statement_aux ctx (stmt:PA.statement) : Stmt.t list =
     let t = conv_term ctx t in
     check_bool_ ctx t;
     [Stmt.Stmt_assert t]
-  | PA.Stmt_check_sat -> [Stmt.Stmt_check_sat]
-  | PA.Stmt_check_sat_assuming _
-  | PA.Stmt_get_assertions 
+  | PA.Stmt_check_sat -> [Stmt.Stmt_check_sat []]
+  | PA.Stmt_check_sat_assuming l ->
+    let l =
+      List.map
+        (fun (t,b) ->
+           let t = conv_term ctx (PA.const t) in
+           check_bool_ ctx t;
+           b,t)
+        l
+    in
+    [Stmt.Stmt_check_sat l]
+  | PA.Stmt_get_assertions
   | PA.Stmt_get_option _
   | PA.Stmt_get_info _
   | PA.Stmt_get_model