Start supporting informative bool types in Coq backend

2 files changed, 201 insertions, 119 deletions

diff --git a/src/pretty_print_coq.ml b/src/pretty_print_coq.ml
index 802957c6..d54b2264 100644
--- a/src/pretty_print_coq.ml
+++ b/src/pretty_print_coq.ml
@@ -84,7 +84,7 @@ type context = {
   kid_renames : kid KBindings.t; (* Plain tyvar -> tyvar renames *)
   kid_id_renames : id KBindings.t; (* tyvar -> argument renames *)
   bound_nvars : KidSet.t;
-  build_ex_return : bool;
+  build_at_return : string option;
   recursive_ids : IdSet.t;
   debug : bool;
 }
@@ -93,7 +93,7 @@ let empty_ctxt = {
   kid_renames = KBindings.empty;
   kid_id_renames = KBindings.empty;
   bound_nvars = KidSet.empty;
-  build_ex_return = false;
+  build_at_return = None;
   recursive_ids = IdSet.empty;
   debug = false;
 }
@@ -272,6 +272,27 @@ let rec orig_nexp (Nexp_aux (nexp, l)) =
   | Nexp_neg n -> rewrap (Nexp_neg (orig_nexp n))
   | _ -> rewrap nexp
 
+let rec orig_nc (NC_aux (nc, l) as full_nc) =
+  let rewrap nc = NC_aux (nc, l) in
+  match nc with
+  | NC_equal (nexp1, nexp2) -> rewrap (NC_equal (orig_nexp nexp1, orig_nexp nexp2))
+  | NC_bounded_ge (nexp1, nexp2) -> rewrap (NC_bounded_ge (orig_nexp nexp1, orig_nexp nexp2))
+  | NC_bounded_le (nexp1, nexp2) -> rewrap (NC_bounded_le (orig_nexp nexp1, orig_nexp nexp2))
+  | NC_not_equal (nexp1, nexp2) -> rewrap (NC_not_equal (orig_nexp nexp1, orig_nexp nexp2))
+  | NC_set (kid,s) -> rewrap (NC_set (orig_kid kid, s))
+  | NC_or (nc1, nc2) -> rewrap (NC_or (orig_nc nc1, orig_nc nc2))
+  | NC_and (nc1, nc2) -> rewrap (NC_and (orig_nc nc1, orig_nc nc2))
+  | NC_app (f,args) -> rewrap (NC_app (f,List.map orig_typ_arg args))
+  | NC_var kid -> rewrap (NC_var (orig_kid kid))
+  | NC_true | NC_false -> full_nc
+and orig_typ_arg (A_aux (arg,l)) =
+  let rewrap a = (A_aux (a,l)) in
+  match arg with
+  | A_nexp nexp -> rewrap (A_nexp (orig_nexp nexp))
+  | A_bool nc -> rewrap (A_bool (orig_nc nc))
+  | A_order _ | A_typ _ ->
+       raise (Reporting.err_unreachable l __POS__ "Tried to pass Type or Order kind to SMT function")
+
 (* Returns the set of type variables that will appear in the Coq output,
    which may be smaller than those in the Sail type.  May need to be
    updated with doc_typ *)
@@ -289,6 +310,7 @@ let rec coq_nvars_of_typ (Typ_aux (t,l)) =
   | Typ_app(Id_aux (Id "implicit", _),_)
   (* TODO: update when complex atom types are sorted out *)
   | Typ_app(Id_aux (Id "atom", _), _) -> KidSet.empty
+  | Typ_app(Id_aux (Id "atom_bool", _), _) -> KidSet.empty
   | Typ_app (_,tas) ->
      List.fold_left (fun s ta -> KidSet.union s (coq_nvars_of_typ_arg ta))
        KidSet.empty tas
@@ -301,71 +323,7 @@ and coq_nvars_of_typ_arg (A_aux (ta,_)) =
   | A_nexp nexp -> tyvars_of_nexp (orig_nexp nexp)
   | A_typ typ -> coq_nvars_of_typ typ
   | A_order _ -> KidSet.empty
