<commit_before>open Syntax
open BatPervasives

include Log.Make(struct let name = "srk.srkZ3" end)

type z3_context = Z3.context
type z3_expr = Z3.Expr.expr
type z3_func_decl = Z3.FuncDecl.func_decl

type 'a open_expr = [
  | `Real of QQ.t
  | `App of z3_func_decl * 'a list
  | `Var of int * typ_fo
  | `Add of 'a list
  | `Mul of 'a list
  | `Binop of [ `Div | `Mod | `Select ] * 'a * 'a
  | `Unop of [ `Floor | `Neg ] * 'a
  | `Store of 'a * 'a * 'a
  | `Tru
  | `Fls
  | `And of 'a list
  | `Or of 'a list
  | `Not of 'a
  | `Ite of 'a * 'a * 'a
  | `Quantify of [`Exists | `Forall] * string * typ_fo * 'a
  | `Atom of [`Eq | `Leq | `Lt] * 'a * 'a
]

let bool_val x =
  match Z3.Boolean.get_bool_value x with
  | Z3enums.L_TRUE -> true
  | Z3enums.L_FALSE -> false
  | Z3enums.L_UNDEF -> invalid_arg "bool_val: not a Boolean"

let rec qq_val ast =
  if Z3.Expr.is_numeral ast then
    QQ.of_string (Z3.Arithmetic.Real.numeral_to_string ast)
  else if Z3.FloatingPoint.is_to_real ast
       || Z3.Arithmetic.is_int2real ast
       || Z3.Arithmetic.is_real2int ast
  then
    qq_val (List.hd (Z3.Expr.get_args ast))
  else invalid_arg "qq_val"

let typ_of_sort sort =
  let open Z3enums in
  match Z3.Sort.get_sort_kind sort with
  | REAL_SORT -> `TyReal
  | INT_SORT -> `TyInt
  | BOOL_SORT -> `TyBool
  | ARRAY_SORT -> `TyArr
  | _ -> invalid_arg "typ_of_sort"

let rec sort_of_typ z3 = function
  | `TyInt -> Z3.Arithmetic.Integer.mk_sort z3
  | `TyReal -> Z3.Arithmetic.Real.mk_sort z3
  | `TyBool -> Z3.Boolean.mk_sort z3
  | `TyArr -> 
    Z3.Z3Array.mk_sort z3 (sort_of_typ z3 `TyInt) (sort_of_typ z3 `TyInt)

