Add top-level pattern match guards internally

Also fix bug in mono analysis with generated variables
Breaks lots of typechecking tests because it generates unnecessary
equality tests on units (and the tests don't have generic equality),
which I'll fix next.

18 files changed, 283 insertions, 148 deletions

diff --git a/src/ast_util.ml b/src/ast_util.ml
index 6480571c..2fd43db5 100644
--- a/src/ast_util.ml
+++ b/src/ast_util.ml
@@ -77,7 +77,7 @@ let mk_lit lit_aux = L_aux (lit_aux, Parse_ast.Unknown)
 
 let mk_lit_exp lit_aux = mk_exp (E_lit (mk_lit lit_aux))
 
-let mk_funcl id pat body = FCL_aux (FCL_Funcl (id, pat, body), no_annot)
+let mk_funcl id pat body = FCL_aux (FCL_Funcl (id, Pat_aux (Pat_exp (pat, body),no_annot)), no_annot)
 
 let mk_qi_nc nc = QI_aux (QI_const nc, Parse_ast.Unknown)
 
@@ -668,7 +668,7 @@ let rec string_of_index_range (BF_aux (ir, _)) =
 
 let id_of_fundef (FD_aux (FD_function (_, _, _, funcls), (l, _))) =
   match (List.fold_right
-           (fun (FCL_aux (FCL_Funcl (id, _, _), _)) id' ->
+           (fun (FCL_aux (FCL_Funcl (id, _), _)) id' ->
              match id' with
              | Some id' -> if string_of_id id' = string_of_id id then Some id'
                            else raise (Reporting_basic.err_typ l
@@ -854,3 +854,13 @@ and undefined_of_typ_args mwords l annot (Typ_arg_aux (typ_arg_aux, _) as typ_ar
   | Typ_arg_nexp n -> [E_aux (E_sizeof n, (l, annot (atom_typ n)))]
   | Typ_arg_typ typ -> [undefined_of_typ mwords l annot typ]
   | Typ_arg_order _ -> []
+
+let destruct_pexp (Pat_aux (pexp,ann)) =
+  match pexp with
+  | Pat_exp (pat,exp) -> pat,None,exp,ann
+  | Pat_when (pat,guard,exp) -> pat,Some guard,exp,ann
+
+let construct_pexp (pat,guard,exp,ann) =
+  match guard with
+  | None -> Pat_aux (Pat_exp (pat,exp),ann)
+  | Some guard -> Pat_aux (Pat_when (pat,guard,exp),ann)
diff --git a/src/ast_util.mli b/src/ast_util.mli
index cbf7deec..a45ca4e9 100644
--- a/src/ast_util.mli
+++ b/src/ast_util.mli
@@ -153,6 +153,7 @@ val quant_split : typquant -> kinded_id list * n_constraint list
 (* Functions to map over the annotations in sub-expressions *)
 val map_exp_annot : ('a annot -> 'b annot) -> 'a exp -> 'b exp
 val map_pat_annot : ('a annot -> 'b annot) -> 'a pat -> 'b pat
+val map_pexp_annot : ('a annot -> 'b annot) -> 'a pexp -> 'b pexp
 val map_lexp_annot : ('a annot -> 'b annot) -> 'a lexp -> 'b lexp
 val map_letbind_annot : ('a annot -> 'b annot) -> 'a letbind -> 'b letbind
 
@@ -271,3 +272,6 @@ val tyvars_of_nexp : nexp -> KidSet.t
 val tyvars_of_typ : typ -> KidSet.t
 
 val undefined_of_typ : bool -> Ast.l -> (typ -> 'annot) -> typ -> 'annot exp
+
+val destruct_pexp : 'a pexp -> 'a pat * ('a exp) option * 'a exp * (Ast.l * 'a)
+val construct_pexp : 'a pat * ('a exp) option * 'a exp * (Ast.l * 'a) ->  'a pexp
diff --git a/src/initial_check.ml b/src/initial_check.ml
index 21e27ac9..6e8af1be 100644
--- a/src/initial_check.ml
+++ b/src/initial_check.ml
@@ -742,7 +742,7 @@ let to_ast_funcl (names,k_env,def_ord) (Parse_ast.FCL_aux(fcl,l) : Parse_ast.fun
   (*let _ = Printf.eprintf "to_ast_funcl\n" in*)
   match fcl with
   | Parse_ast.FCL_Funcl(id,pat,exp) -> 
-    FCL_aux(FCL_Funcl(to_ast_id id, to_ast_pat k_env def_ord pat, to_ast_exp k_env def_ord exp),(l,()))
+    FCL_aux(FCL_Funcl(to_ast_id id, Pat_aux (Pat_exp (to_ast_pat k_env def_ord pat, to_ast_exp k_env def_ord exp),(l,()))),(l,()))
 
 let to_ast_fundef  (names,k_env,def_ord) (Parse_ast.FD_aux(fd,l):Parse_ast.fundef) : (unit fundef) envs_out = 
   match fd with
diff --git a/src/lem_interp/instruction_extractor.lem b/src/lem_interp/instruction_extractor.lem
index b49646ea..11947c17 100644
--- a/src/lem_interp/instruction_extractor.lem
+++ b/src/lem_interp/instruction_extractor.lem
@@ -103,7 +103,8 @@ end
 let rec extract_from_decode decoder = 
   match decoder with
     | [] -> []
-    | (FCL_aux (FCL_Funcl _ pat exp) _)::decoder ->
+    | (FCL_aux (FCL_Funcl _ (Pat_aux pexp _)) _)::decoder ->
+      let exp = match pexp with Pat_exp _ exp -> exp | Pat_when _ _ exp -> exp end in
       (match exp with
 	| E_aux (E_app (Id_aux(Id id) _) parms) (_,(Just (_,Tag_ctor,_,_,_))) ->
 	  Instr_form id (List.map extract_parm parms) []
@@ -112,7 +113,7 @@ end
 
 let rec extract_effects_of_fcl id execute = match execute with
   | [] -> []
-  | FCL_aux (FCL_Funcl _ (P_aux (P_app (Id_aux (Id i) _) _) _) _) (_,(Just(_,_,_,Effect_aux(Effect_set efs) _,_))) :: executes ->
+  | FCL_aux (FCL_Funcl _ (Pat_aux (Pat_exp (P_aux (P_app (Id_aux (Id i) _) _) _) _) _)) (_,(Just(_,_,_,Effect_aux(Effect_set efs) _,_))) :: executes ->
     if i = id 
     then efs
     else extract_effects_of_fcl id executes
@@ -134,7 +135,7 @@ let rec extract_patt_parm (P_aux p (_,tannot)) =
 
 let rec extract_from_execute fcls = match fcls with
   | [] -> []
-  | FCL_aux (FCL_Funcl _ (P_aux (P_app (Id_aux (Id i) _) parms) _) _) (_,Just(_,_,_,Effect_aux(Effect_set efs) _,_))::fcls ->
+  | FCL_aux (FCL_Funcl _ (Pat_aux (Pat_exp (P_aux (P_app (Id_aux (Id i) _) parms) _) _) _)) (_,Just(_,_,_,Effect_aux(Effect_set efs) _,_))::fcls ->
      (Instr_form i (List.map extract_patt_parm parms) efs)::extract_from_execute fcls
   | _ :: fcls ->
      (* AA: Find out what breaks this *)
diff --git a/src/lem_interp/interp.lem b/src/lem_interp/interp.lem
index 486728d9..431c1a08 100644
--- a/src/lem_interp/interp.lem
+++ b/src/lem_interp/interp.lem
@@ -639,7 +639,7 @@ let rec to_fdefs (Defs defs) =
             | FD_aux (FD_function _ _ _ fcls) _ ->
               match fcls with
                 | [] -> to_fdefs (Defs defs)
-                | (FCL_aux (FCL_Funcl name _ _) _)::_  ->
+                | (FCL_aux (FCL_Funcl name _) _)::_  ->
                   Map.insert (get_id name) fcls (to_fdefs (Defs defs)) end end)
         | _ -> to_fdefs (Defs defs) end end
 
@@ -1625,7 +1625,7 @@ val find_funcl : top_level -> list (funcl tannot) -> value -> list (lenv * bool
 let rec find_funcl t_level funcls value =
   match funcls with
   | [] -> []
-  | (FCL_aux (FCL_Funcl id pat exp) _)::funcls ->
+  | (FCL_aux (FCL_Funcl id (Pat_aux (Pat_exp pat exp) _)) _)::funcls ->
     let (is_matching,used_unknown,env) = match_pattern t_level pat value in
     if (is_matching && used_unknown)
     then (env,used_unknown,exp)::(find_funcl t_level funcls value)
diff --git a/src/lem_interp/interp_utilities.lem b/src/lem_interp/interp_utilities.lem
index a178df61..1e6c59ff 100644
--- a/src/lem_interp/interp_utilities.lem
+++ b/src/lem_interp/interp_utilities.lem
@@ -178,7 +178,7 @@ val find_type_def : defs tannot -> id -> maybe (type_def tannot)
 val find_function : defs tannot -> id -> maybe (list (funcl tannot))
 
 let get_funcls id (FD_aux (FD_function _ _ _ fcls) _) =
-  List.filter (fun (FCL_aux (FCL_Funcl name pat exp) _) -> (get_id id) = (get_id name)) fcls
+  List.filter (fun (FCL_aux (FCL_Funcl name pexp) _) -> (get_id id) = (get_id name)) fcls
 
 let rec find_function (Defs defs) id =
   match defs with 
diff --git a/src/lem_interp/pretty_interp.ml b/src/lem_interp/pretty_interp.ml
index 129b74c0..bf42795b 100644
--- a/src/lem_interp/pretty_interp.ml
+++ b/src/lem_interp/pretty_interp.ml
@@ -657,7 +657,7 @@ let doc_effects_opt (Effect_opt_aux(e,_)) = match e with
   | Effect_opt_pure -> string "pure"
   | Effect_opt_effect e -> doc_effects e
 
-let doc_funcl env mem add_red (FCL_aux(FCL_Funcl(id,pat,exp),_)) =
+let doc_funcl env mem add_red (FCL_aux(FCL_Funcl(id,Pat_aux (Pat_exp (pat, exp), _)),_)) =
   group (doc_op equals (separate space [doc_id id; doc_atomic_pat pat]) (doc_exp env mem add_red false exp))
 
 let doc_fundef env mem add_red (FD_aux(FD_function(r, typa, efa, fcls),_)) =
diff --git a/src/monomorphise.ml b/src/monomorphise.ml
index 6350f511..6ae90c23 100644
--- a/src/monomorphise.ml
+++ b/src/monomorphise.ml
@@ -1080,7 +1080,7 @@ let split_defs splits defs =
     else
       let Defs ds = defs in
       match list_extract (function
-      | (DEF_fundef (FD_aux (FD_function (_,_,eff,((FCL_aux (FCL_Funcl (id',_,_),_))::_ as fcls)),_)))
+      | (DEF_fundef (FD_aux (FD_function (_,_,eff,((FCL_aux (FCL_Funcl (id',_),_))::_ as fcls)),_)))
         -> if Id.compare id id' = 0 then Some (eff,fcls) else None
       | _ -> None) ds with
       | None -> None
@@ -1091,7 +1091,7 @@ let split_defs splits defs =
            | [e] -> e
            | _ -> E_aux (E_tuple args,(Generated l,None)) in
          let cases = List.map (function
-           | FCL_aux (FCL_Funcl (_,pat,exp), ann) -> Pat_aux (Pat_exp (pat,exp),ann))
+           | FCL_aux (FCL_Funcl (_,pexp), ann) -> pexp)
            fcls in
          match can_match_with_env env arg cases Bindings.empty Bindings.empty with
          | Some (exp,bindings,kbindings) ->
@@ -1497,20 +1497,8 @@ let split_defs splits defs =
       | LEXP_field (le,id) -> re (LEXP_field (map_lexp le, id))
     in
 
-    let map_funcl (FCL_aux (FCL_Funcl (id,pat,exp),annot)) =
-      match map_pat pat with
-      | NoSplit -> [FCL_aux (FCL_Funcl (id, pat, map_exp exp), annot)]
-      | VarSplit patsubsts ->
-         List.map (fun (pat',substs) ->
-           let exp' = subst_exp substs exp in
-           FCL_aux (FCL_Funcl (id, pat', map_exp exp'), annot))
-           patsubsts
-      | ConstrSplit patnsubsts ->
-         List.map (fun (pat',nsubst) ->
-           let pat' = nexp_subst_pat nsubst pat' in
-           let exp' = nexp_subst_exp nsubst exp in
-           FCL_aux (FCL_Funcl (id, pat', map_exp exp'), annot)
-         ) patnsubsts
+    let map_funcl (FCL_aux (FCL_Funcl (id,pexp),annot)) =
+      List.map (fun pexp -> FCL_aux (FCL_Funcl (id,pexp),annot)) (map_pexp pexp)
     in
 
     let map_fundef (FD_aux (FD_function (r,t,e,fcls),annot)) =
@@ -1673,19 +1661,21 @@ let replace_type env typ =
 
 let rewrite_size_parameters env (Defs defs) =
   let open Rewriter in
-  let size_vars exp =
-    fst (fold_exp
+  let size_vars pexp =
+    fst (fold_pexp
            { (compute_exp_alg KidSet.empty KidSet.union) with
              e_aux = (fun ((s,e),annot) -> KidSet.union s (sizes_of_annot annot), E_aux (e,annot));
              e_let = (fun ((sl,lb),(s2,e2)) -> KidSet.union sl (KidSet.diff s2 (tyvars_bound_in_lb lb)), E_let (lb,e2));
              pat_exp = (fun ((sp,pat),(s,e)) -> KidSet.diff s (tyvars_bound_in_pat pat), Pat_exp (pat,e))}
-           exp)
+           pexp)
   in
-  let sizes_funcl fsizes (FCL_aux (FCL_Funcl (id,pat,exp),(l,_))) =
-    let sizes = size_vars exp in
+  let sizes_funcl fsizes (FCL_aux (FCL_Funcl (id,pexp),(l,_))) =
+    let sizes = size_vars pexp in
+    let pat,guard,exp,pannot = destruct_pexp pexp in
     (* TODO: what, if anything, should sequential be? *)
     let visible_tyvars =
-      KidSet.union (Pretty_print_lem.lem_tyvars_of_typ false true (pat_typ_of pat))
+      KidSet.union
+        (Pretty_print_lem.lem_tyvars_of_typ false true (pat_typ_of pat))
          (Pretty_print_lem.lem_tyvars_of_typ false true (typ_of exp))
     in
     let expose_tyvars = KidSet.diff sizes visible_tyvars in
@@ -1738,11 +1728,12 @@ let rewrite_size_parameters env (Defs defs) =
        E_app (id,args')
     | exception Not_found -> E_app (id,args)
   in
-  let rewrite_funcl (FCL_aux (FCL_Funcl (id,pat,body),(l,annot))) =
-    let pat,body =
+  let rewrite_funcl (FCL_aux (FCL_Funcl (id,pexp),(l,annot))) =
+    let pat,guard,body,(pl,_) = destruct_pexp pexp in
+    let pat,guard,body =
       (* Update pattern and add itself -> nat wrapper to body *)
       match Bindings.find id fn_sizes with
-      | [] -> pat,body
+      | [] -> pat,guard,body
       | to_change ->
          let pat, vars =
            match pat with
@@ -1760,13 +1751,21 @@ let rewrite_size_parameters env (Defs defs) =
                             "Expected multiple parameters at single parameter")
               end
          in
+         (* TODO: only add bindings that are necessary (esp for guards) *)
          let body = List.fold_left add_var_rebind body vars in
-         pat,body
-      | exception Not_found -> pat,body
+         let guard = match guard with
+           | None -> None
+           | Some exp -> Some (List.fold_left add_var_rebind body vars)
+         in
+         pat,guard,body
+      | exception Not_found -> pat,guard,body
     in
     (* Update function applications *)
     let body = fold_exp { id_exp_alg with e_app = rewrite_e_app } body in
-    FCL_aux (FCL_Funcl (id,pat,body),(l,None))
+    let guard = match guard with
+      | None -> None
+      | Some exp -> Some (fold_exp { id_exp_alg with e_app = rewrite_e_app } exp) in
+    FCL_aux (FCL_Funcl (id,construct_pexp (pat,guard,body,(pl,None))),(l,None))
   in
   let rewrite_def = function
     | DEF_fundef (FD_aux (FD_function (recopt,tannopt,effopt,funcls),(l,_))) ->
@@ -2250,10 +2249,9 @@ and analyse_lexp fn_id env assigns deps (LEXP_aux (lexp,_)) =
 let translate_id (Id_aux (_,l) as id) =
   let rec aux l =
     match l with
-    | Range (pos,_) -> id,(pos.Lexing.pos_fname,pos.Lexing.pos_lnum)
+    | Range (pos,_) -> Some (id,(pos.Lexing.pos_fname,pos.Lexing.pos_lnum))
     | Generated l -> aux l
-    | _ ->
-       raise (Reporting_basic.err_general l ("Unable to give location for " ^ string_of_id id))
+    | _ -> None
   in aux l
 
 let initial_env fn_id (TypQ_aux (tq,_)) pat =
@@ -2276,14 +2274,22 @@ let initial_env fn_id (TypQ_aux (tq,_)) pat =
       | P_wild
         -> ArgSet.empty,Bindings.empty,KBindings.empty
       | P_as (pat,id) ->
-         let s,v,k = aux pat in
-         let id' = translate_id id in
-         ArgSet.add id' s, Bindings.add id (Have (ArgSet.singleton id',CallerArgSet.empty,CallerKidSet.empty)) v,k
+         begin
+           let s,v,k = aux pat in
+           match translate_id id with
+           | Some id' -> ArgSet.add id' s, Bindings.add id (Have (ArgSet.singleton id',CallerArgSet.empty,CallerKidSet.empty)) v,k
+           | None -> s, Bindings.add id (Unknown (l, ("Unable to give location for " ^ string_of_id id))) v, k
+         end
       | P_typ (_,pat) -> aux pat
       | P_id id ->
-         let id' = translate_id id in
-         let s = ArgSet.singleton id' in
-         s, Bindings.singleton id (Have (s,CallerArgSet.empty,CallerKidSet.empty)), KBindings.empty
+         begin
+         match translate_id id with
+         | Some id' ->
+            let s = ArgSet.singleton id' in
+            s, Bindings.singleton id (Have (s,CallerArgSet.empty,CallerKidSet.empty)), KBindings.empty
+         | None ->
+            ArgSet.empty, Bindings.singleton id (Unknown (l, ("Unable to give location for " ^ string_of_id id))), KBindings.empty
+         end
       | P_var (pat,kid) ->
          let s,v,k = aux pat in
          s,v,KBindings.add kid (Have (ArgSet.empty,CallerArgSet.singleton (fn_id,i),CallerKidSet.empty)) k
@@ -2330,10 +2336,16 @@ let print_result r =
   let _ = print_endline ("  kid_in_caller: " ^ string_of_callerkidset r.kid_in_caller) in
   ()
 
-let analyse_funcl debug tenv (FCL_aux (FCL_Funcl (id,pat,body),_)) =
+let analyse_funcl debug tenv (FCL_aux (FCL_Funcl (id,pexp),_)) =
+  let pat,guard,body,_ = destruct_pexp pexp in
   let (tq,_) = Env.get_val_spec id tenv in
   let aenv = initial_env id tq pat in
   let _,_,r = analyse_exp id aenv Bindings.empty body in
+  let r = match guard with
+    | None -> r
+    | Some exp -> let _,_,r' = analyse_exp id aenv Bindings.empty exp in
+                  merge r r'
+  in
   let _ =
     if debug > 2 then
       (print_endline (string_of_id id);
diff --git a/src/ocaml_backend.ml b/src/ocaml_backend.ml
index 6254d580..e56060da 100644
--- a/src/ocaml_backend.ml
+++ b/src/ocaml_backend.ml
@@ -365,7 +365,7 @@ and ocaml_atomic_lexp ctx (LEXP_aux (lexp_aux, _) as lexp) =
   | _ -> parens (ocaml_lexp ctx lexp)
 
 let rec get_initialize_registers = function
-  | DEF_fundef (FD_aux (FD_function (_, _, _, [FCL_aux (FCL_Funcl (id, _, E_aux (E_block inits, _)), _)]), _)) :: defs
+  | DEF_fundef (FD_aux (FD_function (_, _, _, [FCL_aux (FCL_Funcl (id, Pat_aux (Pat_exp (_, E_aux (E_block inits, _)),_)), _)]), _)) :: defs
        when Id.compare id (mk_id "initialize_registers") = 0 ->
      inits
   | _ :: defs -> get_initialize_registers defs
@@ -397,11 +397,10 @@ let function_header () =
 
 let funcls_id = function
   | [] -> failwith "Ocaml: empty function"
-  | FCL_aux (FCL_Funcl (id, pat, exp),_) :: _ -> id
+  | FCL_aux (FCL_Funcl (id, _),_) :: _ -> id
 
-let ocaml_funcl_match ctx (FCL_aux (FCL_Funcl (id, pat, exp), _)) =
-  separate space [bar; ocaml_pat ctx pat; string "->"]
-  ^//^ ocaml_exp ctx exp
+let ocaml_funcl_match ctx (FCL_aux (FCL_Funcl (id, pexp), _)) =
+  ocaml_pexp ctx pexp
 
 let rec ocaml_funcl_matches ctx = function
   | [] -> failwith "Ocaml: empty function"
@@ -435,9 +434,14 @@ let ocaml_funcls ctx =
   in
   function
   | [] -> failwith "Ocaml: empty function"
-  | [FCL_aux (FCL_Funcl (id, pat, exp),_)] ->
+  | [FCL_aux (FCL_Funcl (id, pexp),_)] ->
      let Typ_aux (Typ_fn (typ1, typ2, _), _) = Bindings.find id ctx.val_specs in
      let pat_sym = gensym () in
+     let pat, exp =
+       match pexp with
+       | Pat_aux (Pat_exp (pat,exp),_) -> pat,exp
+       | Pat_aux (Pat_when (pat,wh,exp),_) -> failwith "OCaml: top-level pattern guards not supported"
+     in
      let annot_pat =
        let pat = parens (ocaml_pat ctx pat ^^ space ^^ colon ^^ space ^^ ocaml_typ ctx typ1) in
        if !opt_trace_ocaml
diff --git a/src/pretty_print_lem.ml b/src/pretty_print_lem.ml
index 7ced47ee..c76be843 100644
--- a/src/pretty_print_lem.ml
+++ b/src/pretty_print_lem.ml
@@ -1213,20 +1213,34 @@ let doc_fun_body_lem sequential mwords exp =
   then align (string "catch_early_return" ^//^ parens (doc_exp))
   else doc_exp
 
-let doc_funcl_lem sequential mwords (FCL_aux(FCL_Funcl(id,pat,exp),_)) =
+let doc_funcl_lem sequential mwords (FCL_aux(FCL_Funcl(id,pexp),_)) =
+  let pat,guard,exp,(l,_) = destruct_pexp pexp in
+  let _ = match guard with
+    | None -> ()
+    | _ ->
+       raise (Reporting_basic.err_unreachable l
+                "guarded pattern expression should have been rewritten before pretty-printing")
+  in
   group (prefix 3 1 ((doc_pat_lem sequential mwords false pat) ^^ space ^^ arrow)
     (doc_fun_body_lem sequential mwords exp))
 
 let get_id = function
   | [] -> failwith "FD_function with empty list"
-  | (FCL_aux (FCL_Funcl (id,_,_),_))::_ -> id
+  | (FCL_aux (FCL_Funcl (id,_),_))::_ -> id
 
 module StringSet = Set.Make(String)
 
 let doc_fundef_rhs_lem sequential mwords (FD_aux(FD_function(r, typa, efa, funcls),fannot) as fd) =
   match funcls with
   | [] -> failwith "FD_function with empty function list"
-  | [FCL_aux(FCL_Funcl(id,pat,exp),_)] ->
+  | [FCL_aux(FCL_Funcl(id,pexp),_)] ->
+     let pat,guard,exp,(l,_) = destruct_pexp pexp in
+     let _ = match guard with
+       | None -> ()
+       | _ ->
+          raise (Reporting_basic.err_unreachable l
+                   "guarded pattern expression should have been rewritten before pretty-printing")
+     in
     (prefix 2 1)
       ((separate space)
          [doc_id_lem id;
@@ -1251,13 +1265,20 @@ let doc_mutrec_lem sequential mwords = function
 let rec doc_fundef_lem sequential mwords (FD_aux(FD_function(r, typa, efa, fcls),fannot) as fd) =
   match fcls with
   | [] -> failwith "FD_function with empty function list"
-  | FCL_aux (FCL_Funcl(id,_,exp),_) :: _
+  | FCL_aux (FCL_Funcl(id,_),_) :: _
     when string_of_id id = "execute" (*|| string_of_id id = "initial_analysis"*) ->
     let (_,auxiliary_functions,clauses) =
       List.fold_left
         (fun (already_used_fnames,auxiliary_functions,clauses) funcl ->
           match funcl with
-          | FCL_aux (FCL_Funcl (Id_aux (Id _,l),pat,exp),annot) ->
+          | FCL_aux (FCL_Funcl (Id_aux (Id _,l),pexp),annot) ->
+             let pat,guard,exp,(pexp_l,_) = destruct_pexp pexp in
+             let _ = match guard with
+               | None -> ()
+               | _ ->
+                  raise (Reporting_basic.err_unreachable pexp_l
+                           "guarded pattern expression should have been rewritten before pretty-printing")
+             in
              let ctor, l, argspat, pannot = (match pat with
                | P_aux (P_app (Id_aux (Id ctor,l),argspat),pannot) ->
                  (ctor, l, argspat, pannot)
@@ -1273,7 +1294,8 @@ let rec doc_fundef_lem sequential mwords (FD_aux(FD_function(r, typa, efa, fcls)
              let aux_fname = pick_name_not_clashing_with already_used_fnames (string_of_id id ^ "_" ^ ctor) in
              let already_used_fnames = StringSet.add aux_fname already_used_fnames in
              let fcl = FCL_aux (FCL_Funcl (Id_aux (Id aux_fname,l),
-                                           P_aux (P_tup argspat,pannot),exp),annot) in
+                                           Pat_aux (Pat_exp (P_aux (P_tup argspat,pannot),exp),(pexp_l,None)))
+                                  ,annot) in
              let auxiliary_functions =
                 auxiliary_functions ^^ hardline ^^ hardline ^^
                   doc_fundef_lem sequential mwords (FD_aux (FD_function(r,typa,efa,[fcl]),fannot)) in
@@ -1304,8 +1326,8 @@ let rec doc_fundef_lem sequential mwords (FD_aux(FD_function(r, typa, efa, fcls)
     (prefix 2 1)
       ((separate space) [string "let" ^^ doc_rec_lem r ^^ doc_id_lem id;equals;string "function"])
       (clauses ^/^ string "end")
-  | FCL_aux (FCL_Funcl(id,_,exp),_) :: _
-    when not (Env.is_extern id (env_of exp) "lem") ->
+  | FCL_aux (FCL_Funcl(id,_),annot) :: _
+    when not (Env.is_extern id (env_of_annot annot) "lem") ->
     string "let" ^^ (doc_rec_lem r) ^^ (doc_fundef_rhs_lem sequential mwords fd)
   | _ -> empty
 
diff --git a/src/pretty_print_lem_ast.ml b/src/pretty_print_lem_ast.ml
index 020b8dbd..3b7a7345 100644
--- a/src/pretty_print_lem_ast.ml
+++ b/src/pretty_print_lem_ast.ml
@@ -579,9 +579,9 @@ let pp_lem_effects_opt ppf (Effect_opt_aux(e,l)) =
   | Effect_opt_pure -> fprintf ppf "(Effect_opt_aux Effect_opt_pure %a)" pp_lem_l l
   | Effect_opt_effect e -> fprintf ppf "(Effect_opt_aux (Effect_opt_effect %a) %a)" pp_lem_effects e pp_lem_l l
 
-let pp_lem_funcl ppf (FCL_aux(FCL_Funcl(id,pat,exp),(l,annot))) =
-  fprintf ppf "@[<0>(FCL_aux (%a %a %a %a) (%a,%a))@]@\n"
-    kwd "FCL_Funcl" pp_lem_id id pp_lem_pat pat pp_lem_exp exp pp_lem_l l pp_annot annot
+let pp_lem_funcl ppf (FCL_aux(FCL_Funcl(id,pexp),(l,annot))) =
+  fprintf ppf "@[<0>(FCL_aux (%a %a %a) (%a,%a))@]@\n"
+    kwd "FCL_Funcl" pp_lem_id id pp_lem_case pexp pp_lem_l l pp_annot annot
 
 let pp_lem_fundef ppf (FD_aux(FD_function(r, typa, efa, fcls),(l,annot))) =
   let pp_funcls ppf funcl = fprintf ppf "%a %a" pp_lem_funcl funcl kwd ";" in
diff --git a/src/pretty_print_sail.ml b/src/pretty_print_sail.ml
index 4ef506c5..05dbb9ee 100644
--- a/src/pretty_print_sail.ml
+++ b/src/pretty_print_sail.ml
@@ -485,8 +485,13 @@ let doc_effects_opt (Effect_opt_aux(e,_)) = match e with
   | Effect_opt_pure -> string "pure"
   | Effect_opt_effect e -> doc_effects e
 
-let doc_funcl (FCL_aux(FCL_Funcl(id,pat,exp),_)) =
-  group (doc_op equals (separate space [doc_id id; doc_atomic_pat pat]) (doc_exp exp))
+let doc_funcl (FCL_aux(FCL_Funcl(id,pexp),_)) =
+  match pexp with
+  | Pat_aux (Pat_exp (pat,exp),_) ->
+     group (doc_op equals (separate space [doc_id id; doc_atomic_pat pat]) (doc_exp exp))
+  | Pat_aux (Pat_when (pat,wh,exp),_) ->
+     group (doc_op equals (separate space [doc_id id; doc_atomic_pat pat; string "if"; doc_exp wh])
+              (doc_exp exp))
 
 let doc_fundef (FD_aux(FD_function(r, typa, efa, fcls),_)) =
   match fcls with
diff --git a/src/pretty_print_sail2.ml b/src/pretty_print_sail2.ml
index 07259a6f..4d05e04d 100644
--- a/src/pretty_print_sail2.ml
+++ b/src/pretty_print_sail2.ml
@@ -394,12 +394,12 @@ and doc_atomic_lexp (LEXP_aux (l_aux, _) as lexp) =
   | LEXP_vector_range (lexp, exp1, exp2) -> doc_atomic_lexp lexp ^^ brackets (separate space [doc_exp exp1; string ".."; doc_exp exp2])
   | LEXP_memory (id, exps) -> doc_id id ^^ parens (separate_map (comma ^^ space) doc_exp exps)
   | _ -> parens (doc_lexp lexp)
-and doc_pexps pexps = surround 2 0 lbrace (separate_map (comma ^^ hardline) doc_pexp pexps) rbrace
-and doc_pexp (Pat_aux (pat_aux, _)) =
+and doc_pexps pexps = surround 2 0 lbrace (separate_map (comma ^^ hardline) (doc_pexp "=>") pexps) rbrace
+and doc_pexp sym (Pat_aux (pat_aux, _)) =
   match pat_aux with
-  | Pat_exp (pat, exp) -> separate space [doc_pat pat; string "=>"; doc_exp exp]
+  | Pat_exp (pat, exp) -> separate space [doc_pat pat; string sym; doc_exp exp]
   | Pat_when (pat, wh, exp) ->
-     separate space [doc_pat pat; string "if"; doc_exp wh; string "=>"; doc_exp exp]
+     separate space [doc_pat pat; string "if"; doc_exp wh; string sym; doc_exp exp]
 and doc_letbind (LB_aux (lb_aux, _)) =
   match lb_aux with
   | LB_val (pat, exp) ->
@@ -407,8 +407,8 @@ and doc_letbind (LB_aux (lb_aux, _)) =
 
 let doc_funcl funcl = string "FUNCL"
 
-let doc_funcl (FCL_aux  (FCL_Funcl (id, pat, exp), _)) =
-  group (separate space [doc_id id; doc_pat pat; equals; doc_exp exp])
+let doc_funcl (FCL_aux  (FCL_Funcl (id, pexp), _)) =
+  group (separate space [doc_id id; doc_pexp "=" pexp])
 
 let doc_default (DT_aux(df,_)) = match df with
   | DT_kind(bk,v) -> string "DT_kind" (* separate space [string "default"; doc_bkind bk; doc_var v] *)
diff --git a/src/rewriter.ml b/src/rewriter.ml
index 4d5e7967..abe351d4 100644
--- a/src/rewriter.ml
+++ b/src/rewriter.ml
@@ -257,6 +257,14 @@ let vector_string_to_bit_list l lit =
                    | '1' -> L_aux (L_one, gen_loc l)
                    | _ -> raise (Reporting_basic.err_unreachable (gen_loc l) "binary had non-zero or one")) s_bin
 
+let rewrite_pexp rewriters =
+  let rewrite = rewriters.rewrite_exp rewriters in
+  function
+  | (Pat_aux (Pat_exp(p, e), pannot)) ->
+     Pat_aux (Pat_exp(rewriters.rewrite_pat rewriters p, rewrite e), pannot)
+  | (Pat_aux (Pat_when(p, e, e'), pannot)) ->
+     Pat_aux (Pat_when(rewriters.rewrite_pat rewriters p, rewrite e, rewrite e'), pannot)
+
 let rewrite_pat rewriters (P_aux (pat,(l,annot))) =
   let rewrap p = P_aux (p,(l,annot)) in
   let rewrite = rewriters.rewrite_pat rewriters in
@@ -322,12 +330,7 @@ let rewrite_exp rewriters (E_aux (exp,(l,annot))) =
                                               FE_aux(FE_Fexp(id,rewrite e),fannot)) fexps, bool), fannot))))
   | E_field(exp,id) -> rewrap (E_field(rewrite exp,id))
   | E_case (exp,pexps) ->
-    let rewrite_pexp = function
-    | (Pat_aux (Pat_exp(p, e), pannot)) ->
-      Pat_aux (Pat_exp(rewriters.rewrite_pat rewriters p, rewrite e), pannot)
-    | (Pat_aux (Pat_when(p, e, e'), pannot)) ->
-      Pat_aux (Pat_when(rewriters.rewrite_pat rewriters p, rewrite e, rewrite e'), pannot) in
-    rewrap (E_case (rewrite exp, List.map rewrite_pexp pexps))
+    rewrap (E_case (rewrite exp, List.map (rewrite_pexp rewriters) pexps))
   | E_let (letbind,body) -> rewrap (E_let(rewriters.rewrite_let rewriters letbind,rewrite body))
   | E_assign (lexp,exp) -> rewrap (E_assign(rewriters.rewrite_lexp rewriters lexp,rewrite exp))
   | E_sizeof n -> rewrap (E_sizeof n)
@@ -364,9 +367,8 @@ let rewrite_lexp rewriters (LEXP_aux(lexp,(l,annot))) =
   | LEXP_field (lexp,id) -> rewrap (LEXP_field (rewriters.rewrite_lexp rewriters lexp,id))
 
 let rewrite_fun rewriters (FD_aux (FD_function(recopt,tannotopt,effectopt,funcls),(l,fdannot))) = 
-  let rewrite_funcl (FCL_aux (FCL_Funcl(id,pat,exp),(l,annot))) =
-    (FCL_aux (FCL_Funcl (id,rewriters.rewrite_pat rewriters pat,
-                         rewriters.rewrite_exp rewriters exp),(l,annot))) 
+  let rewrite_funcl (FCL_aux (FCL_Funcl(id,pexp),(l,annot))) =
+    (FCL_aux (FCL_Funcl (id,rewrite_pexp rewriters pexp),(l,annot)))
   in FD_aux (FD_function(recopt,tannotopt,effectopt,List.map rewrite_funcl funcls),(l,fdannot))
 
 let rewrite_def rewriters d = match d with
diff --git a/src/rewriter.mli b/src/rewriter.mli
index cd293ca5..6f8c6396 100644
--- a/src/rewriter.mli
+++ b/src/rewriter.mli
@@ -176,6 +176,10 @@ val fold_exp : ('a,'exp,'exp_aux,'lexp,'lexp_aux,'fexp,'fexp_aux,'fexps,'fexps_a
       'opt_default_aux,'opt_default,'pexp,'pexp_aux,'letbind_aux,'letbind,
       'pat,'pat_aux,'fpat,'fpat_aux) exp_alg -> 'a exp -> 'exp
 
+val fold_pexp : ('a,'exp,'exp_aux,'lexp,'lexp_aux,'fexp,'fexp_aux,'fexps,'fexps_aux,
+      'opt_default_aux,'opt_default,'pexp,'pexp_aux,'letbind_aux,'letbind,
+      'pat,'pat_aux,'fpat,'fpat_aux) exp_alg -> 'a pexp -> 'pexp
+
 val id_pat_alg : ('a,'a pat, 'a pat_aux, 'a fpat, 'a fpat_aux) pat_alg
 val id_exp_alg :
   ('a,'a exp,'a exp_aux,'a lexp,'a lexp_aux,'a fexp,
diff --git a/src/rewrites.ml b/src/rewrites.ml
index 333f87cb..de280529 100644
--- a/src/rewrites.ml
+++ b/src/rewrites.ml
@@ -391,15 +391,21 @@ let rewrite_sizeof (Defs defs) =
 
   let rewrite_sizeof_fun params_map
                          (FD_aux (FD_function (rec_opt,tannot,eff,funcls),((l,_) as annot))) =
-    let rewrite_funcl_body (FCL_aux (FCL_Funcl (id,pat,exp), annot)) (funcls,nvars) =
-      let body_env = env_of exp in
-      let body_typ = typ_of exp in
+    let rewrite_funcl_body (FCL_aux (FCL_Funcl (id,pexp), annot)) (funcls,nvars) =
+      let pat,guard,exp,pannot = destruct_pexp pexp in
       let nmap = nexps_from_params pat in
       (* first rewrite calls to other functions... *)
       let exp' = fst (fold_exp { copy_exp_alg with e_aux = e_app_aux params_map } exp) in
       (* ... then rewrite sizeof expressions in current function body *)
       let exp'' = fold_exp { id_exp_alg with e_sizeof = e_sizeof nmap } exp' in
-      (FCL_aux (FCL_Funcl (id,pat,exp''), annot) :: funcls,
+      let guard' = match guard with
+        | Some guard ->
+           (* As above *)
+           let guard' = fst (fold_exp { copy_exp_alg with e_aux = e_app_aux params_map } guard) in
+           Some (fold_exp { id_exp_alg with e_sizeof = e_sizeof nmap } guard')
+        | None -> None in
+      let pexp' = construct_pexp (pat,guard',exp'',pannot) in
+      (FCL_aux (FCL_Funcl (id,pexp'), annot) :: funcls,
        KidSet.union nvars (sizeof_frees exp'')) in
     let (funcls, nvars) = List.fold_right rewrite_funcl_body funcls ([], KidSet.empty) in
     (* Add a parameter for each remaining free type-level variable in a
@@ -414,7 +420,7 @@ let rewrite_sizeof (Defs defs) =
     let kid_typs = List.map kid_typ (KidSet.elements nvars) in
     let kid_pats = List.map kid_pat (KidSet.elements nvars) in
     let kid_nmap = List.map (fun kid -> (nvar kid, kid_eaux kid)) (KidSet.elements nvars) in
-    let rewrite_funcl_params (FCL_aux (FCL_Funcl (id, pat, exp), annot) as funcl) =
+    let rewrite_funcl_params (FCL_aux (FCL_Funcl (id, pexp), annot) as funcl) =
       let rec rewrite_pat (P_aux (pat, ((l, _) as pannot)) as paux) =
         let penv = env_of_annot pannot in
         let peff = effect_of_annot (snd pannot) in
@@ -440,8 +446,14 @@ let rewrite_sizeof (Defs defs) =
           | ptyp ->
              let ptyp' = Typ_aux (Typ_tup (kid_typs @ [ptyp]), l) in
              P_aux (P_tup (kid_pats @ [paux]), (l, Some (penv, ptyp', peff))) in
+      let pat,guard,exp,pannot = destruct_pexp pexp in
+      let pat' = rewrite_pat pat in
+      let guard' = match guard with
+        | Some guard -> Some (fold_exp { id_exp_alg with e_sizeof = e_sizeof kid_nmap } guard)
+        | None -> None in
       let exp' = fold_exp { id_exp_alg with e_sizeof = e_sizeof kid_nmap } exp in
-      FCL_aux (FCL_Funcl (id, rewrite_pat pat, exp'), annot) in
+      let pexp' = construct_pexp (pat',guard',exp',pannot) in
+      FCL_aux (FCL_Funcl (id, pexp'), annot) in
     let funcls = List.map rewrite_funcl_params funcls in
     let fd = FD_aux (FD_function (rec_opt,tannot,eff,funcls),annot) in
     let params_map =
@@ -812,10 +824,15 @@ let rewrite_exp_remove_vector_concat_pat rewriters (E_aux (exp,(l,annot)) as ful
 
 let rewrite_fun_remove_vector_concat_pat
       rewriters (FD_aux (FD_function(recopt,tannotopt,effectopt,funcls),(l,fdannot))) = 
-  let rewrite_funcl (FCL_aux (FCL_Funcl(id,pat,exp),(l,annot))) =
+  let rewrite_funcl (FCL_aux (FCL_Funcl(id,pexp),(l,annot))) =
+    let pat,guard,exp,pannot = destruct_pexp pexp in
     let (pat',_,decls) = remove_vector_concat_pat pat in
+    let guard' = match guard with
+      | Some exp -> Some (decls (rewriters.rewrite_exp rewriters exp))
+      | None -> None in
     let exp' = decls (rewriters.rewrite_exp rewriters exp) in
-    (FCL_aux (FCL_Funcl (id,pat',exp'),(l,annot))) 
+    let pexp' = construct_pexp (pat',guard',exp',pannot) in
+    (FCL_aux (FCL_Funcl (id,pexp'),(l,annot)))
   in FD_aux (FD_function(recopt,tannotopt,effectopt,List.map rewrite_funcl funcls),(l,fdannot))
 
 let rewrite_defs_remove_vector_concat (Defs defs) =
@@ -1238,13 +1255,18 @@ let rewrite_fun_remove_bitvector_pat
       rewriters (FD_aux (FD_function(recopt,tannotopt,effectopt,funcls),(l,fdannot))) =
   let _ = reset_fresh_name_counter () in
   let funcls = match funcls with
-    | (FCL_aux (FCL_Funcl(id,_,_),_) :: _) ->
-        let clause (FCL_aux (FCL_Funcl(_,pat,exp),annot)) =
+    | (FCL_aux (FCL_Funcl(id,_),_) :: _) ->
+        let clause (FCL_aux (FCL_Funcl(_,pexp),annot)) =
+          let pat,fguard,exp,pannot = destruct_pexp pexp in
           let (pat,(guard,decls,_)) = remove_bitvector_pat pat in
+          let guard = match guard,fguard with
+            | None,e | e,None -> e
+            | Some g, Some wh ->
+               Some (bitwise_and_exp g (decls (rewriters.rewrite_exp rewriters wh)))
+          in
           let exp = decls (rewriters.rewrite_exp rewriters exp) in
-          (pat,guard,exp,annot) in
-        let cs = rewrite_guarded_clauses l (List.map clause funcls) in
-        List.map (fun (pat,exp,annot) -> FCL_aux (FCL_Funcl(id,pat,exp),annot)) cs
+          FCL_aux (FCL_Funcl (id,construct_pexp (pat,guard,exp,annot)),annot) in
+        List.map clause funcls
     | _ -> funcls in
   FD_aux (FD_function(recopt,tannotopt,effectopt,funcls),(l,fdannot))
 
@@ -1272,7 +1294,7 @@ let rewrite_defs_remove_bitvector_pats (Defs defs) =
 
 
 (* Remove pattern guards by rewriting them to if-expressions within the
-   pattern expression. Shares code with the rewriting of bitvector patterns. *)
+   pattern expression. *)
 let rewrite_exp_guarded_pats rewriters (E_aux (exp,(l,annot)) as full_exp) =
   let rewrap e = E_aux (e,(l,annot)) in
   let rewrite_rec = rewriters.rewrite_exp rewriters in
@@ -1300,8 +1322,22 @@ let rewrite_exp_guarded_pats rewriters (E_aux (exp,(l,annot)) as full_exp) =
     else case_exp e (typ_of full_exp) clauses
   | _ -> rewrite_base full_exp
 
+let rewrite_fun_guarded_pats rewriters (FD_aux (FD_function (r,t,e,funcls),(l,fdannot))) =
+   let funcls = match funcls with
+    | (FCL_aux (FCL_Funcl(id,_),_) :: _) ->
+       let clause (FCL_aux (FCL_Funcl(_,pexp),annot)) =
+         let pat,guard,exp,_ = destruct_pexp pexp in
+         let exp = rewriters.rewrite_exp rewriters exp in
+         (pat,guard,exp,annot) in
+       let cs = rewrite_guarded_clauses l (List.map clause funcls) in
+       List.map (fun (pat,exp,annot) ->
+         FCL_aux (FCL_Funcl(id,construct_pexp (pat,None,exp,(Unknown,None))),annot)) cs
+     | _ -> funcls (* TODO is the empty list possible here? *) in
+   FD_aux (FD_function(r,t,e,funcls),(l,fdannot))
+
 let rewrite_defs_guarded_pats =
-  rewrite_defs_base { rewriters_base with rewrite_exp = rewrite_exp_guarded_pats }
+  rewrite_defs_base { rewriters_base with rewrite_exp = rewrite_exp_guarded_pats;
+                    rewrite_fun = rewrite_fun_guarded_pats }
 
 
 let id_is_local_var id env = match Env.lookup_id id env with
@@ -1652,7 +1688,8 @@ let rewrite_defs_early_return (Defs defs) =
       E_aux (E_app (mk_id "early_return", [exp']), (l, annot'))
     | _ -> full_exp in
 
-  let rewrite_funcl_early_return _ (FCL_aux (FCL_Funcl (id, pat, exp), a)) =
+  let rewrite_funcl_early_return _ (FCL_aux (FCL_Funcl (id, pexp), a)) =
+    let pat,guard,exp,pannot = destruct_pexp pexp in
     let exp =
       exp
       (* Pull early returns out as far as possible *)
@@ -1665,7 +1702,7 @@ let rewrite_defs_early_return (Defs defs) =
     | (l, Some (env, typ, eff)) ->
       (l, Some (env, typ, union_effects eff (effect_of exp)))
     | _ -> a in
-    FCL_aux (FCL_Funcl (id, pat, exp), a) in
+    FCL_aux (FCL_Funcl (id, construct_pexp (pat, guard, exp, pannot)), a) in
 
   let rewrite_fun_early_return rewriters
     (FD_aux (FD_function (rec_opt, tannot_opt, effect_opt, funcls), a)) =
@@ -1716,7 +1753,9 @@ let rewrite_fix_val_specs (Defs defs) =
 
   let rewrite_exp val_specs = fold_exp { id_exp_alg with e_aux = e_aux val_specs } in
 
-  let rewrite_funcl (val_specs, funcls) (FCL_aux (FCL_Funcl (id, pat, exp), (l, annot))) =
+  let rewrite_funcl (val_specs, funcls) (FCL_aux (FCL_Funcl (id, pexp), (l, annot))) =
+    let pat,guard,exp,pannot = destruct_pexp pexp in
+    (* Assumes there are no effects in guards *)
     let exp = propagate_exp_effect (rewrite_exp val_specs exp) in
     let vs, eff = match find_vs (env_of_annot (l, annot)) val_specs id with
       | (tq, Typ_aux (Typ_fn (args_t, ret_t, eff), a)) ->
@@ -1727,14 +1766,17 @@ let rewrite_fix_val_specs (Defs defs) =
     in
     let annot = add_effect_annot annot eff in
     (Bindings.add id vs val_specs,
-     funcls @ [FCL_aux (FCL_Funcl (id, pat, exp), (l, annot))])
+     funcls @ [FCL_aux (FCL_Funcl (id, construct_pexp (pat, guard, exp, pannot)), (l, annot))])
   in
 
   let rewrite_fundef (val_specs, FD_aux (FD_function (recopt, tannotopt, effopt, funcls), a)) =
     let (val_specs, funcls) = List.fold_left rewrite_funcl (val_specs, []) funcls in
     (* Repeat once to cross-propagate effects between clauses *)
     let (val_specs, funcls) = List.fold_left rewrite_funcl (val_specs, []) funcls in
-    let is_funcl_rec (FCL_aux (FCL_Funcl (id, _, exp), _)) =
+    let is_funcl_rec (FCL_aux (FCL_Funcl (id, pexp), _)) =
+      let pat,guard,exp,pannot = destruct_pexp pexp in
+      let exp = match guard with None -> exp
+        | Some exp' -> E_aux (E_block [exp';exp],(Parse_ast.Unknown,None)) in
       fst (fold_exp
         { (compute_exp_alg false (||) ) with
           e_app = (fun (f, es) ->
@@ -2140,7 +2182,7 @@ let rewrite_defs_letbind_effects  =
   and n_fexpL (fexps : 'a fexp list) (k : 'a fexp list -> 'a exp) : 'a exp =
     mapCont n_fexp fexps k
 
-  and n_pexp (newreturn : bool) (pexp : 'a pexp) (k : 'a pexp -> 'a exp) : 'a exp =
+  and n_pexp : 'b. bool -> 'a pexp -> ('a pexp -> 'b) -> 'b = fun newreturn pexp k ->
     match pexp with
     | Pat_aux (Pat_exp (pat,exp),annot) ->
       k (fix_eff_pexp (Pat_aux (Pat_exp (pat,n_exp_term newreturn exp), annot)))
@@ -2341,11 +2383,11 @@ let rewrite_defs_letbind_effects  =
     | E_internal_plet _ -> failwith "E_internal_plet should not be here yet" in
 
   let rewrite_fun _ (FD_aux (FD_function(recopt,tannotopt,effectopt,funcls),fdannot)) = 
-    let effectful_funcl (FCL_aux (FCL_Funcl(_, _, exp), _)) = effectful exp in
+    let effectful_funcl (FCL_aux (FCL_Funcl(_, pexp), _)) = effectful_pexp pexp in
     let newreturn = List.exists effectful_funcl funcls in
-    let rewrite_funcl (FCL_aux (FCL_Funcl(id,pat,exp),annot)) =
+    let rewrite_funcl (FCL_aux (FCL_Funcl(id,pexp),annot)) =
       let _ = reset_fresh_name_counter () in
-      FCL_aux (FCL_Funcl (id,pat,n_exp_term newreturn exp),annot)
+      FCL_aux (FCL_Funcl (id,n_pexp newreturn pexp (fun x -> x)),annot)
     in FD_aux (FD_function(recopt,tannotopt,effectopt,List.map rewrite_funcl funcls),fdannot) in
   let rewrite_def rewriters def =
     (* let _ = Pretty_print_sail.pp_defs stderr (Defs [def]) in *)
diff --git a/src/spec_analysis.ml b/src/spec_analysis.ml
index 078fc647..b35bc48f 100644
--- a/src/spec_analysis.ml
+++ b/src/spec_analysis.ml
@@ -445,15 +445,22 @@ let fv_of_tannot_opt consider_var (Typ_annot_opt_aux (t,_)) =
     (mt, mt)
 
 (*Unlike the other fv, the bound returns are the names bound by the pattern for use in the exp*)
-let fv_of_funcl consider_var base_bounds (FCL_aux(FCL_Funcl(id,pat,exp),l)) =
-  let pat_bs,pat_ns = pat_bindings consider_var base_bounds mt pat in
-  let _, exp_ns, exp_sets = fv_of_exp consider_var pat_bs pat_ns mt exp in
-  (pat_bs,exp_ns,exp_sets)
+let fv_of_funcl consider_var base_bounds (FCL_aux(FCL_Funcl(id,pexp),l)) =
+  match pexp with
+  | Pat_aux(Pat_exp (pat,exp),_) ->
+     let pat_bs,pat_ns = pat_bindings consider_var base_bounds mt pat in
+     let _, exp_ns, exp_sets = fv_of_exp consider_var pat_bs pat_ns mt exp in
+     (pat_bs,exp_ns,exp_sets)
+  | Pat_aux(Pat_when (pat,guard,exp),_) ->
+     let pat_bs,pat_ns = pat_bindings consider_var base_bounds mt pat in
+     let bound_g, exp_ns, exp_sets = fv_of_exp consider_var pat_bs pat_ns mt exp in
+     let _, exp_ns, exp_sets = fv_of_exp consider_var bound_g exp_ns exp_sets exp in
+     (pat_bs,exp_ns,exp_sets)
 
 let fv_of_fun consider_var (FD_aux (FD_function(rec_opt,tannot_opt,_,funcls),_) as fd) =
   let fun_name = match funcls with
     | [] -> failwith "fv_of_fun fell off the end looking for the function name"
-    | FCL_aux(FCL_Funcl(id,_,_),_)::_ -> string_of_id id in
+    | FCL_aux(FCL_Funcl(id,_),_)::_ -> string_of_id id in
   let base_bounds = match rec_opt with
     | Rec_aux(Ast.Rec_rec,_) -> init_env fun_name
     | _ -> mt in
@@ -464,10 +471,18 @@ let fv_of_fun consider_var (FD_aux (FD_function(rec_opt,tannot_opt,_,funcls),_)
       bound, fv_of_typ consider_var bound mt typ
     | Typ_annot_opt_aux(Typ_annot_opt_none, _) ->
       base_bounds, mt in
-  let ns = List.fold_right (fun (FCL_aux(FCL_Funcl(_,pat,exp),_)) ns ->
-      let pat_bs,pat_ns = pat_bindings consider_var base_bounds ns pat in
-      let _, exp_ns,_ = fv_of_exp consider_var pat_bs pat_ns Nameset.empty exp in
-      exp_ns) funcls mt in
+  let ns = List.fold_right (fun (FCL_aux(FCL_Funcl(_,pexp),_)) ns ->
+    match pexp with
+    | Pat_aux(Pat_exp (pat,exp),_) ->
+       let pat_bs,pat_ns = pat_bindings consider_var base_bounds ns pat in
+       let _, exp_ns,_ = fv_of_exp consider_var pat_bs pat_ns Nameset.empty exp in
+       exp_ns
+    | Pat_aux(Pat_when (pat,guard,exp),_) ->
+       let pat_bs,pat_ns = pat_bindings consider_var base_bounds ns pat in
+       let guard_bs, guard_ns,_ = fv_of_exp consider_var pat_bs pat_ns Nameset.empty guard in
+       let _, exp_ns,_ = fv_of_exp consider_var guard_bs guard_ns Nameset.empty exp in
+       exp_ns
+  ) funcls mt in
   let ns_vs = init_env ("val:" ^ (string_of_id (id_of_fundef fd))) in
   (* let _ = Printf.eprintf "Function %s uses %s\n" fun_name (set_to_string (Nameset.union ns ns_r)) in *)
   init_env fun_name, Nameset.union ns_vs (Nameset.union ns ns_r)
@@ -481,7 +496,7 @@ let rec find_scattered_of name = function
   | DEF_scattered (SD_aux(sda,_) as sd):: defs ->
     (match sda with
      | SD_scattered_function(_,_,_,id)
-     | SD_scattered_funcl(FCL_aux(FCL_Funcl(id,_,_),_))
+     | SD_scattered_funcl(FCL_aux(FCL_Funcl(id,_),_))
      | SD_scattered_unioncl(id,_) ->
        if name = string_of_id id
        then [sd] else []
@@ -498,13 +513,25 @@ let rec fv_of_scattered consider_var consider_scatter_as_one all_defs (SD_aux(sd
         | Typ_annot_opt_aux(Typ_annot_opt_none, _) ->
           mt, mt) in
     init_env (string_of_id id),ns
-  | SD_scattered_funcl (FCL_aux(FCL_Funcl(id,pat,exp),_)) ->
-    let pat_bs,pat_ns = pat_bindings consider_var mt mt pat in
-    let _,exp_ns,_ = fv_of_exp consider_var pat_bs pat_ns Nameset.empty exp in
-    let scattered_binds = match pat with
-      | P_aux(P_app(pid,_),_) -> init_env ((string_of_id id) ^ "/" ^ (string_of_id pid))
-      | _ -> mt in
-    scattered_binds, exp_ns
+  | SD_scattered_funcl (FCL_aux(FCL_Funcl(id,pexp),_)) ->
+     begin
+       match pexp with
+       | Pat_aux(Pat_exp (pat,exp),_) ->
+          let pat_bs,pat_ns = pat_bindings consider_var mt mt pat in
+          let _,exp_ns,_ = fv_of_exp consider_var pat_bs pat_ns Nameset.empty exp in
+          let scattered_binds = match pat with
+            | P_aux(P_app(pid,_),_) -> init_env ((string_of_id id) ^ "/" ^ (string_of_id pid))
+            | _ -> mt in
+          scattered_binds, exp_ns
+       | Pat_aux(Pat_when (pat,guard,exp),_) ->
+          let pat_bs,pat_ns = pat_bindings consider_var mt mt pat in
+          let guard_bs, guard_ns,_ = fv_of_exp consider_var pat_bs pat_ns Nameset.empty guard in
+          let _, exp_ns,_ = fv_of_exp consider_var guard_bs guard_ns Nameset.empty exp in
+          let scattered_binds = match pat with
+            | P_aux(P_app(pid,_),_) -> init_env ((string_of_id id) ^ "/" ^ (string_of_id pid))
+            | _ -> mt in
+          scattered_binds, exp_ns
+     end
   | SD_scattered_variant (id,_,_) ->
     let name = string_of_id id in
     let uses =
diff --git a/src/type_check.ml b/src/type_check.ml
index 2cfba5fd..f60424ba 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -1979,6 +1979,7 @@ let irule r env exp =
 
 let strip_exp : 'a exp -> unit exp = function exp -> map_exp_annot (fun (l, _) -> (l, ())) exp
 let strip_pat : 'a pat -> unit pat = function pat -> map_pat_annot (fun (l, _) -> (l, ())) pat
+let strip_pexp : 'a pexp -> unit pexp = function pexp -> map_pexp_annot (fun (l, _) -> (l, ())) pexp
 let strip_lexp : 'a lexp -> unit lexp = function lexp -> map_lexp_annot (fun (l, _) -> (l, ())) lexp
 
 let rec check_exp env (E_aux (exp_aux, (l, ())) as exp : unit exp) (Typ_aux (typ_aux, _) as typ) : tannot exp =
@@ -2016,24 +2017,8 @@ let rec check_exp env (E_aux (exp_aux, (l, ())) as exp : unit exp) (Typ_aux (typ
      end
   | E_case (exp, cases), _ ->
      let inferred_exp = irule infer_exp env exp in
-     let rec check_case pat typ = match pat with
-       | Pat_aux (Pat_exp (P_aux (P_lit lit, _) as pat, case), annot) ->
-          let guard = mk_exp (E_app_infix (mk_exp (E_id (mk_id "p#")), mk_id "==", mk_exp (E_lit lit))) in
-          check_case (Pat_aux (Pat_when (mk_pat (P_id (mk_id "p#")), guard, case), annot)) typ
-       | Pat_aux (Pat_when (P_aux (P_lit lit, _) as pat, guard, case), annot) ->
-          let guard' = mk_exp (E_app_infix (mk_exp (E_id (mk_id "p#")), mk_id "==", mk_exp (E_lit lit))) in
-          check_case (Pat_aux (Pat_when (mk_pat (P_id (mk_id "p#")), mk_exp (E_app_infix (guard, mk_id "&", guard')), case), annot)) typ
-       | Pat_aux (Pat_exp (pat, case), (l, _)) ->
-          let tpat, env = bind_pat env pat (typ_of inferred_exp) in
-          Pat_aux (Pat_exp (tpat, crule check_exp env case typ), (l, None))
-       | Pat_aux (Pat_when (pat, guard, case), (l, _)) ->
-          let tpat, env = bind_pat env pat (typ_of inferred_exp) in
-          let checked_guard = check_exp env guard bool_typ in
-          let flows, constrs = infer_flow env checked_guard in
-          let env' = add_constraints constrs (add_flows true flows env) in
-          Pat_aux (Pat_when (tpat, checked_guard, crule check_exp env' case typ), (l, None))
-     in
-     annot_exp (E_case (inferred_exp, List.map (fun case -> check_case case typ) cases)) typ
+     let inferred_typ = typ_of inferred_exp in
+     annot_exp (E_case (inferred_exp, List.map (fun case -> check_case env inferred_typ case typ) cases)) typ
   | E_try (exp, cases), _ ->
      let checked_exp = crule check_exp env exp typ in
      let check_case pat typ = match pat with
@@ -2198,6 +2183,23 @@ let rec check_exp env (E_aux (exp_aux, (l, ())) as exp : unit exp) (Typ_aux (typ
      let inferred_exp = irule infer_exp env exp in
      type_coercion env inferred_exp typ
 
+and check_case env pat_typ pat typ = match pat with
+  | Pat_aux (Pat_exp (P_aux (P_lit lit, _) as pat, case), annot) ->
+     let guard = mk_exp (E_app_infix (mk_exp (E_id (mk_id "p#")), mk_id "==", mk_exp (E_lit lit))) in
+     check_case env pat_typ (Pat_aux (Pat_when (mk_pat (P_id (mk_id "p#")), guard, case), annot)) typ
+  | Pat_aux (Pat_when (P_aux (P_lit lit, _) as pat, guard, case), annot) ->
+     let guard' = mk_exp (E_app_infix (mk_exp (E_id (mk_id "p#")), mk_id "==", mk_exp (E_lit lit))) in
+     check_case env pat_typ (Pat_aux (Pat_when (mk_pat (P_id (mk_id "p#")), mk_exp (E_app_infix (guard, mk_id "&", guard')), case), annot)) typ
+  | Pat_aux (Pat_exp (pat, case), (l, _)) ->
+     let tpat, env = bind_pat env pat pat_typ in
+     Pat_aux (Pat_exp (tpat, crule check_exp env case typ), (l, None))
+  | Pat_aux (Pat_when (pat, guard, case), (l, _)) ->
+     let tpat, env = bind_pat env pat pat_typ in
+     let checked_guard = check_exp env guard bool_typ in
+     let flows, constrs = infer_flow env checked_guard in
+     let env' = add_constraints constrs (add_flows true flows env) in
+     Pat_aux (Pat_when (tpat, checked_guard, crule check_exp env' case typ), (l, None))
+
 (* type_coercion env exp typ takes a fully annoted (i.e. already type
    checked) expression exp, and attempts to cast (coerce) it to the
    type typ by inserting a coercion function that transforms the
@@ -3334,11 +3336,10 @@ let check_letdef env (LB_aux (letbind, (l, _))) =
        [DEF_val (LB_aux (LB_val (tpat, inferred_bind), (l, None)))], env
   end
 
-let check_funcl env (FCL_aux (FCL_Funcl (id, pat, exp), (l, _))) typ =
+let check_funcl env (FCL_aux (FCL_Funcl (id, pexp), (l, _))) typ =
   match typ with
   | Typ_aux (Typ_fn (typ_arg, typ_ret, eff), _) ->
      begin
-       let typed_pat, env = bind_pat env (strip_pat pat) typ_arg in
        let env = Env.add_ret_typ typ_ret env in
        (* We want to forbid polymorphic undefined values in all cases,
           except when type checking the specific undefined_(type)
@@ -3351,12 +3352,12 @@ let check_funcl env (FCL_aux (FCL_Funcl (id, pat, exp), (l, _))) typ =
          then Env.allow_polymorphic_undefineds env
          else env
        in
-       let exp = propagate_exp_effect (crule check_exp env (strip_exp exp) typ_ret) in
-       FCL_aux (FCL_Funcl (id, typed_pat, exp), (l, Some (env, typ, effect_of exp)))
+       let typed_pexp, prop_eff = propagate_pexp_effect (check_case env typ_arg (strip_pexp pexp) typ_ret) in
+       FCL_aux (FCL_Funcl (id, typed_pexp), (l, Some (env, typ, prop_eff)))
      end
   | _ -> typ_error l ("Function clause must have function type: " ^ string_of_typ typ ^ " is not a function type")
 
-let funcl_effect (FCL_aux (FCL_Funcl (id, typed_pat, exp), (l, annot))) =
+let funcl_effect (FCL_aux (FCL_Funcl (id, typed_pexp), (l, annot))) =
   match annot with
   | Some (_, _, eff) -> eff
   | None -> no_effect (* Maybe could be assert false. This should never happen *)
@@ -3373,7 +3374,8 @@ let infer_funtyp l env tannotopt funcls =
          | _ -> typ_error l ("Cannot infer type from pattern " ^ string_of_pat pat)
        in
        match funcls with
-       | [FCL_aux (FCL_Funcl (_, pat, _), _)] ->
+       | [FCL_aux (FCL_Funcl (_, Pat_aux (pexp,_)), _)] ->
+          let pat = match pexp with Pat_exp (pat,_) | Pat_when (pat,_,_) -> pat in
           let arg_typ = typ_from_pat pat in
           let fn_typ = mk_typ (Typ_fn (arg_typ, ret_typ, Effect_aux (Effect_set [], Parse_ast.Unknown))) in
           (quant, fn_typ)
@@ -3393,7 +3395,7 @@ let check_tannotopt env typq ret_typ = function
 let check_fundef env (FD_aux (FD_function (recopt, tannotopt, effectopt, funcls), (l, _)) as fd_aux) =
   let id =
     match (List.fold_right
-             (fun (FCL_aux (FCL_Funcl (id, _, _), _)) id' ->
+             (fun (FCL_aux (FCL_Funcl (id, _), _)) id' ->
                match id' with
                | Some id' -> if string_of_id id' = string_of_id id then Some id'
                              else typ_error l ("Function declaration expects all definitions to have the same name, "