-
-(* Follows Coq precedence levels *)
-let rec doc_nc_prop ctx nc =
-  let rec l85 (NC_aux (nc,_) as nc_full) =
-  match nc with
-  | NC_or (nc1, nc2) -> doc_op (string "\\/") (doc_nc_prop ctx nc1) (doc_nc_prop ctx nc2)
-  | _ -> l80 nc_full
-  and l80 (NC_aux (nc,_) as nc_full) =
-  match nc with
-  | NC_and (nc1, nc2) -> doc_op (string "/\\") (doc_nc_prop ctx nc1) (doc_nc_prop ctx nc2)
-  | _ -> l70 nc_full
-  and l70 (NC_aux (nc,_) as nc_full) =
-  match nc with
-  | NC_equal (ne1, ne2) -> doc_op equals (doc_nexp ctx ne1) (doc_nexp ctx ne2)
-  | NC_bounded_ge (ne1, ne2) -> doc_op (string ">=") (doc_nexp ctx ne1) (doc_nexp ctx ne2)
-  | NC_bounded_le (ne1, ne2) -> doc_op (string "<=") (doc_nexp ctx ne1) (doc_nexp ctx ne2)
-  | NC_not_equal (ne1, ne2) -> doc_op (string "<>") (doc_nexp ctx ne1) (doc_nexp ctx ne2)
-  | _ -> l10 nc_full
-  and l10 (NC_aux (nc,_) as nc_full) =
-  match nc with
-  | NC_set (kid, is) ->
-     separate space [string "In"; doc_var ctx kid;
-                     brackets (separate (string "; ")
-                                 (List.map (fun i -> string (Nat_big_num.to_string i)) is))]
-  | NC_true -> string "True"
-  | NC_false -> string "False"
-  | NC_or _
-  | NC_and _
-  | NC_equal _
-  | NC_bounded_ge _
-  | NC_bounded_le _
-  | NC_not_equal _ -> parens (l85 nc_full)
-  in l85 nc
-
-(* Follows Coq precedence levels *)
-let doc_nc_exp ctx nc =
-  let rec l70 (NC_aux (nc,_) as nc_full) =
-    match nc with
-    | NC_equal (ne1, ne2) -> doc_op (string "=?") (doc_nexp ctx ne1) (doc_nexp ctx ne2)
-    | NC_bounded_ge (ne1, ne2) -> doc_op (string ">=?") (doc_nexp ctx ne1) (doc_nexp ctx ne2)
-    | NC_bounded_le (ne1, ne2) -> doc_op (string "<=?") (doc_nexp ctx ne1) (doc_nexp ctx ne2)
-    | _ -> l50 nc_full
-  and l50 (NC_aux (nc,_) as nc_full) =
-    match nc with
-    | NC_or (nc1, nc2) -> doc_op (string "||") (l50 nc1) (l40 nc2)
-    | _ -> l40 nc_full
-  and l40 (NC_aux (nc,_) as nc_full) =
-    match nc with
-    | NC_and (nc1, nc2) -> doc_op (string "&&") (l40 nc1) (l10 nc2)
-    | _ -> l10 nc_full
-  and l10 (NC_aux (nc,_) as nc_full) =
-    match nc with
-    | NC_not_equal (ne1, ne2) -> string "negb" ^^ space ^^ parens (doc_op (string "=?") (doc_nexp ctx ne1) (doc_nexp ctx ne2))
-    | NC_set (kid, is) ->
-       separate space [string "member_Z_list"; doc_var ctx kid;
-                       brackets (separate (string "; ")
-                                   (List.map (fun i -> string (Nat_big_num.to_string i)) is))]
-    | NC_true -> string "true"
-    | NC_false -> string "false"
-    | NC_equal _
-    | NC_bounded_ge _
-    | NC_bounded_le _
-    | NC_or _
-    | NC_and _ -> parens (l70 nc_full)
-  in l70 nc
+  | A_bool nc -> tyvars_of_constraint (orig_nc nc)
 
 let maybe_expand_range_type (Typ_aux (typ,l) as full_typ) =
   match typ with
