From 1dca0cb37543a3913054c4d5c4395c0fcdcdf3c0 Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Tue, 20 Oct 2020 02:38:36 -0700
Subject: [PATCH] [sledge] Evaluate function symbols applied to constants

Summary:
Function symbols when applied to literal values can be simplified by
evaluating them. This can be done even for function symbols that are
otherwise uninterpreted. This is not very strong, but is important in
some cases, and can prevent accumulating large complex terms that are
equal to literal constants.

Reviewed By: jvillard

Differential Revision: D24306044

fbshipit-source-id: 8c34d1ef2
---
 sledge/src/fol.ml           | 21 ++++++++-
 sledge/src/ses/funsym.ml    | 90 +++++++++++++++++++++++++++++++++++++
 sledge/src/ses/funsym.mli   | 12 +++++
 sledge/src/test/fol_test.ml | 14 +++---
 4 files changed, 129 insertions(+), 8 deletions(-)

diff --git a/sledge/src/fol.ml b/sledge/src/fol.ml
index ac3e2043a..913862dac 100644
--- a/sledge/src/fol.ml
+++ b/sledge/src/fol.ml
@@ -58,8 +58,21 @@ let equal_trm x y =
       Int.equal i j
   | _ -> equal_trm x y
 
+(* destructors *)
+
+let get_z = function Z z -> Some z | _ -> None
+let get_q = function Q q -> Some q | Z z -> Some (Q.of_z z) | _ -> None
+
+(* constructors *)
+
+(* statically allocated since they are tested with == *)
 let zero = Z Z.zero
 let one = Z Z.one
+
+let _Z z =
+  if Z.equal Z.zero z then zero else if Z.equal Z.one z then one else Z z
+
+let _Q q = if Z.equal Z.one (Q.den q) then _Z (Q.num q) else Q q
 let _Neg x = Neg x
 
 let _Add x y =
@@ -81,7 +94,13 @@ let _Select idx rcd = Select {idx; rcd}
 let _Update idx rcd elt = Update {idx; rcd; elt}
 let _Record es = Record es
 let _Ancestor i = Ancestor i