let rec eval alg ast =
  let open Z3 in
  let open Z3enums in
  match AST.get_ast_kind (Expr.ast_of_expr ast) with
  | APP_AST -> begin
      let decl = Expr.get_func_decl ast in
      let args = List.map (eval alg) (Expr.get_args ast) in
      match FuncDecl.get_decl_kind decl, args with
      | (OP_UNINTERPRETED, args) -> alg (`App (decl, args))
      | (OP_ADD, args) -> alg (`Add args)
      | (OP_MUL, args) -> alg (`Mul args)
      | (OP_SUB, [x;y]) -> alg (`Add [x; alg (`Unop (`Neg, y))])
      | (OP_UMINUS, [x]) -> alg (`Unop (`Neg, x))
      | (OP_MOD, [x;y]) -> alg (`Binop (`Mod, x, y))
      | (OP_IDIV, [x;y]) -> alg (`Unop (`Floor, alg (`Binop (`Div, x, y))))
      | (OP_DIV, [x;y]) -> alg (`Binop (`Div, x, y))
      | (OP_TO_REAL, [x]) -> x
      | (OP_TO_INT, [x]) -> alg (`Unop (`Floor, x))

      | (OP_STORE, [a; i; v]) -> alg (`Store (a, i, v))
      | (OP_SELECT, [a; i]) -> alg (`Binop(`Select, a, i))
      
      | (OP_TRUE, []) -> alg `Tru
      | (OP_FALSE, []) -> alg `Fls
      | (OP_AND, args) -> alg (`And args)
      | (OP_OR, args) -> alg (`Or args)
      | (OP_IMPLIES, [phi;psi]) -> alg (`Or [alg (`Not phi); psi])
      | (OP_IFF, [phi;psi]) ->
        alg (`Or [alg (`And [phi; psi]);
                  alg (`Not (alg (`Or [phi; psi])))])
      | (OP_NOT, [phi]) -> alg (`Not phi)
      | (OP_EQ, [s; t]) -> alg (`Atom (`Eq, s, t))
      | (OP_LE, [s; t]) -> alg (`Atom (`Leq, s, t))
      | (OP_GE, [s; t]) -> alg (`Atom (`Leq, t, s))
      | (OP_LT, [s; t]) -> alg (`Atom (`Lt, s, t))
      | (OP_GT, [s; t]) -> alg (`Atom (`Lt, t, s))
      | (OP_ITE, [cond; s; t]) -> alg (`Ite (cond, s, t))
      | (_, _) -> invalid_arg ("eval: unknown application: "
                               ^ (Expr.to_string ast))
    end
  | NUMERAL_AST ->
    alg (`Real (qq_val ast))
  | VAR_AST -><commit_msg>Translation for distinct and xor<commit_after>open Syntax
open BatPervasives

include Log.Make(struct let name = "srk.srkZ3" end)

type z3_context = Z3.context
type z3_expr = Z3.Expr.expr
type z3_func_decl = Z3.FuncDecl.func_decl

type 'a open_expr = [
  | `Real of QQ.t
  | `App of z3_func_decl * 'a list
  | `Var of int * typ_fo
  | `Add of 'a list
  | `Mul of 'a list
  | `Binop of [ `Div | `Mod | `Select ] * 'a * 'a
  | `Unop of [ `Floor | `Neg ] * 'a
  | `Store of 'a * 'a * 'a
  | `Tru
  | `Fls
  | `And of 'a list
  | `Or of 'a list
  | `Not of 'a
  | `Ite of 'a * 'a * 'a
  | `Quantify of [`Exists | `Forall] * string * typ_fo * 'a
  | `Atom of [`Eq | `Leq | `Lt] * 'a * 'a
]

let bool_val x =
  match Z3.Boolean.get_bool_value x with
  | Z3enums.L_TRUE -> true
  | Z3enums.L_FALSE -> false
  | Z3enums.L_UNDEF -> invalid_arg "bool_val: not a Boolean"

let rec qq_val ast =
  if Z3.Expr.is_numeral ast then
    QQ.of_string (Z3.Arithmetic.Real.numeral_to_string ast)
  else if Z3.FloatingPoint.is_to_real ast
       || Z3.Arithmetic.is_int2real ast
       || Z3.Arithmetic.is_real2int ast
  then
    qq_val (List.hd (Z3.Expr.get_args ast))
  else invalid_arg "qq_val"

let typ_of_sort sort =
  let open Z3enums in
  match Z3.Sort.get_sort_kind sort with
  | REAL_SORT -> `TyReal
  | INT_SORT -> `TyInt
  | BOOL_SORT -> `TyBool
  | ARRAY_SORT -> `TyArr
  | _ -> invalid_arg "typ_of_sort"

let rec sort_of_typ z3 = function
  | `TyInt -> Z3.Arithmetic.Integer.mk_sort z3
  | `TyReal -> Z3.Arithmetic.Real.mk_sort z3
  | `TyBool -> Z3.Boolean.mk_sort z3
  | `TyArr -> 
    Z3.Z3Array.mk_sort z3 (sort_of_typ z3 `TyInt) (sort_of_typ z3 `TyInt)

let rec eval alg ast =
  let open Z3 in
  let open Z3enums in
  match AST.get_ast_kind (Expr.ast_of_expr ast) with
  | APP_AST -> begin
      let decl = Expr.get_func_decl ast in
      let args = List.map (eval alg) (Expr.get_args ast) in
      match FuncDecl.get_decl_kind decl, args with
      | (OP_UNINTERPRETED, args) -> alg (`App (decl, args))
      | (OP_ADD, args) -> alg (`Add args)
      | (OP_MUL, args) -> alg (`Mul args)
      | (OP_SUB, [x;y]) -> alg (`Add [x; alg (`Unop (`Neg, y))])
      | (OP_UMINUS, [x]) -> alg (`Unop (`Neg, x))
      | (OP_MOD, [x;y]) -> alg (`Binop (`Mod, x, y))
      | (OP_IDIV, [x;y]) -> alg (`Unop (`Floor, alg (`Binop (`Div, x, y))))
      | (OP_DIV, [x;y]) -> alg (`Binop (`Div, x, y))
      | (OP_TO_REAL, [x]) -> x
      | (OP_TO_INT, [x]) -> alg (`Unop (`Floor, x))

      | (OP_STORE, [a; i; v]) -> alg (`Store (a, i, v))
      | (OP_SELECT, [a; i]) -> alg (`Binop(`Select, a, i))
      
      | (OP_TRUE, []) -> alg `Tru
      | (OP_FALSE, []) -> alg `Fls
      | (OP_AND, args) -> alg (`And args)
      | (OP_OR, args) -> alg (`Or args)
      | (OP_IMPLIES, [phi;psi]) -> alg (`Or [alg (`Not phi); psi])
      | (OP_IFF, [phi;psi]) ->
        alg (`Or [alg (`And [phi; psi]);
                  alg (`Not (alg (`Or [phi; psi])))])
      | (OP_NOT, [phi]) -> alg (`Not phi)
      | (OP_EQ, [s; t]) -> alg (`Atom (`Eq, s, t))
      | (OP_LE, [s; t]) -> alg (`Atom (`Leq, s, t))
      | (OP_GE, [s; t]) -> alg (`Atom (`Leq, t, s))
      | (OP_LT, [s; t]) -> alg (`Atom (`Lt, s, t))
      | (OP_GT, [s; t]) -> alg (`Atom (`Lt, t, s))
      | (OP_ITE, [cond; s; t]) -> alg (`Ite (cond, s, t))
      | (OP_DISTINCT, xs) ->
        let rec exclusive acc = function
          | y::ys ->
            exclusive
              (List.fold_left
                 (fun acc z -> alg (`Not (alg (`And [y; z])))::acc)
                 acc
                 ys)
              ys
          | [] -> acc
        in
        alg (`And (exclusive [alg (`Or xs)] xs))
      | (OP_XOR, xs) ->
        let xor x y =
          alg (`Or [alg (`And [alg (`Not x); y]);
                    alg (`And [x; alg (`Not y)])])
        in
        BatList.reduce xor xs

      | (_, _) -> invalid_arg ("eval: unknown application: "
                               ^ (Expr.to_string ast))
    end
  | NUMERAL_AST ->
    alg (`Real (qq_val ast))
  | VAR_AST ->
{'fraction_non_alphanumeric': 0.1821954484605087, 'fraction_numerical': 0.00963855421686747, 'mean_word_length': 2.2810715854194115, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 90, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4963723', 'n_tokens_mistral': 3266, 'n_tokens_neox': 3133, 'n_words': 863}