@@ -385,18 +343,6 @@ let maybe_expand_range_type (Typ_aux (typ,l) as full_typ) =
 
 let expand_range_type typ = Util.option_default typ (maybe_expand_range_type typ)
 
-let doc_arithfact ctxt ?(exists = []) ?extra nc =
-  let prop = doc_nc_prop ctxt nc in
-  let prop = match extra with
-    | None -> prop
-    | Some pp -> separate space [pp; string "/\\"; prop]
-  in
-  let prop =
-    match exists with
-    | [] -> prop
-    | _ -> separate space ([string "exists"]@(List.map (doc_var ctxt) exists)@[comma; prop])
-  in
-  string "ArithFact" ^^ space ^^ parens prop
 
 let nice_and nc1 nc2 =
 match nc1, nc2 with
@@ -404,9 +350,20 @@ match nc1, nc2 with
 | _, NC_aux (NC_true,_) -> nc1
 | _,_ -> nc_and nc1 nc2
 
+(* n_constraint functions are currently just Z3 functions *)
+let doc_nc_fn_prop id =
+  match string_of_id id with
+  | "not" -> string "not"
+  | s -> string s
+
+(* n_constraint functions are currently just Z3 functions *)
+let doc_nc_fn id =
+  match string_of_id id with
+  | "not" -> string "negb"
+  | s -> string s
+
 (* When making changes here, check whether they affect coq_nvars_of_typ *)
-let doc_typ, doc_atomic_typ =
-  let fns ctx =
+let rec doc_typ_fns ctx =
   (* following the structure of parser for precedence *)
   let rec typ ty = fn_typ true ty
     and typ' ty = fn_typ false ty
@@ -448,6 +405,10 @@ let doc_typ, doc_atomic_typ =
          (string "Z")
       | Typ_app(Id_aux (Id "atom", _), [A_aux(A_nexp n,_)]) ->
          (string "Z")
+      | Typ_app(Id_aux (Id "atom_bool", _), [_]) -> string "bool"
+      | Typ_app (Id_aux (Id "atom#bool",_), [A_aux (A_bool nc,_)]) ->
+         let tpp = string "Bool" ^^ space ^^ doc_nc_prop ~top:false ctx nc in
+         if atyp_needed then parens tpp else tpp
       | Typ_app(id,args) ->
          let tpp = (doc_id_type id) ^^ space ^^ (separate_map space doc_typ_arg args) in
          if atyp_needed then parens tpp else tpp
@@ -536,8 +497,106 @@ let doc_typ, doc_atomic_typ =
       | A_typ t -> app_typ true t
       | A_nexp n -> doc_nexp ctx n
       | A_order o -> empty