-let _Apply f es = Apply (f, es)
+
+let _Apply f es =
+  match
+    Funsym.eval ~equal:equal_trm ~get_z ~ret_z:_Z ~get_q ~ret_q:_Q f es
+  with
+  | Some c -> c
+  | None -> Apply (f, es)
 
 (*
  * Formulas
diff --git a/sledge/src/ses/funsym.ml b/sledge/src/ses/funsym.ml
index f00a5086e..295f82a3c 100644
--- a/sledge/src/ses/funsym.ml
+++ b/sledge/src/ses/funsym.ml
@@ -39,3 +39,93 @@ let pp ppf f =
   | Uninterp sym -> pf "%s" sym
 
 let uninterp s = Uninterp s
+
+let eval ~equal ~get_z ~ret_z ~get_q ~ret_q:_ f xs =
+  match (f, xs) with
+  | Rem, [|x; y|] -> (
+    match get_z y with
+    (* x % 1 ==> 0 *)
+    | Some j when Z.equal Z.one j -> Some (ret_z Z.zero)
+    | Some j when not (Z.equal Z.zero j) -> (
+      match get_z x with
+      (* i % j *)
+      | Some i -> Some (ret_z (Z.rem i j))
+      | None -> (
+        match get_q x with
+        (* (n/d) % i ==> (n / d) % i *)
+        | Some {Q.num; den} -> Some (ret_z (Z.rem (Z.div num den) j))
+        | None -> None ) )
+    | _ -> None )
+  | BitAnd, [|x; y|] -> (
+    match (get_z x, get_z y) with
+    (* i && j *)
+    | Some i, Some j -> Some (ret_z (Z.logand i j))
+    (* x && true ==> x *)
+    | _, Some z when Z.is_true z -> Some x
+    | Some z, _ when Z.is_true z -> Some y
+    (* x && false ==> false *)
+    | _, Some z when Z.is_false z -> Some y
+    | Some z, _ when Z.is_false z -> Some x
+    (* x && x ==> x *)
+    | _ when equal x y -> Some x
+    | _ -> None )
+  | BitOr, [|x; y|] -> (
+    match (get_z x, get_z y) with
+    (* i || j *)
+    | Some i, Some j -> Some (ret_z (Z.logor i j))
+    (* x || true ==> true *)
+    | _, Some z when Z.is_true z -> Some y
+    | Some z, _ when Z.is_true z -> Some x
+    (* x || false ==> x *)
+    | _, Some z when Z.is_false z -> Some x
+    | Some z, _ when Z.is_false z -> Some y
+    (* x || x ==> x *)
+    | _ when equal x y -> Some x
+    | _ -> None )
+  | BitShl, [|x; y|] -> (
+    match get_z y with
+    (* x shl 0 ==> x *)
+    | Some z when Z.equal Z.zero z -> Some x
+    | get_z_y -> (
+      match (get_z x, get_z_y) with
+      (* i shl j *)
+      | Some i, Some j when Z.sign j >= 0 -> (
+        match Z.to_int j with
+        | n -> Some (ret_z (Z.shift_left i n))
+        | exception Z.Overflow -> None )
+      | _ -> None ) )
+  | BitLshr, [|x; y|] -> (
+    match get_z y with
+    (* x lshr 0 ==> x *)
+    | Some z when Z.equal Z.zero z -> Some x
+    | get_z_y -> (
+      match (get_z x, get_z_y) with
+      (* i lshr j *)
+      | Some i, Some j when Z.sign j >= 0 -> (
+        match Z.to_int j with
+        | n -> Some (ret_z (Z.shift_right_trunc i n))
+        | exception Z.Overflow -> None )
+      | _ -> None ) )
+  | BitAshr, [|x; y|] -> (
+    match get_z y with
+    (* x ashr 0 ==> x *)
+    | Some z when Z.equal Z.zero z -> Some x
+    | get_z_y -> (
+      match (get_z x, get_z_y) with
+      (* i ashr j *)
+      | Some i, Some j when Z.sign j >= 0 -> (
+        match Z.to_int j with
+        | n -> Some (ret_z (Z.shift_right i n))
+        | exception Z.Overflow -> None )
+      | _ -> None ) )
+  | Signed n, [|x|] -> (
+    match get_z x with
+    (* (sN)i *)
+    | Some i -> Some (ret_z (Z.signed_extract i 0 n))
+    | _ -> None )
+  | Unsigned n, [|x|] -> (
+    match get_z x with
+    (* (uN)i *)
+    | Some i -> Some (ret_z (Z.extract i 0 n))
+    | _ -> None )
+  | _ -> None
diff --git a/sledge/src/ses/funsym.mli b/sledge/src/ses/funsym.mli
index 5913296dc..c65e06187 100644
--- a/sledge/src/ses/funsym.mli
+++ b/sledge/src/ses/funsym.mli
@@ -36,3 +36,15 @@ type t =
 
 val pp : t pp
 val uninterp : string -> t
+
+val eval :
+     equal:('a -> 'a -> bool)
+  -> get_z:('a -> Z.t option)
+  -> ret_z:(Z.t -> 'a)
+  -> get_q:('a -> Q.t option)
+  -> ret_q:(Q.t -> 'a)
+  -> t
+  -> 'a array
+  -> 'a option
+(** [eval ~equal ~get_z ~ret_z ~get_q ~ret_q f xs] evaluates the application
+    [f xs] using the provided embeddings of arguments into values. *)
diff --git a/sledge/src/test/fol_test.ml b/sledge/src/test/fol_test.ml
index 6818497f4..c359a3f93 100644
--- a/sledge/src/test/fol_test.ml
+++ b/sledge/src/test/fol_test.ml
@@ -537,31 +537,31 @@ let%test_module _ =
 
     let%expect_test _ =
       pp_exp Llair.(Exp.signed 1 !!1 ~to_:Typ.bool) ;
-      [%expect {| exp= ((i1)(s1) 1); term= (s1)(1) |}]
+      [%expect {| exp= ((i1)(s1) 1); term= -1 |}]
 
     let%expect_test _ =
       pp_exp Llair.(Exp.unsigned 1 !!(-1) ~to_:Typ.byt) ;
-      [%expect {| exp= ((i8)(u1) -1); term= (u1)(-1) |}]
+      [%expect {| exp= ((i8)(u1) -1); term= 1 |}]
 
     let%expect_test _ =
       pp_exp Llair.(Exp.signed 8 !!(-1) ~to_:Typ.int) ;
-      [%expect {| exp= ((i32)(s8) -1); term= (s8)(-1) |}]
+      [%expect {| exp= ((i32)(s8) -1); term= -1 |}]
 
     let%expect_test _ =
       pp_exp Llair.(Exp.unsigned 8 !!(-1) ~to_:Typ.int) ;
-      [%expect {| exp= ((i32)(u8) -1); term= (u8)(-1) |}]
+      [%expect {| exp= ((i32)(u8) -1); term= 255 |}]
 
     let%expect_test _ =
       pp_exp Llair.(Exp.signed 8 !!255 ~to_:Typ.byt) ;
-      [%expect {| exp= ((i8)(s8) 255); term= (s8)(255) |}]
+      [%expect {| exp= ((i8)(s8) 255); term= -1 |}]
 
     let%expect_test _ =
       pp_exp Llair.(Exp.signed 7 !!255 ~to_:Typ.byt) ;
-      [%expect {| exp= ((i8)(s7) 255); term= (s7)(255) |}]
+      [%expect {| exp= ((i8)(s7) 255); term= -1 |}]
 
     let%expect_test _ =
       pp_exp Llair.(Exp.unsigned 7 !!255 ~to_:Typ.byt) ;
-      [%expect {| exp= ((i8)(u7) 255); term= (u7)(255) |}]
+      [%expect {| exp= ((i8)(u7) 255); term= 127 |}]
 
     let%expect_test _ =
       pp_exp