(*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *)

let pp_boxed fs fmt =
  Format.pp_open_box fs 2 ;
  Format.kfprintf (fun fs -> Format.pp_close_box fs ()) fs fmt

module Funsym = Ses.Funsym
module Predsym = Ses.Predsym

(*
 * Terms
 *)

(** Terms, denoting functions from structures to values, built from
    variables and applications of function symbols from various theories. *)
type trm =
  (* variables *)
  | Var of {id: int; name: string}
  (* arithmetic *)
  | Z of Z.t
  | Q of Q.t
  | Neg of trm
  | Add of trm * trm
  | Sub of trm * trm
  | Mulq of Q.t * trm
  (* sequences (of non-fixed size) *)
  | Splat of trm
  | Sized of {seq: trm; siz: trm}
  | Extract of {seq: trm; off: trm; len: trm}
  | Concat of trm array
  (* records (with fixed indices) *)
  | Select of {idx: int; rcd: trm}
  | Update of {idx: int; rcd: trm; elt: trm}
  | Record of trm array
  | Ancestor of int
  (* uninterpreted *)
  | Apply of Funsym.t * trm array
[@@deriving compare, equal, sexp]

let compare_trm x y =
  if x == y then 0
  else
    match (x, y) with
    | Var {id= i; name= _}, Var {id= j; name= _} when i > 0 && j > 0 ->
        Int.compare i j
    | _ -> compare_trm x y

let equal_trm x y =
  x == y
  ||
  match (x, y) with
  | Var {id= i; name= _}, Var {id= j; name= _} when i > 0 && j > 0 ->
      Int.equal i j
  | _ -> equal_trm x y

let zero = Z Z.zero
let one = Z Z.one
let _Neg x = Neg x

let _Add x y =
  match (x, y) with
  | _, Q q when Q.sign q = 0 -> x
  | Q q, _ when Q.sign q = 0 -> y
  | _ -> Add (x, y)

let _Sub x y = Sub (x, y)

let _Mulq q x =
  if Q.equal Q.one q then x else if Q.sign q = 0 then zero else Mulq (q, x)

let _Splat x = Splat x
let _Sized seq siz = Sized {seq; siz}
let _Extract seq off len = Extract {seq; off; len}
let _Concat es = Concat es
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)

(*
 * Formulas
 *)

(** Formulas, denoting sets of structures, built from propositional
    variables, applications of predicate symbols from various theories, and
    first-order logic connectives. *)
module Fml : sig
  type fml = private
    (* propositional constants *)
    | Tt
    | Ff
    (* equality *)
    | Eq of trm * trm
    | Dq of trm * trm
    (* arithmetic *)
    | Eq0 of trm  (** [Eq0(x)] iff x = 0 *)
    | Dq0 of trm  (** [Dq0(x)] iff x ≠ 0 *)
    | Gt0 of trm  (** [Gt0(x)] iff x > 0 *)
    | Ge0 of trm  (** [Ge0(x)] iff x ≥ 0 *)
    | Lt0 of trm  (** [Lt0(x)] iff x < 0 *)
    | Le0 of trm  (** [Le0(x)] iff x ≤ 0 *)
    (* propositional connectives *)
    | And of fml * fml
    | Or of fml * fml
    | Iff of fml * fml
    | Xor of fml * fml
    | Cond of {cnd: fml; pos: fml; neg: fml}
    (* uninterpreted literals *)
    | UPosLit of Predsym.t * trm array
    | UNegLit of Predsym.t * trm array
  [@@deriving compare, equal, sexp]

  val _Tt : fml
  val _Ff : fml
  val _Eq : trm -> trm -> fml
  val _Dq : trm -> trm -> fml
  val _Eq0 : trm -> fml
  val _Dq0 : trm -> fml
  val _Gt0 : trm -> fml
  val _Ge0 : trm -> fml
  val _Lt0 : trm -> fml
  val _Le0 : trm -> fml
  val _Not : fml -> fml
  val _And : fml -> fml -> fml
  val _Or : fml -> fml -> fml
  val _Iff : fml -> fml -> fml
  val _Xor : fml -> fml -> fml
  val _Cond : fml -> fml -> fml -> fml
  val _UPosLit : Predsym.t -> trm array -> fml
  val _UNegLit : Predsym.t -> trm array -> fml