-  in typ', atomic_typ
-  in (fun ctx -> (fst (fns ctx))), (fun ctx -> (snd (fns ctx)))
+      | A_bool nc -> doc_nc_prop ~top:false ctx nc
+  in typ', atomic_typ, doc_typ_arg
+and doc_typ ctx = let f,_,_ = doc_typ_fns ctx in f
+and doc_atomic_typ ctx = let _,f,_ = doc_typ_fns ctx in f
+and doc_typ_arg ctx = let _,_,f = doc_typ_fns ctx in f
+
+and doc_arithfact ctxt ?(exists = []) ?extra nc =
+  let prop = doc_nc_prop ctxt nc in
+  let prop = match extra with
+    | None -> prop
+    | Some pp -> separate space [pp; string "/\\"; prop]
+  in
+  let prop =
+    match exists with
+    | [] -> prop
+    | _ -> separate space ([string "exists"]@(List.map (doc_var ctxt) exists)@[comma; prop])
+  in
+  string "ArithFact" ^^ space ^^ parens prop
+
+(* Follows Coq precedence levels *)
+and doc_nc_prop ?(top = true) ctx nc =
+  let rec l85 (NC_aux (nc,_) as nc_full) =
+  match nc with
+  | NC_or (nc1, nc2) -> doc_op (string "\\/") (doc_nc_prop ctx nc1) (doc_nc_prop ctx nc2)
+  | _ -> l80 nc_full
+  and l80 (NC_aux (nc,_) as nc_full) =
+  match nc with
+  | NC_and (nc1, nc2) -> doc_op (string "/\\") (doc_nc_prop ctx nc1) (doc_nc_prop ctx nc2)
+  | _ -> l70 nc_full
+  and l70 (NC_aux (nc,_) as nc_full) =
+  match nc with
+  | NC_equal (ne1, ne2) -> doc_op equals (doc_nexp ctx ne1) (doc_nexp ctx ne2)
+  | NC_var kid -> doc_op equals (doc_nexp ctx (nvar kid)) (string "true")
+  | NC_bounded_ge (ne1, ne2) -> doc_op (string ">=") (doc_nexp ctx ne1) (doc_nexp ctx ne2)
+  | NC_bounded_le (ne1, ne2) -> doc_op (string "<=") (doc_nexp ctx ne1) (doc_nexp ctx ne2)
+  | NC_not_equal (ne1, ne2) -> doc_op (string "<>") (doc_nexp ctx ne1) (doc_nexp ctx ne2)
+  | _ -> l10 nc_full
+  and l10 (NC_aux (nc,_) as nc_full) =
+  match nc with
+  | NC_set (kid, is) ->
+     separate space [string "In"; doc_var ctx kid;
+                     brackets (separate (string "; ")
+                                 (List.map (fun i -> string (Nat_big_num.to_string i)) is))]
+  | NC_app (f,args) -> separate space (doc_nc_fn_prop f::List.map (doc_typ_arg ctx) args)
+  | _ -> l0 nc_full
+  and l0 (NC_aux (nc,_) as nc_full) =
+  match nc with
+  | NC_true -> string "True"
+  | NC_false -> string "False"
+  | NC_set _
+  | NC_app _
+  | NC_var _
+  | NC_or _
+  | NC_and _
+  | NC_equal _
+  | NC_bounded_ge _
+  | NC_bounded_le _
+  | NC_not_equal _ -> parens (l85 nc_full)
+  in if top then l85 nc else l0 nc
+
+(* Follows Coq precedence levels *)
+let rec doc_nc_exp ctx env nc =
+  let nc = Env.expand_constraint_synonyms env nc in
+  let rec l70 (NC_aux (nc,_) as nc_full) =
+    match nc with
+    | NC_equal (ne1, ne2) -> doc_op (string "=?") (doc_nexp ctx ne1) (doc_nexp ctx ne2)
+    | NC_bounded_ge (ne1, ne2) -> doc_op (string ">=?") (doc_nexp ctx ne1) (doc_nexp ctx ne2)
+    | NC_bounded_le (ne1, ne2) -> doc_op (string "<=?") (doc_nexp ctx ne1) (doc_nexp ctx ne2)
+    | _ -> l50 nc_full
+  and l50 (NC_aux (nc,_) as nc_full) =
+    match nc with
+    | NC_or (nc1, nc2) -> doc_op (string "||") (l50 nc1) (l40 nc2)
+    | _ -> l40 nc_full
+  and l40 (NC_aux (nc,_) as nc_full) =
+    match nc with
+    | NC_and (nc1, nc2) -> doc_op (string "&&") (l40 nc1) (l10 nc2)
+    | _ -> l10 nc_full
+  and l10 (NC_aux (nc,_) as nc_full) =
+    match nc with
+    | NC_not_equal (ne1, ne2) -> string "negb" ^^ space ^^ parens (doc_op (string "=?") (doc_nexp ctx ne1) (doc_nexp ctx ne2))
+    | NC_set (kid, is) ->
+       separate space [string "member_Z_list"; doc_var ctx kid;
+                       brackets (separate (string "; ")
+                                   (List.map (fun i -> string (Nat_big_num.to_string i)) is))]
+    | NC_true -> string "true"
+    | NC_false -> string "false"
+    | NC_app (f,args) -> separate space (doc_nc_fn f::List.map (doc_typ_arg_exp ctx env) args)
+    | NC_var kid -> doc_nexp ctx (nvar kid)
+    | NC_equal _
+    | NC_bounded_ge _
+    | NC_bounded_le _
+    | NC_or _
+    | NC_and _ -> parens (l70 nc_full)
+  in l70 nc
+and doc_typ_arg_exp ctx env (A_aux (arg,l)) =
+  match arg with
+  | A_nexp nexp -> doc_nexp ctx nexp
+  | A_bool nc -> doc_nc_exp ctx env nc
+  | A_order _ | A_typ _ ->
+       raise (Reporting.err_unreachable l __POS__ "Tried to pass Type or Order kind to SMT function")
 
 (* Check for variables in types that would be pretty-printed and are not
    bound in the val spec of the function. *)
@@ -619,6 +678,7 @@ let doc_quant_item_id ctx delimit (QI_aux (qi,_)) =
     | K_type -> Some (delimit (separate space [doc_var ctx kid; colon; string "Type"]))
     | K_int -> Some (delimit (separate space [doc_var ctx kid; colon; string "Z"]))
     | K_order -> None
+    | K_bool -> Some (delimit (separate space [doc_var ctx kid; colon; string "bool"]))
   end
   | QI_const nc -> None
 
@@ -630,6 +690,7 @@ let quant_item_id_name ctx (QI_aux (qi,_)) =
     | K_type -> Some (doc_var ctx kid)
     | K_int -> Some (doc_var ctx kid)
     | K_order -> None
+    | K_bool -> Some (doc_var ctx kid)
   end
   | QI_const nc -> None
 
@@ -888,13 +949,15 @@ let condition_produces_constraint exp =
    dependent pair with a proof that the result is the expected integer.  This
    is redundant for basic arithmetic functions and functions which we unfold
    in the constraint solver. *)
-let no_Z_proof_fns = ["Z.add"; "Z.sub"; "Z.opp"; "Z.mul"; "length_mword"; "length"]
+let no_proof_fns = ["Z.add"; "Z.sub"; "Z.opp"; "Z.mul"; "length_mword"; "length";
+                      "negb"; "andb"; "orb";
+                      "Z.leb"; "Z.geb"; "Z.ltb"; "Z.gtb"; "Z.eqb"]
 
-let is_no_Z_proof_fn env id =
+let is_no_proof_fn env id =
   if Env.is_extern id env "coq"
   then
     let s = Env.get_extern id env "coq" in
-    List.exists (fun x -> String.compare x s == 0) no_Z_proof_fns
+    List.exists (fun x -> String.compare x s == 0) no_proof_fns
   else false
 
 let replace_atom_return_type ret_typ =
@@ -902,8 +965,12 @@ let replace_atom_return_type ret_typ =
   match ret_typ with
   | Typ_aux (Typ_app (Id_aux (Id "atom",_), [A_aux (A_nexp nexp,_)]),l) ->
      let kid = mk_kid "_retval" in (* TODO: collision avoidance *)