end = struct
  type fml =
    | Tt
    | Ff
    | Eq of trm * trm
    | Dq of trm * trm
    | Eq0 of trm
    | Dq0 of trm
    | Gt0 of trm
    | Ge0 of trm
    | Lt0 of trm
    | Le0 of trm
    | And of fml * fml
    | Or of fml * fml
    | Iff of fml * fml
    | Xor of fml * fml
    | Cond of {cnd: fml; pos: fml; neg: fml}
    | UPosLit of Predsym.t * trm array
    | UNegLit of Predsym.t * trm array
  [@@deriving compare, equal, sexp]

  let sort_fml x y = if compare_fml x y <= 0 then (x, y) else (y, x)

  (** Some normalization is necessary for [embed_into_fml] (defined below)
      to be left inverse to [embed_into_cnd]. Essentially
      [0 ≠ (p ? 1 : 0)] needs to normalize to [p], by way of
      [0 ≠ (p ? 1 : 0)] ==> [(p ? 0 ≠ 1 : 0 ≠ 0)] ==> [(p ? tt : ff)]
      ==> [p]. *)

  let _Tt = Tt
  let _Ff = Ff

  (** classification of terms as either semantically equal or disequal, or
      if semantic relationship is unknown, as either syntactically less than
      or greater than *)
  type compare_semantic_syntactic = SemEq | SemDq | SynLt | SynGt

  let compare_semantic_syntactic d e =
    match (d, e) with
    | Z y, Z z -> if Z.equal y z then SemEq else SemDq
    | Q q, Q r -> if Q.equal q r then SemEq else SemDq
    | Z z, Q q | Q q, Z z -> if Q.equal (Q.of_z z) q then SemEq else SemDq
    | _ ->
        let ord = compare_trm d e in
        if ord < 0 then SynLt else if ord = 0 then SemEq else SynGt

  let _Eq0 x =
    match compare_semantic_syntactic zero x with
    (* 0 = 0 ==> tt *)
    | SemEq -> Tt
    (* 0 = N ==> ff for N ≢ 0 *)
    | SemDq -> Ff
    | SynLt | SynGt -> Eq0 x

  let _Dq0 x =
    match compare_semantic_syntactic zero x with
    (* 0 ≠ 0 ==> ff *)
    | SemEq -> Ff
    (* 0 ≠ N ==> tt for N ≢ 0 *)
    | SemDq -> Tt
    | SynLt | SynGt -> Dq0 x

  let _Eq x y =
    if x == zero then _Eq0 y
    else if y == zero then _Eq0 x
    else
      match compare_semantic_syntactic x y with
      | SemEq -> Tt
      | SemDq -> Ff
      | SynLt -> Eq (x, y)
      | SynGt -> Eq (y, x)

  let _Dq x y =
    if x == zero then _Dq0 y
    else if y == zero then _Dq0 x
    else
      match compare_semantic_syntactic x y with
      | SemEq -> Ff
      | SemDq -> Tt
      | SynLt -> Dq (x, y)
      | SynGt -> Dq (y, x)

  let _Gt0 = function
    | Z z -> if Z.gt z Z.zero then Tt else Ff
    | Q q -> if Q.gt q Q.zero then Tt else Ff
    | x -> Gt0 x

  let _Ge0 = function
    | Z z -> if Z.geq z Z.zero then Tt else Ff
    | Q q -> if Q.geq q Q.zero then Tt else Ff
    | x -> Ge0 x

  let _Lt0 = function
    | Z z -> if Z.lt z Z.zero then Tt else Ff
    | Q q -> if Q.lt q Q.zero then Tt else Ff
    | x -> Lt0 x

  let _Le0 = function
    | Z z -> if Z.leq z Z.zero then Tt else Ff
    | Q q -> if Q.leq q Q.zero then Tt else Ff
    | x -> Le0 x

  let _UPosLit p xs = UPosLit (p, xs)
  let _UNegLit p xs = UNegLit (p, xs)

  let is_negative = function
    | Ff | Dq _ | Dq0 _ | Lt0 _ | Le0 _ | Or _ | Xor _ | UNegLit _ -> true
    | Tt | Eq _ | Eq0 _ | Gt0 _ | Ge0 _ | And _ | Iff _ | UPosLit _ | Cond _
      ->
        false

  type equal_or_opposite = Equal | Opposite | Unknown

  let rec equal_or_opposite p q =
    match (p, q) with
    | Tt, Ff | Ff, Tt -> Opposite
    | Eq (a, b), Dq (a', b') | Dq (a, b), Eq (a', b') ->
        if equal_trm a a' && equal_trm b b' then Opposite else Unknown
    | Eq0 a, Dq0 a'
     |Dq0 a, Eq0 a'
     |Gt0 a, Le0 a'
     |Ge0 a, Lt0 a'
     |Lt0 a, Ge0 a'
     |Le0 a, Gt0 a' ->
        if equal_trm a a' then Opposite else Unknown
    | And (a, b), Or (a', b') | Or (a', b'), And (a, b) -> (
      match equal_or_opposite a a' with
      | Opposite -> (
        match equal_or_opposite b b' with
        | Opposite -> Opposite
        | _ -> Unknown )
      | _ -> Unknown )
    | Iff (p, q), Xor (p', q') | Xor (p, q), Iff (p', q') ->
        if equal_fml p p' && equal_fml q q' then Opposite else Unknown
    | Cond {cnd= c; pos= p; neg= n}, Cond {cnd= c'; pos= p'; neg= n'} ->
        if equal_fml c c' then
          match equal_or_opposite p p' with
          | Opposite -> (
            match equal_or_opposite n n' with
            | Opposite -> Opposite
            | _ -> Unknown )
          | Equal -> if equal_fml n n' then Equal else Unknown
          | Unknown -> Unknown
        else Unknown
    | UPosLit (p, xs), UNegLit (p', xs') | UNegLit (p, xs), UPosLit (p', xs')
      ->
        if Predsym.equal p p' && Array.equal equal_trm xs xs' then Opposite
        else Unknown
    | _ -> if equal_fml p q then Equal else Unknown

  let _And p q =
    match (p, q) with
    | Tt, p | p, Tt -> p
    | Ff, _ | _, Ff -> Ff
    | _ -> (
      match equal_or_opposite p q with
      | Equal -> p
      | Opposite -> Ff
      | Unknown ->
          let p, q = sort_fml p q in
          And (p, q) )

  let _Or p q =
    match (p, q) with
    | Ff, p | p, Ff -> p
    | Tt, _ | _, Tt -> Tt
    | _ -> (
      match equal_or_opposite p q with
      | Equal -> p
      | Opposite -> Tt
      | Unknown ->
          let p, q = sort_fml p q in
          Or (p, q) )

  let rec _Iff p q =
    match (p, q) with
    | Tt, p | p, Tt -> p
    | Ff, p | p, Ff -> _Not p
    | _ -> (
      match equal_or_opposite p q with
      | Equal -> Tt
      | Opposite -> Ff
      | Unknown ->
          let p, q = sort_fml p q in
          Iff (p, q) )

  and _Xor p q =
    match (p, q) with
    | Tt, p | p, Tt -> _Not p
    | Ff, p | p, Ff -> p
    | _ -> (
      match equal_or_opposite p q with
      | Equal -> Ff
      | Opposite -> Tt
      | Unknown ->
          let p, q = sort_fml p q in
          Xor (p, q) )

  and _Not = function
    | Tt -> _Ff
    | Ff -> _Tt
    | Eq (x, y) -> _Dq x y
    | Dq (x, y) -> _Eq x y
    | Eq0 x -> _Dq0 x
    | Dq0 x -> _Eq0 x
    | Gt0 x -> _Le0 x
    | Ge0 x -> _Lt0 x
    | Lt0 x -> _Ge0 x
    | Le0 x -> _Gt0 x
    | And (x, y) -> _Or (_Not x) (_Not y)
    | Or (x, y) -> _And (_Not x) (_Not y)
    | Iff (x, y) -> _Xor x y
    | Xor (x, y) -> _Iff x y
    | Cond {cnd; pos; neg} -> _Cond cnd (_Not pos) (_Not neg)
    | UPosLit (p, xs) -> _UNegLit p xs
    | UNegLit (p, xs) -> _UPosLit p xs

  and _Cond cnd pos neg =
    match (cnd, pos, neg) with
    (* (tt ? p : n) ==> p *)
    | Tt, _, _ -> pos
    (* (ff ? p : n) ==> n *)
    | Ff, _, _ -> neg
    (* (c ? tt : ff) ==> c *)
    | _, Tt, Ff -> cnd
    (* (c ? ff : tt) ==> ¬c *)
    | _, Ff, Tt -> _Not cnd
    (* (c ? p : ff) ==> c ∧ p *)
    | _, _, Ff -> _And cnd pos
    (* (c ? ff : n) ==> ¬c ∧ n *)
    | _, Ff, _ -> _And (_Not cnd) neg
    (* (c ? tt : n) ==> c ∨ n *)
    | _, Tt, _ -> _Or cnd neg
    (* (c ? p : tt) ==> ¬c ∨ p *)
    | _, _, Tt -> _Or (_Not cnd) pos
    | _ -> (
      match equal_or_opposite pos neg with
      (* (c ? p : p) ==> c *)
      | Equal -> cnd
      (* (c ? p : ¬p) ==> c <=> p *)
      | Opposite -> _Iff cnd pos
      (* (c ? p : n) <=> (¬c ? n : p) *)
      | Unknown when is_negative cnd ->
          Cond {cnd= _Not cnd; pos= neg; neg= pos}
      (* (c ? p : n) *)
      | _ -> Cond {cnd; pos; neg} )
end

open Fml

(*
 * Conditional terms
 *)

(** Conditional terms, denoting functions from structures to values, taking
    the form of trees with internal nodes labeled with formulas and leaves
    labeled with terms. *)
type cnd = [`Ite of fml * cnd * cnd | `Trm of trm]
[@@deriving compare, equal, sexp]

(*
 * Expressions
 *)

(** Expressions, which are partitioned into terms, conditional terms, and
    formulas. *)
type exp = [cnd | `Fml of fml] [@@deriving compare, equal, sexp]

(*
 * Variables
 *)

(** Variable terms *)
module Var : sig
  include Ses.Var_intf.VAR with type t = private trm

  val of_ : trm -> t
end = struct
  module T = struct
    type t = trm [@@deriving compare, equal, sexp]
  end

  let invariant (x : trm) =
    let@ () = Invariant.invariant [%here] x [%sexp_of: trm] in
    match x with
    | Var _ -> ()
    | _ -> fail "non-var: %a" Sexp.pp_hum (sexp_of_trm x) ()

  include Ses.Var0.Make (struct
    include T

    let make ~id ~name = Var {id; name} |> check invariant
    let id = function Var v -> v.id | x -> violates invariant x
    let name = function Var v -> v.name | x -> violates invariant x
  end)

  let of_ v = v |> check invariant
end

type var = Var.t

(*
 * Representation operations
 *)

(** pp *)

let rec ppx_t strength fs trm =
  let rec pp fs trm =
    let pf fmt = pp_boxed fs fmt in
    match trm with
    | Var _ as v -> Var.ppx strength fs (Var.of_ v)
    | Z z -> Trace.pp_styled `Magenta "%a" fs Z.pp z
    | Q q -> Trace.pp_styled `Magenta "%a" fs Q.pp q
    | Neg x -> pf "(- %a)" pp x
    | Add (x, y) -> pf "(%a@ + %a)" pp x pp y
    | Sub (x, y) -> pf "(%a@ - %a)" pp x pp y
    | Mulq (q, x) -> pf "(%a@ @<2>× %a)" Q.pp q pp x
    | Splat x -> pf "%a^" pp x
    | Sized {seq; siz} -> pf "@<1>⟨%a,%a@<1>⟩" pp siz pp seq
    | Extract {seq; off; len} -> pf "%a[%a,%a)" pp seq pp off pp len
    | Concat [||] -> pf "@<2>⟨⟩"
    | Concat xs -> pf "(%a)" (Array.pp "@,^" pp) xs
    | Select {idx; rcd} -> pf "%a[%i]" pp rcd idx
    | Update {idx; rcd; elt} -> pf "[%a@ @[| %i → %a@]]" pp rcd idx pp elt
    | Record xs -> pf "{%a}" (ppx_record strength) xs
    | Ancestor i -> pf "(ancestor %i)" i
    | Apply (f, [||]) -> pf "%a" Funsym.pp f
    | Apply
        ( ( ( Mul | Div | Rem | BitAnd | BitOr | BitXor | BitShl | BitLshr
            | BitAshr ) as f )
        , [|x; y|] ) ->
        pf "(%a@ %a@ %a)" pp x Funsym.pp f pp y
    | Apply (f, es) ->
        pf "%a(%a)" Funsym.pp f (Array.pp ",@ " (ppx_t strength)) es
  in
  pp fs trm

and ppx_record strength fs elts =
  [%Trace.fprintf
    fs "%a"
      (fun fs elts ->
        let exception Not_a_string in
        match
          String.init (Array.length elts) ~f:(fun i ->
              match elts.(i) with
              | Z c -> Char.of_int_exn (Z.to_int c)
              | _ -> raise Not_a_string )
        with
        | s -> Format.fprintf fs "%S" s
        | exception (Not_a_string | Z.Overflow | Failure _) ->
            Format.fprintf fs "@[<h>%a@]"
              (Array.pp ",@ " (ppx_t strength))
              elts )
      elts]

let ppx_f strength fs fml =
  let pp_t = ppx_t strength in
  let rec pp fs fml =
    let pf fmt = pp_boxed fs fmt in
    match (fml : fml) with
    | Tt -> pf "tt"
    | Ff -> pf "ff"
    | Eq (x, y) -> pf "(%a@ = %a)" pp_t x pp_t y
    | Dq (x, y) -> pf "(%a@ @<2>≠ %a)" pp_t x pp_t y
    | Eq0 x -> pf "(0 = %a)" pp_t x
    | Dq0 x -> pf "(0 @<2>≠ %a)" pp_t x
    | Gt0 x -> pf "(0 < %a)" pp_t x
    | Ge0 x -> pf "(0 @<2>≤ %a)" pp_t x
    | Lt0 x -> pf "(0 > %a)" pp_t x
    | Le0 x -> pf "(0 @<2>≥ %a)" pp_t x
    | And (x, y) -> pf "(%a@ @<2>∧ %a)" pp x pp y
    | Or (x, y) -> pf "(%a@ @<2>∨ %a)" pp x pp y
    | Iff (x, y) -> pf "(%a@ <=> %a)" pp x pp y
    | Xor (x, y) -> pf "(%a@ xor %a)" pp x pp y
    | Cond {cnd; pos; neg} ->
        pf "@[<hv 1>(%a@ ? %a@ : %a)@]" pp cnd pp pos pp neg
    | UPosLit (p, xs) -> pf "%a(%a)" Predsym.pp p (Array.pp ",@ " pp_t) xs
    | UNegLit (p, xs) ->
        pf "@<1>¬%a(%a)" Predsym.pp p (Array.pp ",@ " pp_t) xs
  in
  pp fs fml

let pp_f = ppx_f (fun _ -> None)

let ppx_c strength fs ct =
  let pp_t = ppx_t strength in
  let pp_f = ppx_f strength in
  let rec pp fs ct =
    let pf fmt = pp_boxed fs fmt in
    match ct with
    | `Ite (cnd, thn, els) -> pf "(%a@ ? %a@ : %a)" pp_f cnd pp thn pp els
    | `Trm t -> pp_t fs t
  in
  pp fs ct

let ppx strength fs = function
  | #cnd as c -> ppx_c strength fs c
  | `Fml f -> ppx_f strength fs f

let pp = ppx (fun _ -> None)

(** fold_vars *)

let rec fold_vars_t e ~init ~f =
  match e with
  | Var _ as v -> f init (Var.of_ v)
  | Z _ | Q _ | Ancestor _ -> init
  | Neg x | Mulq (_, x) | Splat x | Select {rcd= x} ->
      fold_vars_t ~f x ~init
  | Add (x, y)
   |Sub (x, y)
   |Sized {seq= x; siz= y}
   |Update {rcd= x; elt= y} ->
      fold_vars_t ~f x ~init:(fold_vars_t ~f y ~init)
  | Extract {seq= x; off= y; len= z} ->
      fold_vars_t ~f x
        ~init:(fold_vars_t ~f y ~init:(fold_vars_t ~f z ~init))
  | Concat xs | Record xs | Apply (_, xs) ->
      Array.fold ~f:(fun init -> fold_vars_t ~f ~init) xs ~init

let rec fold_vars_f ~init p ~f =
  match (p : fml) with
  | Tt | Ff -> init
  | Eq (x, y) | Dq (x, y) -> fold_vars_t ~f x ~init:(fold_vars_t ~f y ~init)
  | Eq0 x | Dq0 x | Gt0 x | Ge0 x | Lt0 x | Le0 x -> fold_vars_t ~f x ~init
  | And (x, y) | Or (x, y) | Iff (x, y) | Xor (x, y) ->
      fold_vars_f ~f x ~init:(fold_vars_f ~f y ~init)
  | Cond {cnd; pos; neg} ->
      fold_vars_f ~f cnd
        ~init:(fold_vars_f ~f pos ~init:(fold_vars_f ~f neg ~init))
  | UPosLit (_, xs) | UNegLit (_, xs) ->
      Array.fold ~f:(fun init -> fold_vars_t ~f ~init) xs ~init

let rec fold_vars_c ~init ~f = function
  | `Ite (cnd, thn, els) ->
      fold_vars_f ~f cnd
        ~init:(fold_vars_c ~f thn ~init:(fold_vars_c ~f els ~init))
  | `Trm t -> fold_vars_t ~f t ~init

let fold_vars ~init e ~f =
  match e with
  | `Fml p -> fold_vars_f ~f ~init p
  | #cnd as c -> fold_vars_c ~f ~init c

(** map *)

let map1 f e cons x =
  let x' = f x in
  if x == x' then e else cons x'

let map2 f e cons x y =
  let x' = f x in
  let y' = f y in
  if x == x' && y == y' then e else cons x' y'

let map3 f e cons x y z =
  let x' = f x in
  let y' = f y in
  let z' = f z in
  if x == x' && y == y' && z == z' then e else cons x' y' z'

let mapN f e cons xs =
  let xs' = Array.map_endo ~f xs in
  if xs' == xs then e else cons xs'

(** map_trms *)

let rec map_trms_f ~f b =
  match b with
  | Tt | Ff -> b
  | Eq (x, y) -> map2 f b _Eq x y
  | Dq (x, y) -> map2 f b _Dq x y
  | Eq0 x -> map1 f b _Eq0 x
  | Dq0 x -> map1 f b _Dq0 x
  | Gt0 x -> map1 f b _Gt0 x
  | Ge0 x -> map1 f b _Ge0 x
  | Lt0 x -> map1 f b _Lt0 x
  | Le0 x -> map1 f b _Le0 x
  | And (x, y) -> map2 (map_trms_f ~f) b _And x y
  | Or (x, y) -> map2 (map_trms_f ~f) b _Or x y
  | Iff (x, y) -> map2 (map_trms_f ~f) b _Iff x y
  | Xor (x, y) -> map2 (map_trms_f ~f) b _Xor x y
  | Cond {cnd; pos; neg} -> map3 (map_trms_f ~f) b _Cond cnd pos neg
  | UPosLit (p, xs) -> mapN f b (_UPosLit p) xs
  | UNegLit (p, xs) -> mapN f b (_UNegLit p) xs

(** map_vars *)

let rec map_vars_t ~f e =
  match e with
  | Var _ as v -> (f (Var.of_ v) : var :> trm)
  | Z _ | Q _ | Ancestor _ -> e
  | Neg x -> map1 (map_vars_t ~f) e _Neg x
  | Add (x, y) -> map2 (map_vars_t ~f) e _Add x y
  | Sub (x, y) -> map2 (map_vars_t ~f) e _Sub x y
  | Mulq (q, x) -> map1 (map_vars_t ~f) e (_Mulq q) x
  | Splat x -> map1 (map_vars_t ~f) e _Splat x
  | Sized {seq; siz} -> map2 (map_vars_t ~f) e _Sized seq siz
  | Extract {seq; off; len} -> map3 (map_vars_t ~f) e _Extract seq off len
  | Concat xs -> mapN (map_vars_t ~f) e _Concat xs
  | Select {idx; rcd} -> map1 (map_vars_t ~f) e (_Select idx) rcd
  | Update {idx; rcd; elt} -> map2 (map_vars_t ~f) e (_Update idx) rcd elt
  | Record xs -> mapN (map_vars_t ~f) e _Record xs
  | Apply (g, xs) -> mapN (map_vars_t ~f) e (_Apply g) xs

let map_vars_f ~f = map_trms_f ~f:(map_vars_t ~f)

let rec map_vars_c ~f c =
  match c with
  | `Ite (cnd, thn, els) ->
      let cnd' = map_vars_f ~f cnd in
      let thn' = map_vars_c ~f thn in
      let els' = map_vars_c ~f els in
      if cnd' == cnd && thn' == thn && els' == els then c
      else `Ite (cnd', thn', els')
  | `Trm t ->
      let t' = map_vars_t ~f t in
      if t' == t then c else `Trm t'

let map_vars ~f = function
  | `Fml p -> `Fml (map_vars_f ~f p)
  | #cnd as c -> (map_vars_c ~f c :> exp)

(*
 * Core construction functions
 *
 * Support functions for constructing expressions as if terms and formulas
 * could be freely mixed, instead of being strictly partitioned into terms
 * and formulas stratified below conditional terms and then expressions.
 *)

(** Map a unary function on terms over the leaves of a conditional term,
    rebuilding the tree of conditionals with the supplied ite construction
    function. *)
let rec map_cnd : (fml -> 'a -> 'a -> 'a) -> (trm -> 'a) -> cnd -> 'a =
 fun f_ite f_trm -> function
  | `Trm trm -> f_trm trm
  | `Ite (cnd, thn, els) ->
      let thn' = map_cnd f_ite f_trm thn in
      let els' = map_cnd f_ite f_trm els in
      f_ite cnd thn' els'

(** Embed a formula into a conditional term (associating true with 1 and
    false with 0), identity on conditional terms. *)
let embed_into_cnd : exp -> cnd = function
  | #cnd as c -> c
  (* p ==> (p ? 1 : 0) *)
  | `Fml fml -> `Ite (fml, `Trm one, `Trm zero)

(** Project out a formula that is embedded into a conditional term.

    - [project_out_fml] is left inverse to [embed_into_cnd] in the sense
      that [project_out_fml (embed_into_cnd (`Fml f)) = Some f]. *)
let project_out_fml : cnd -> fml option = function
  (* (p ? 1 : 0) ==> p *)
  | `Ite (cnd, `Trm one', `Trm zero') when one == one' && zero == zero' ->
      Some cnd
  | _ -> None

(** Embed a conditional term into a formula (associating 0 with false and
    non-0 with true, lifted over the tree mapping conditional terms to
    conditional formulas), identity on formulas.

    - [embed_into_fml] is left inverse to [embed_into_cnd] in the sense that
      [embed_into_fml ((embed_into_cnd (`Fml f)) :> exp) = f].
    - [embed_into_fml] is not right inverse to [embed_into_cnd] since
      [embed_into_fml] can only preserve one bit of information from its
      argument. So in general
      [(embed_into_cnd (`Fml (embed_into_fml x)) :> exp)] is not equivalent
      to [x].
    - The weaker condition that
      [0 ≠ (embed_into_cnd (`Fml (embed_into_fml x)) :> exp)] iff
      [0 ≠ x] holds. *)
let embed_into_fml : exp -> fml = function
  | `Fml fml -> fml
  | #cnd as c -> map_cnd _Cond _Dq0 c

(** Construct a conditional term, or formula if possible precisely. *)
let ite : fml -> exp -> exp -> exp =
 fun cnd thn els ->
  match (thn, els) with
  | `Fml pos, `Fml neg -> `Fml (_Cond cnd pos neg)
  | _ -> (
      let c = `Ite (cnd, embed_into_cnd thn, embed_into_cnd els) in
      match project_out_fml c with Some f -> `Fml f | None -> c )

(** Map a unary function on terms over an expression. *)
let ap1 : (trm -> exp) -> exp -> exp =
 fun f x -> map_cnd ite f (embed_into_cnd x)

let ap1t : (trm -> trm) -> exp -> exp = fun f -> ap1 (fun x -> `Trm (f x))

let ap1f : (trm -> fml) -> exp -> fml =
 fun f x -> map_cnd _Cond f (embed_into_cnd x)

(** Map a binary function on terms over conditional terms. This yields a
    conditional tree with the structure from the first argument where each
    leaf has been replaced by a conditional tree with the structure from the
    second argument where each leaf has been replaced by the application of
    the argument binary function to the corresponding leaves from the first
    and second argument. *)
let map2_cnd :
    (fml -> 'a -> 'a -> 'a) -> (trm -> trm -> 'a) -> cnd -> cnd -> 'a =
 fun f_ite f_trm x y ->
  map_cnd f_ite (fun x' -> map_cnd f_ite (fun y' -> f_trm x' y') y) x

(** Map a binary function on terms over expressions. *)
let ap2 : (trm -> trm -> exp) -> exp -> exp -> exp =
 fun f x y -> map2_cnd ite f (embed_into_cnd x) (embed_into_cnd y)

let ap2t : (trm -> trm -> trm) -> exp -> exp -> exp =
 fun f -> ap2 (fun x y -> `Trm (f x y))

let ap2f : (trm -> trm -> fml) -> exp -> exp -> fml =
 fun f x y -> map2_cnd _Cond f (embed_into_cnd x) (embed_into_cnd y)

(** Map a ternary function on terms over conditional terms. *)
let map3_cnd :
       (fml -> 'a -> 'a -> 'a)
    -> (trm -> trm -> trm -> 'a)
    -> cnd
    -> cnd
    -> cnd
    -> 'a =
 fun f_ite f_trm x y z ->
  map_cnd f_ite
    (fun x' ->
      map_cnd f_ite (fun y' -> map_cnd f_ite (fun z' -> f_trm x' y' z') z) y
      )
    x

(** Map a ternary function on terms over expressions. *)
let ap3 : (trm -> trm -> trm -> exp) -> exp -> exp -> exp -> exp =
 fun f x y z ->
  map3_cnd ite f (embed_into_cnd x) (embed_into_cnd y) (embed_into_cnd z)

let ap3t : (trm -> trm -> trm -> trm) -> exp -> exp -> exp -> exp =
 fun f -> ap3 (fun x y z -> `Trm (f x y z))

(** Reverse-map an nary function on terms over conditional terms. *)
let rev_mapN_cnd :
    (fml -> 'a -> 'a -> 'a) -> (trm list -> 'a) -> cnd list -> 'a =
 fun f_ite f_trms rev_xs ->
  let rec loop xs' = function
    | x :: xs -> map_cnd f_ite (fun x' -> loop (x' :: xs') xs) x
    | [] -> f_trms xs'
  in
  loop [] rev_xs

(** Map an nary function on terms over expressions. *)
let apNt : (trm array -> trm) -> exp array -> exp =
 fun f xs ->
  rev_mapN_cnd ite
    (fun xs -> `Trm (f (Array.of_list xs)))
    (Array.to_list_rev_map ~f:embed_into_cnd xs)

let apNf : (trm array -> fml) -> exp array -> fml =
 fun f xs ->
  rev_mapN_cnd _Cond
    (fun xs -> f (Array.of_list xs))
    (Array.to_list_rev_map ~f:embed_into_cnd xs)

(*
 * Terms: exposed interface
 *)

module Term = struct
  (* Exposed terms are represented as expressions, which allow formulas to
     appear at toplevel, although semantically these are redundant with
     their [inject]ion into [trm] proper. This redundancy of representation
     is allowed in order to avoid churning formulas back and forth between
     [fml] and [cnd] via [inject] and [project] in cases where formulas only
     transiently pass through term contexts. The construction functions will
     convert such a formula [p] into [(p ? 1 : 0)] as soon as it is used as
     a subterm, so this redundancy is only lazily delaying normalization by
     one step. *)
  module T = struct
    type t = exp [@@deriving compare, equal, sexp]
  end

  include T
  module Map = Map.Make (T)

  let ppx = ppx
  let pp = pp

  (* variables *)

  let var v = `Trm (v : var :> trm)

  (* constants *)

  let zero = `Trm zero
  let one = `Trm one

  let integer z =
    if Z.equal Z.zero z then zero
    else if Z.equal Z.one z then one
    else `Trm (Z z)

  let rational q = `Trm (Q q)

  (* arithmetic *)

  let neg = ap1t _Neg
  let add = ap2t _Add
  let sub = ap2t _Sub
  let mulq q = ap1t (_Mulq q)

  let mul =
    ap2 (fun x y ->
        match x with
        | Z z -> mulq (Q.of_z z) (`Trm y)
        | Q q -> mulq q (`Trm y)
        | _ -> (
          match y with
          | Z z -> mulq (Q.of_z z) (`Trm x)
          | Q q -> mulq q (`Trm x)
          | _ -> ap2t (fun x y -> Apply (Mul, [|x; y|])) (`Trm x) (`Trm y) ) )

  (* sequences *)

  let splat = ap1t _Splat
  let sized ~seq ~siz = ap2t _Sized seq siz
  let extract ~seq ~off ~len = ap3t _Extract seq off len
  let concat elts = apNt _Concat elts

  (* records *)

  let select ~rcd ~idx = ap1t (_Select idx) rcd
  let update ~rcd ~idx ~elt = ap2t (_Update idx) rcd elt
  let record elts = apNt _Record elts
  let ancestor i = `Trm (_Ancestor i)

  (* if-then-else *)

  let ite ~cnd ~thn ~els = ite cnd thn els

  (* uninterpreted *)

  let apply sym args = apNt (_Apply sym) args

  (** Destruct *)

  let d_int = function `Trm (Z z) -> Some z | _ -> None

  (** Access *)

  let const_of x =
    let rec const_of t =
      let neg = Option.map ~f:Q.neg in
      let add = Option.map2 ~f:Q.add in
      match t with
      | Z z -> Some (Q.of_z z)
      | Q q -> Some q
      | Neg x -> neg (const_of x)
      | Add (x, y) -> add (const_of x) (const_of y)
      | Sub (x, y) -> add (const_of x) (neg (const_of y))
      | _ -> None
    in
    match x with `Trm t -> const_of t | _ -> None

  (** Traverse *)

  let fold_vars = fold_vars

  (** Transform *)

  let map_vars = map_vars

  let fold_map_vars e ~init ~f =
    let s = ref init in
    let f x =
      let s', x' = f !s x in
      s := s' ;
      x'
    in
    let e' = map_vars ~f e in
    (!s, e')

  let rename s e = map_vars ~f:(Var.Subst.apply s) e

  (** Query *)

  let fv e = fold_vars e ~f:Var.Set.add ~init:Var.Set.empty
end

(*
 * Formulas: exposed interface
 *)

module Formula = struct
  type t = fml [@@deriving compare, equal, sexp]

  let inject f = `Fml f
  let project = function `Fml f -> Some f | #cnd as c -> project_out_fml c
  let ppx = ppx_f
  let pp = pp_f

  (* constants *)

  let tt = _Tt
  let ff = _Ff

  (* comparisons *)

  let eq = ap2f _Eq
  let dq = ap2f _Dq
  let eq0 = ap1f _Eq0
  let dq0 = ap1f _Dq0
  let gt0 = ap1f _Gt0
  let ge0 = ap1f _Ge0
  let lt0 = ap1f _Lt0
  let le0 = ap1f _Le0

  let gt a b =
    if a == Term.zero then lt0 b
    else if b == Term.zero then gt0 a
    else gt0 (Term.sub a b)

  let ge a b =
    if a == Term.zero then le0 b
    else if b == Term.zero then ge0 a
    else ge0 (Term.sub a b)

  let lt a b = gt b a
  let le a b = ge b a

  (* uninterpreted *)

  let uposlit p es = apNf (_UPosLit p) es
  let uneglit p es = apNf (_UNegLit p) es

  (* connectives *)

  let and_ = _And
  let andN = function [] -> tt | b :: bs -> List.fold ~init:b ~f:and_ bs
  let or_ = _Or
  let orN = function [] -> ff | b :: bs -> List.fold ~init:b ~f:or_ bs
  let iff = _Iff
  let xor = _Xor
  let cond ~cnd ~pos ~neg = _Cond cnd pos neg
  let not_ = _Not

  (** Query *)

  let fv e = fold_vars_f e ~f:Var.Set.add ~init:Var.Set.empty

  (** Traverse *)

  let fold_vars = fold_vars_f

  (** Transform *)

  let map_vars = map_vars_f

  let rec map_terms ~f b =
    let lift_map1 : (exp -> exp) -> t -> (trm -> t) -> trm -> t =
     fun f b cons x -> map1 f b (ap1f cons) (`Trm x)
    in
    let lift_map2 :
        (exp -> exp) -> t -> (trm -> trm -> t) -> trm -> trm -> t =
     fun f b cons x y -> map2 f b (ap2f cons) (`Trm x) (`Trm y)
    in
    let lift_mapN : (exp -> exp) -> t -> (trm array -> t) -> trm array -> t
        =
     fun f b cons xs ->
      mapN f b (apNf cons) (Array.map ~f:(fun x -> `Trm x) xs)
    in
    match b with
    | Tt | Ff -> b
    | Eq (x, y) -> lift_map2 f b _Eq x y
    | Dq (x, y) -> lift_map2 f b _Dq x y
    | Eq0 x -> lift_map1 f b _Eq0 x
    | Dq0 x -> lift_map1 f b _Dq0 x
    | Gt0 x -> lift_map1 f b _Gt0 x
    | Ge0 x -> lift_map1 f b _Ge0 x
    | Lt0 x -> lift_map1 f b _Lt0 x
    | Le0 x -> lift_map1 f b _Le0 x
    | And (x, y) -> map2 (map_terms ~f) b _And x y
    | Or (x, y) -> map2 (map_terms ~f) b _Or x y
    | Iff (x, y) -> map2 (map_terms ~f) b _Iff x y
    | Xor (x, y) -> map2 (map_terms ~f) b _Xor x y
    | Cond {cnd; pos; neg} -> map3 (map_terms ~f) b _Cond cnd pos neg
    | UPosLit (p, xs) -> lift_mapN f b (_UPosLit p) xs
    | UNegLit (p, xs) -> lift_mapN f b (_UNegLit p) xs

  let fold_map_vars ~init e ~f =
    let s = ref init in
    let f x =
      let s', x' = f !s x in
      s := s' ;
      x'
    in
    let e' = map_vars ~f e in
    (!s, e')

  let rename s e = map_vars ~f:(Var.Subst.apply s) e

  let fold_dnf :
         meet1:('literal -> 'conjunction -> 'conjunction)
      -> join1:('conjunction -> 'disjunction -> 'disjunction)
      -> top:'conjunction
      -> bot:'disjunction
      -> 'formula
      -> 'disjunction =
   fun ~meet1 ~join1 ~top ~bot fml ->
    let rec add_conjunct (cjn, splits) fml =
      match fml with
      | Tt | Ff | Eq _ | Dq _ | Eq0 _ | Dq0 _ | Gt0 _ | Ge0 _ | Lt0 _
       |Le0 _ | Iff _ | Xor _ | UPosLit _ | UNegLit _ ->
          (meet1 fml cjn, splits)
      | And (p, q) -> add_conjunct (add_conjunct (cjn, splits) p) q
      | Or (p, q) -> (cjn, [p; q] :: splits)
      | Cond {cnd; pos; neg} ->
          (cjn, [and_ cnd pos; and_ (not_ cnd) neg] :: splits)
    in
    let rec add_disjunct (cjn, splits) djn fml =
      let cjn, splits = add_conjunct (cjn, splits) fml in
      match splits with
      | split :: splits ->
          List.fold ~f:(add_disjunct (cjn, splits)) ~init:djn split
      | [] -> join1 cjn djn
    in
    add_disjunct (top, []) bot fml
end

(*
 * Convert to Ses
 *)

let v_to_ses : var -> Ses.Var.t =
 fun v -> Ses.Var.identified ~id:(Var.id v) ~name:(Var.name v)

let vs_to_ses : Var.Set.t -> Ses.Var.Set.t =
 fun vs ->
  Var.Set.fold vs ~init:Ses.Var.Set.empty ~f:(fun vs v ->
      Ses.Var.Set.add vs (v_to_ses v) )

let rec t_to_ses : trm -> Ses.Term.t = function
  | Var {name; id} -> Ses.Term.var (Ses.Var.identified ~name ~id)
  | Z z -> Ses.Term.integer z
  | Q q -> Ses.Term.rational q
  | Neg x -> Ses.Term.neg (t_to_ses x)
  | Add (x, y) -> Ses.Term.add (t_to_ses x) (t_to_ses y)
  | Sub (x, y) -> Ses.Term.sub (t_to_ses x) (t_to_ses y)
  | Mulq (q, x) -> Ses.Term.mulq q (t_to_ses x)
  | Splat x -> Ses.Term.splat (t_to_ses x)
  | Sized {seq; siz} ->
      Ses.Term.sized ~seq:(t_to_ses seq) ~siz:(t_to_ses siz)
  | Extract {seq; off; len} ->
      Ses.Term.extract ~seq:(t_to_ses seq) ~off:(t_to_ses off)
        ~len:(t_to_ses len)
  | Concat es -> Ses.Term.concat (Array.map ~f:t_to_ses es)
  | Select {idx; rcd} -> Ses.Term.select ~rcd:(t_to_ses rcd) ~idx
  | Update {idx; rcd; elt} ->
      Ses.Term.update ~rcd:(t_to_ses rcd) ~idx ~elt:(t_to_ses elt)
  | Record es ->
      Ses.Term.record (IArray.of_array (Array.map ~f:t_to_ses es))
  | Ancestor i -> Ses.Term.rec_record i
  | Apply (Mul, [|x; y|]) -> Ses.Term.mul (t_to_ses x) (t_to_ses y)
  | Apply (Div, [|x; y|]) -> Ses.Term.div (t_to_ses x) (t_to_ses y)
  | Apply (Rem, [|x; y|]) -> Ses.Term.rem (t_to_ses x) (t_to_ses y)
  | Apply (BitAnd, [|x; y|]) -> Ses.Term.and_ (t_to_ses x) (t_to_ses y)
  | Apply (BitOr, [|x; y|]) -> Ses.Term.or_ (t_to_ses x) (t_to_ses y)
  | Apply (BitXor, [|x; y|]) -> Ses.Term.dq (t_to_ses x) (t_to_ses y)
  | Apply (BitShl, [|x; y|]) -> Ses.Term.shl (t_to_ses x) (t_to_ses y)
  | Apply (BitLshr, [|x; y|]) -> Ses.Term.lshr (t_to_ses x) (t_to_ses y)
  | Apply (BitAshr, [|x; y|]) -> Ses.Term.ashr (t_to_ses x) (t_to_ses y)
  | Apply (Signed n, [|x|]) -> Ses.Term.signed n (t_to_ses x)
  | Apply (Unsigned n, [|x|]) -> Ses.Term.unsigned n (t_to_ses x)
  | Apply (sym, xs) ->
      Ses.Term.apply sym (IArray.of_array (Array.map ~f:t_to_ses xs))

let rec f_to_ses : fml -> Ses.Term.t = function
  | Tt -> Ses.Term.true_
  | Ff -> Ses.Term.false_
  | Eq (x, y) -> Ses.Term.eq (t_to_ses x) (t_to_ses y)
  | Dq (x, y) -> Ses.Term.dq (t_to_ses x) (t_to_ses y)
  | Eq0 x -> Ses.Term.eq Ses.Term.zero (t_to_ses x)
  | Dq0 x -> Ses.Term.dq Ses.Term.zero (t_to_ses x)
  | Gt0 x -> Ses.Term.lt Ses.Term.zero (t_to_ses x)
  | Ge0 x -> Ses.Term.le Ses.Term.zero (t_to_ses x)
  | Lt0 x -> Ses.Term.lt (t_to_ses x) Ses.Term.zero
  | Le0 x -> Ses.Term.le (t_to_ses x) Ses.Term.zero
  | And (p, q) -> Ses.Term.and_ (f_to_ses p) (f_to_ses q)
  | Or (p, q) -> Ses.Term.or_ (f_to_ses p) (f_to_ses q)
  | Iff (p, q) -> Ses.Term.eq (f_to_ses p) (f_to_ses q)
  | Xor (p, q) -> Ses.Term.dq (f_to_ses p) (f_to_ses q)
  | Cond {cnd; pos; neg} ->
      Ses.Term.conditional ~cnd:(f_to_ses cnd) ~thn:(f_to_ses pos)
        ~els:(f_to_ses neg)
  | UPosLit (sym, args) ->
      Ses.Term.poslit sym (IArray.of_array (Array.map ~f:t_to_ses args))
  | UNegLit (sym, args) ->
      Ses.Term.neglit sym (IArray.of_array (Array.map ~f:t_to_ses args))

let rec to_ses : exp -> Ses.Term.t = function
  | `Ite (cnd, thn, els) ->
      Ses.Term.conditional ~cnd:(f_to_ses cnd)
        ~thn:(to_ses (thn :> exp))
        ~els:(to_ses (els :> exp))
  | `Trm t -> t_to_ses t
  | `Fml f -> f_to_ses f

(*
 * Convert from Ses
 *)

let v_of_ses : Ses.Var.t -> var =
 fun v -> Var.identified ~id:(Ses.Var.id v) ~name:(Ses.Var.name v)

let vs_of_ses : Ses.Var.Set.t -> Var.Set.t =
 fun vs ->
  Ses.Var.Set.fold vs ~init:Var.Set.empty ~f:(fun vs v ->
      Var.Set.add vs (v_of_ses v) )

let uap1 f = ap1t (fun x -> Apply (f, [|x|]))
let uap2 f = ap2t (fun x y -> Apply (f, [|x; y|]))
let uposN p = apNf (_UPosLit p)
let unegN p = apNf (_UNegLit p)

let rec uap_tt f a = uap1 f (of_ses a)
and uap_ttt f a b = uap2 f (of_ses a) (of_ses b)

and ap2 mk_f mk_t a b =
  match (of_ses a, of_ses b) with
  | `Fml p, `Fml q -> `Fml (mk_f p q)
  | x, y -> mk_t x y

and ap2_f mk_f mk_t a b = ap2 mk_f (fun x y -> `Fml (mk_t x y)) a b

and apN mk_f mk_t mk_unit es =
  match
    Ses.Term.Set.fold ~init:(None, None) es ~f:(fun (fs, ts) e ->
        match of_ses e with
        | `Fml f ->
            (Some (match fs with None -> f | Some g -> mk_f f g), ts)
        | t -> (fs, Some (match ts with None -> t | Some u -> mk_t t u)) )
  with
  | Some f, Some t -> mk_t t (Formula.inject f)
  | Some f, None -> `Fml f
  | None, Some t -> t
  | None, None -> `Fml mk_unit

and of_ses : Ses.Term.t -> exp =
 fun t ->
  let open Term in
  let open Formula in
  match t with
  | Var {id; name} -> var (Var.identified ~id ~name)
  | Integer {data} -> integer data
  | Rational {data} -> rational data
  | Ap1 (Signed {bits}, e) -> uap_tt (Signed bits) e
  | Ap1 (Unsigned {bits}, e) -> uap_tt (Unsigned bits) e
  | Ap2 (Eq, d, e) -> ap2_f iff eq d e
  | Ap2 (Dq, d, e) -> ap2_f xor dq d e
  | Ap2 (Lt, d, e) -> ap2_f (fun p q -> and_ (not_ p) q) lt d e
  | Ap2 (Le, d, e) -> ap2_f (fun p q -> or_ (not_ p) q) le d e
  | PosLit (p, es) ->
      `Fml (uposN p (IArray.to_array (IArray.map ~f:of_ses es)))
  | NegLit (p, es) ->
      `Fml (unegN p (IArray.to_array (IArray.map ~f:of_ses es)))
  | Add sum -> (
    match Ses.Term.Qset.pop_min_elt sum with
    | None -> zero
    | Some (e, q, sum) ->
        let mul e q =
          if Q.equal Q.one q then of_ses e
          else
            match of_ses e with
            | `Trm (Z z) -> rational (Q.mul q (Q.of_z z))
            | `Trm (Q r) -> rational (Q.mul q r)
            | t -> mulq q t
        in
        Ses.Term.Qset.fold sum ~init:(mul e q) ~f:(fun e q s ->
            add (mul e q) s ) )
  | Mul prod -> (
    match Ses.Term.Qset.pop_min_elt prod with
    | None -> one
    | Some (e, q, prod) ->
        let rec expn e n =
          let p = Z.pred n in
          if Z.sign p = 0 then e else uap2 Mul e (expn e p)
        in
        let exp e q =
          let n = Q.num q in
          let sn = Z.sign n in
          if sn = 0 then of_ses e
          else if sn > 0 then expn (of_ses e) n
          else uap2 Div one (expn (of_ses e) (Z.neg n))
        in
        Ses.Term.Qset.fold prod ~init:(exp e q) ~f:(fun e q s ->
            uap2 Mul (exp e q) s ) )
  | Ap2 (Div, d, e) -> uap_ttt Div d e
  | Ap2 (Rem, d, e) -> uap_ttt Rem d e
  | And es -> apN and_ (uap2 BitAnd) tt es
  | Or es -> apN or_ (uap2 BitOr) ff es
  | Ap2 (Xor, d, e) -> ap2 xor (uap2 BitXor) d e
  | Ap2 (Shl, d, e) -> uap_ttt BitShl d e
  | Ap2 (Lshr, d, e) -> uap_ttt BitLshr d e
  | Ap2 (Ashr, d, e) -> uap_ttt BitAshr d e
  | Ap3 (Conditional, cnd, thn, els) -> (
      let cnd = embed_into_fml (of_ses cnd) in
      match (of_ses thn, of_ses els) with
      | `Fml pos, `Fml neg -> `Fml (cond ~cnd ~pos ~neg)
      | thn, els -> ite ~cnd ~thn ~els )
  | Ap1 (Splat, byt) -> splat (of_ses byt)
  | Ap3 (Extract, seq, off, len) ->
      extract ~seq:(of_ses seq) ~off:(of_ses off) ~len:(of_ses len)
  | Ap2 (Sized, siz, seq) -> sized ~seq:(of_ses seq) ~siz:(of_ses siz)
  | ApN (Concat, args) ->
      concat (Array.map ~f:of_ses (IArray.to_array args))
  | Ap1 (Select idx, rcd) -> select ~rcd:(of_ses rcd) ~idx
  | Ap2 (Update idx, rcd, elt) ->
      update ~rcd:(of_ses rcd) ~idx ~elt:(of_ses elt)
  | ApN (Record, elts) ->
      record (Array.map ~f:of_ses (IArray.to_array elts))
  | RecRecord i -> ancestor i
  | Apply (sym, args) ->
      apply sym (Array.map ~f:of_ses (IArray.to_array args))

let f_of_ses e = embed_into_fml (of_ses e)

let v_map_ses : (var -> var) -> Ses.Var.t -> Ses.Var.t =
 fun f x ->
  let v = v_of_ses x in
  let v' = f v in
  if v' == v then x else v_to_ses v'

let ses_map : (Ses.Term.t -> Ses.Term.t) -> exp -> exp =
 fun f x -> of_ses (f (to_ses x))

let f_ses_map : (Ses.Term.t -> Ses.Term.t) -> fml -> fml =
 fun f x -> f_of_ses (f (f_to_ses x))

(*
 * Contexts
 *)

module Context = struct
  type t = Ses.Equality.t [@@deriving sexp]

  let invariant = Ses.Equality.invariant

  (* Query *)

  let fold_vars ~init x ~f =
    Ses.Equality.fold_vars x ~init ~f:(fun s v -> f s (v_of_ses v))

  let fv e = fold_vars e ~f:Var.Set.add ~init:Var.Set.empty
  let is_empty x = Ses.Equality.is_true x
  let is_unsat x = Ses.Equality.is_false x
  let implies x b = Ses.Equality.implies x (f_to_ses b)

  let refutes x b =
    Ses.Term.is_false (Ses.Equality.normalize x (f_to_ses b))

  let normalize x e = ses_map (Ses.Equality.normalize x) e

  (* Classes *)

  let class_of x e = List.map ~f:of_ses (Ses.Equality.class_of x (to_ses e))

  let classes x =
    Ses.Term.Map.fold (Ses.Equality.classes x) ~init:Term.Map.empty
      ~f:(fun ~key:rep ~data:cls clss ->
        let rep' = of_ses rep in
        let cls' = List.map ~f:of_ses cls in
        Term.Map.set ~key:rep' ~data:cls' clss )

  let diff_classes r s =
    Term.Map.filter_mapi (classes r) ~f:(fun ~key:rep ~data:cls ->
        match
          List.filter cls ~f:(fun exp ->
              not (implies s (Formula.eq rep exp)) )
        with
        | [] -> None
        | cls -> Some cls )

  (* Pretty printing *)

  let pp_raw = Ses.Equality.pp
  let ppx_cls x = List.pp "@ = " (Term.ppx x)

  let ppx_classes x fs clss =
    List.pp "@ @<2>∧ "
      (fun fs (rep, cls) ->
        Format.fprintf fs "@[%a@ = %a@]" (Term.ppx x) rep (ppx_cls x)
          (List.sort ~compare:Term.compare cls) )
      fs (Term.Map.to_alist clss)

  let pp fs r = ppx_classes (fun _ -> None) fs (classes r)

  let ppx var_strength fs clss noneqs =
    let first = Term.Map.is_empty clss in
    if not first then Format.fprintf fs "  " ;
    ppx_classes var_strength fs clss ;
    List.pp
      ~pre:(if first then "@[  " else "@ @[@<2>∧ ")
      "@ @<2>∧ "
      (Formula.ppx var_strength)
      fs noneqs ~suf:"@]" ;
    first && List.is_empty noneqs

  let ppx_diff var_strength fs parent_ctx fml ctx =
    let fml' = f_ses_map (Ses.Equality.normalize ctx) fml in
    ppx var_strength fs
      (diff_classes ctx parent_ctx)
      (if Formula.(equal tt fml') then [] else [fml'])

  (* Construct *)

  let empty = Ses.Equality.true_

  let add vs f x =
    let vs', x' = Ses.Equality.and_term (vs_to_ses vs) (f_to_ses f) x in
    (vs_of_ses vs', x')

  let union vs x y =
    let vs', z = Ses.Equality.and_ (vs_to_ses vs) x y in
    (vs_of_ses vs', z)

  let inter vs x y =
    let vs', z = Ses.Equality.or_ (vs_to_ses vs) x y in
    (vs_of_ses vs', z)

  let interN vs xs =
    let vs', z = Ses.Equality.orN (vs_to_ses vs) xs in
    (vs_of_ses vs', z)

  let dnf f =
    let meet1 a (vs, p, x) =
      let vs, x = add vs a x in
      (vs, Formula.and_ p a, x)
    in
    let join1 = Iter.cons in
    let top = (Var.Set.empty, Formula.tt, empty) in
    let bot = Iter.empty in
    Formula.fold_dnf ~meet1 ~join1 ~top ~bot f

  let rename x sub = Ses.Equality.rename x (v_map_ses (Var.Subst.apply sub))

  (* Substs *)

  module Subst = struct
    type t = Ses.Equality.Subst.t [@@deriving sexp]

    let pp = Ses.Equality.Subst.pp
    let is_empty = Ses.Equality.Subst.is_empty

    let fold s ~init ~f =
      Ses.Equality.Subst.fold s ~init ~f:(fun ~key ~data ->
          f ~key:(of_ses key) ~data:(of_ses data) )

    let subst s = ses_map (Ses.Equality.Subst.subst s)

    let partition_valid vs s =
      let t, ks, u = Ses.Equality.Subst.partition_valid (vs_to_ses vs) s in
      (t, vs_of_ses ks, u)
  end

  let apply_subst vs s x =
    let vs', x' = Ses.Equality.apply_subst (vs_to_ses vs) s x in
    (vs_of_ses vs', x')

  let solve_for_vars vss x =
    Ses.Equality.solve_for_vars (List.map ~f:vs_to_ses vss) x

  let elim vs x = Ses.Equality.elim (vs_to_ses vs) x

  (* Replay debugging *)

  type call =
    | Add of Var.Set.t * fml * t
    | Union of Var.Set.t * t * t
    | Inter of Var.Set.t * t * t
    | InterN of Var.Set.t * t list
    | Rename of t * Var.Subst.t
    | Is_unsat of t
    | Implies of t * fml
    | Refutes of t * fml
    | Normalize of t * exp
    | Apply_subst of Var.Set.t * Subst.t * t
    | Solve_for_vars of Var.Set.t list * t
    | Elim of Var.Set.t * t
  [@@deriving sexp]

  let replay c =
    match call_of_sexp (Sexp.of_string c) with
    | Add (us, e, r) -> add us e r |> ignore
    | Union (us, r, s) -> union us r s |> ignore
    | Inter (us, r, s) -> inter us r s |> ignore
    | InterN (us, rs) -> interN us rs |> ignore
    | Rename (r, s) -> rename r s |> ignore
    | Is_unsat r -> is_unsat r |> ignore
    | Implies (r, f) -> implies r f |> ignore
    | Refutes (r, f) -> refutes r f |> ignore
    | Normalize (r, e) -> normalize r e |> ignore
    | Apply_subst (us, s, r) -> apply_subst us s r |> ignore
    | Solve_for_vars (vss, r) -> solve_for_vars vss r |> ignore
    | Elim (ks, r) -> elim ks r |> ignore

  (* Debug wrappers *)

  let report ~name ~elapsed ~aggregate ~count =
    Format.eprintf "%15s time: %12.3f ms  %12.3f ms  %12d calls@." name
      elapsed aggregate count

  let dump_threshold = ref 1000.

  let wrap tmr f call =
    let f () =
      Timer.start tmr ;
      let r = f () in
      Timer.stop_report tmr (fun ~name ~elapsed ~aggregate ~count ->
          report ~name ~elapsed ~aggregate ~count ;
          if Float.(elapsed > !dump_threshold) then (
            dump_threshold := 2. *. !dump_threshold ;
            Format.eprintf "@\n%a@\n@." Sexp.pp_hum (sexp_of_call (call ()))
            ) ) ;
      r
    in
    if not [%debug] then f ()
    else
      try f ()
      with exn ->
        let bt = Printexc.get_raw_backtrace () in
        let sexp = sexp_of_call (call ()) in
        raise (Replay (exn, bt, sexp))

  let add_tmr = Timer.create "add" ~at_exit:report
  let union_tmr = Timer.create "union" ~at_exit:report
  let inter_tmr = Timer.create "inter" ~at_exit:report
  let interN_tmr = Timer.create "interN" ~at_exit:report
  let rename_tmr = Timer.create "rename" ~at_exit:report
  let is_unsat_tmr = Timer.create "is_unsat" ~at_exit:report
  let implies_tmr = Timer.create "implies" ~at_exit:report
  let refutes_tmr = Timer.create "refutes" ~at_exit:report
  let normalize_tmr = Timer.create "normalize" ~at_exit:report
  let apply_subst_tmr = Timer.create "apply_subst" ~at_exit:report
  let solve_for_vars_tmr = Timer.create "solve_for_vars" ~at_exit:report
  let elim_tmr = Timer.create "elim" ~at_exit:report

  let add us e r =
    wrap add_tmr (fun () -> add us e r) (fun () -> Add (us, e, r))

  let union us r s =
    wrap union_tmr (fun () -> union us r s) (fun () -> Union (us, r, s))

  let inter us r s =
    wrap inter_tmr (fun () -> inter us r s) (fun () -> Inter (us, r, s))

  let interN us rs =
    wrap interN_tmr (fun () -> interN us rs) (fun () -> InterN (us, rs))

  let rename r s =
    wrap rename_tmr (fun () -> rename r s) (fun () -> Rename (r, s))

  let is_unsat r =
    wrap is_unsat_tmr (fun () -> is_unsat r) (fun () -> Is_unsat r)

  let implies r f =
    wrap implies_tmr (fun () -> implies r f) (fun () -> Implies (r, f))

  let refutes r f =
    wrap refutes_tmr (fun () -> refutes r f) (fun () -> Refutes (r, f))

  let normalize r e =
    wrap normalize_tmr (fun () -> normalize r e) (fun () -> Normalize (r, e))

  let apply_subst us s r =
    wrap apply_subst_tmr
      (fun () -> apply_subst us s r)
      (fun () -> Apply_subst (us, s, r))

  let solve_for_vars vss r =
    wrap solve_for_vars_tmr
      (fun () -> solve_for_vars vss r)
      (fun () -> Solve_for_vars (vss, r))

  let elim ks r =
    wrap elim_tmr (fun () -> elim ks r) (fun () -> Elim (ks, r))
end