-     true, Typ_aux (Typ_exist ([mk_kopt K_int kid], nc_eq (nvar kid) nexp, atom_typ (nvar kid)),Parse_ast.Generated l)
-  | _ -> false, ret_typ
+     Some "build_ex", Typ_aux (Typ_exist ([mk_kopt K_int kid], nc_eq (nvar kid) nexp, atom_typ (nvar kid)),Parse_ast.Generated l)
+  (* For informative booleans tweak the type name so that doc_typ knows that the
+     constraint should be output. *)
+  | Typ_aux (Typ_app (Id_aux (Id "atom_bool",il), ([A_aux (A_bool _,_)] as args)),l) ->
+     Some "build_Bool", Typ_aux (Typ_app (Id_aux (Id "atom#bool",il), args),l)
+  | _ -> None, ret_typ
 
 let is_range_from_atom env (Typ_aux (argty,_)) (Typ_aux (fnty,_)) =
   match argty, fnty with
@@ -1083,14 +1150,12 @@ let doc_exp, doc_let =
             match args with
             | [from_exp; to_exp; step_exp; ord_exp; vartuple; body] ->
                let loopvar, body = match body with
-                 | E_aux (E_let (LB_aux (LB_val (_, _), _),
-                   E_aux (E_let (LB_aux (LB_val (_, _), _),
-                   E_aux (E_if (_,
+                 | E_aux (E_if (_,
                    E_aux (E_let (LB_aux (LB_val (
                      ((P_aux (P_typ (_, P_aux (P_var (P_aux (P_id id, _), _), _)), _))
                      | (P_aux (P_var (P_aux (P_id id, _), _), _))
                      | (P_aux (P_id id, _))), _), _),
-                     body), _), _), _)), _)), _) -> id, body
+                     body), _), _), _) -> id, body
                  | _ -> raise (Reporting.err_unreachable l __POS__ ("Unable to find loop variable in " ^ string_of_exp body)) in
                let dir = match ord_exp with
                  | E_aux (E_lit (L_aux (L_false, _)), _) -> "_down"
@@ -1169,9 +1234,9 @@ let doc_exp, doc_let =
             match args with
             | [exp] ->
                let exp_pp =
-                 if ctxt.build_ex_return
-                 then parens (string "build_ex" ^/^ expY exp)
-                 else expY exp
+                 match ctxt.build_at_return with
+                 | Some s -> parens (string s ^/^ expY exp)
+                 | None -> expY exp
                in
                let epp = separate space [string "early_return"; exp_pp] in
                let aexp_needed, tepp =
@@ -1307,12 +1372,12 @@ let doc_exp, doc_let =
           let ret_typ_inst =
                subst_unifiers inst ret_typ
           in
-          let packeff,unpack,autocast =
+          let packeff,unpack,autocast,projbool =
             let ann_typ = Env.expand_synonyms env (general_typ_of_annot (l,annot)) in
             let ann_typ = expand_range_type ann_typ in
             let ret_typ_inst = expand_range_type (Env.expand_synonyms env ret_typ_inst) in
             let ret_typ_inst =
-              if is_no_Z_proof_fn env f then ret_typ_inst
+              if is_no_proof_fn env f then ret_typ_inst
               else snd (replace_atom_return_type ret_typ_inst) in
             let () = 
               debug ctxt (lazy (" type returned " ^ string_of_typ ret_typ_inst));
@@ -1336,13 +1401,19 @@ let doc_exp, doc_let =
                 Typ_aux (Typ_app (_,[A_aux (A_nexp n2,_);_;_]),_) ->
                  not (similar_nexps ctxt env n1 n2)
               | _ -> false
-            in pack,unpack,autocast
+            in
+            let projbool =
+              match in_typ with
+              | Typ_aux (Typ_app (Id_aux (Id "atom#bool",_),_),_) -> true
+              | _ -> false
+            in pack,unpack,autocast,projbool
           in
           let autocast_id, proj_id =
             if effectful eff
             then "autocast_m", "projT1_m"
             else "autocast",   "projT1" in
           let epp = if unpack && not (effectful eff) then string proj_id ^^ space ^^ parens epp else epp in
+          let epp = if projbool && not (effectful eff) then string "projBool" ^^ space ^^ parens epp else epp in
           let epp = if autocast then string autocast_id ^^ space ^^ parens epp else epp in
           let epp =
             if effectful eff && packeff && not unpack
@@ -1665,9 +1736,9 @@ let doc_exp, doc_let =
     | E_return r ->
       let ret_monad = " : MR" in
       let exp_pp =
-        if ctxt.build_ex_return
-        then parens (string "build_ex" ^/^ expY r)
-        else expY r
+        match ctxt.build_at_return with
+        | Some s -> parens (string s ^/^ expY r)
+        | None -> expY r
       in
       let ta =
         if contains_t_pp_var ctxt (typ_of full_exp) || contains_t_pp_var ctxt (typ_of r)
@@ -1677,7 +1748,7 @@ let doc_exp, doc_let =
           parens (doc_typ ctxt (typ_of full_exp));
           parens (doc_typ ctxt (typ_of r))] in
       align (parens (string "early_return" ^//^ exp_pp ^//^ ta))
-    | E_constraint nc -> wrap_parens (doc_nc_exp ctxt nc)
+    | E_constraint nc -> wrap_parens (doc_nc_exp ctxt (env_of full_exp) nc)
     | E_internal_value _ ->
       raise (Reporting.err_unreachable l __POS__
         "unsupported internal expression encountered while pretty-printing")
@@ -1818,6 +1889,8 @@ let doc_typdef generic_eq_types (TD_aux(td, (l, annot))) = match td with
                         doc_typquant_items empty_ctxt parens typq;
                         colon; string "Type"])
        (doc_typschm empty_ctxt false typschm) ^^ dot
+  | TD_abbrev _ -> empty (* TODO? *)
+  | TD_bitfield _ -> empty (* TODO? *)
   | TD_record(id,typq,fs,_) ->
     let fname fid = if prefix_recordtype && string_of_id id <> "regstate"
                     then concat [doc_id id;string "_";doc_id_type fid;]
@@ -1916,7 +1989,6 @@ let doc_typdef generic_eq_types (TD_aux(td, (l, annot))) = match td with
            string "forall (x y : " ^^ id_pp ^^ string "), Decidable (x = y) :=" ^/^
            string "Decidable_eq_from_dec " ^^ id_pp ^^ string "_eq_dec." in
           typ_pp ^^ dot ^^ hardline ^^ eq1_pp ^^ hardline ^^ eq2_pp ^^ hardline)
-    | _ -> raise (Reporting.err_unreachable l __POS__ "register with non-constant indices")
 
 let args_of_typ l env typs =
   let arg i typ =
@@ -2065,18 +2137,23 @@ let merge_kids_atoms pats =
   let try_eliminate (gone,map,seen) = function
     | P_aux (P_id id, ann), typ
     | P_aux (P_typ (_,P_aux (P_id id, ann)),_), typ -> begin
-      match Type_check.destruct_atom_nexp (env_of_annot ann) typ with
-      | Some (Nexp_aux (Nexp_var kid,l)) ->
-         if KidSet.mem kid seen then
-           let () =
-             Reporting.print_err l "merge_kids_atoms"
-               ("want to merge tyvar and argument for " ^ string_of_kid kid ^
+        let merge kid l =
+          if KidSet.mem kid seen then
+            let () =
+              Reporting.print_err l "merge_kids_atoms"
+                ("want to merge tyvar and argument for " ^ string_of_kid kid ^
                    " but rearranging arguments isn't supported yet") in
-           gone,map,seen
-         else
-           KidSet.add kid gone, KBindings.add kid id map, KidSet.add kid seen
-      | _ -> gone,map,KidSet.union seen (tyvars_of_typ typ)
-    end
+            gone,map,seen
+          else
+            KidSet.add kid gone, KBindings.add kid id map, KidSet.add kid seen
+        in
+        match Type_check.destruct_atom_nexp (env_of_annot ann) typ with
+        | Some (Nexp_aux (Nexp_var kid,l)) -> merge kid l
+        | _ ->
+           match Type_check.destruct_atom_bool (env_of_annot ann) typ with
+           | Some (NC_aux (NC_var kid,l)) -> merge kid l
+           | _ -> gone,map,KidSet.union seen (tyvars_of_typ typ)
+      end
     | _, typ -> gone,map,KidSet.union seen (tyvars_of_typ typ)
   in
   let gone,map,_ = List.fold_left try_eliminate (KidSet.empty, KBindings.empty, KidSet.empty) pats in
@@ -2100,7 +2177,7 @@ let doc_funcl rec_opt (FCL_aux(FCL_Funcl(id, pexp), annot)) =
   in
   let build_ex, ret_typ = replace_atom_return_type ret_typ in
   let build_ex = match destruct_exist_plain (Env.expand_synonyms env (expand_range_type ret_typ)) with
-    | Some _ -> true
+    | Some _ -> Some "build_ex"
     | _ -> build_ex
   in
   let ids_to_avoid = all_ids pexp in
@@ -2129,15 +2206,18 @@ let doc_funcl rec_opt (FCL_aux(FCL_Funcl(id, pexp), annot)) =
       kid_renames = mk_kid_renames ids_to_avoid kids_used;
       kid_id_renames = kid_to_arg_rename;
       bound_nvars = bound_kids;
-      build_ex_return = effectful eff && build_ex;
+      build_at_return = if effectful eff then build_ex else None;
       recursive_ids = recursive_ids;
       debug = List.mem (string_of_id id) (!opt_debug_on)
     } in
   let () =
     debug ctxt (lazy ("Function " ^ string_of_id id));
     debug ctxt (lazy (" return type " ^ string_of_typ ret_typ));
-    debug ctxt (lazy (" build_ex " ^ if build_ex then "needed" else "not needed"));
-    debug ctxt (lazy (if effectful eff then " effectful" else " pure"))
+    debug ctxt (lazy (" build_ex " ^ match build_ex with Some s -> s ^ " needed" | _ -> "not needed"));
+    debug ctxt (lazy (if effectful eff then " effectful" else " pure"));
+    debug ctxt (lazy (" kid_id_renames " ^ String.concat ", " (List.map
+                        (fun (kid,id) -> string_of_kid kid ^ " |-> " ^ string_of_id id)
+                        (KBindings.bindings kid_to_arg_rename))))
   in
   (* Put the constraints after pattern matching so that any type variable that's
      been replaced by one of the term-level arguments is bound. *)
@@ -2228,7 +2308,7 @@ let doc_funcl rec_opt (FCL_aux(FCL_Funcl(id, pexp), annot)) =
        raise (Reporting.err_unreachable l __POS__
                "guarded pattern expression should have been rewritten before pretty-printing") in
   let bodypp = doc_fun_body ctxt exp in
-  let bodypp = if effectful eff || not build_ex then bodypp else string "build_ex" ^^ parens bodypp in
+  let bodypp = if effectful eff then bodypp else match build_ex with Some s -> string s ^^ parens bodypp | None -> bodypp in
   let bodypp = separate (break 1) fixupspp ^/^ bodypp in
   group (prefix 3 1
     (flow (break 1) ([intropp; idpp] @ quantspp @ [patspp] @ constrspp @ [atom_constr_pp] @ accpp) ^/^
diff --git a/src/type_check.mli b/src/type_check.mli
index 7a5a3446..522074b3 100644
--- a/src/type_check.mli
+++ b/src/type_check.mli
@@ -358,6 +358,8 @@ val expected_typ_of : Ast.l * tannot -> typ option
 
 val destruct_atom_nexp : Env.t -> typ -> nexp option
 
+val destruct_atom_bool : Env.t -> typ -> n_constraint option
+
 val destruct_range : Env.t -> typ -> (kid list * n_constraint * nexp * nexp) option
 
 val destruct_numeric : ?name:string option -> typ -> (kid list * n_constraint * nexp) option