Find buried set constraints in asserts

1 files changed, 35 insertions, 28 deletions

diff --git a/src/monomorphise.ml b/src/monomorphise.ml
index ca595278..873154cb 100644
--- a/src/monomorphise.ml
+++ b/src/monomorphise.ml
@@ -2799,29 +2799,42 @@ let merge_set_asserts_by_kid sets1 sets2 =
 
 (* Set constraints in assertions don't always use the set syntax, so we also
    handle assert('N == 1 | ...) style set constraints *)
-let set_from_exp e =
-  let mykid = ref None in
-  let check_kid kid =
-    match !mykid with
-    | None -> mykid := Some kid
-    | Some kid' -> if Kid.compare kid kid' == 0 then ()
+let rec sets_from_assert e =
+  let set_from_or_exps (E_aux (_,(l,_)) as e) =
+    let mykid = ref None in
+    let check_kid kid =
+      match !mykid with
+      | None -> mykid := Some kid
+      | Some kid' -> if Kid.compare kid kid' == 0 then ()
         else raise Not_found
+    in
+    let rec aux (E_aux (e,_)) =
+      match e with
+      | E_app (Id_aux (Id "or_bool",_),[e1;e2]) ->
+         aux e1 @ aux e2
+      | E_app (Id_aux (Id "eq_atom",_),
+               [E_aux (E_sizeof (Nexp_aux (Nexp_var kid,_)),_);
+                E_aux (E_lit (L_aux (L_num i,_)),_)]) ->
+         (check_kid kid; [i])
+      | _ -> raise Not_found
+    in try
+         let is = aux e in
+         match !mykid with
+         | None -> KBindings.empty
+         | Some kid -> KBindings.singleton kid (l,is)
+      with Not_found -> KBindings.empty
   in
-  let rec aux (E_aux (e,_)) =
-    match e with
-    | E_app (Id_aux (Id "or_bool",_),[e1;e2]) ->
-       aux e1 @ aux e2
-    | E_app (Id_aux (Id "eq_atom",_),
-             [E_aux (E_sizeof (Nexp_aux (Nexp_var kid,_)),_);
-              E_aux (E_lit (L_aux (L_num i,_)),_)]) ->
-       (check_kid kid; [i])
-    | _ -> raise Not_found
-  in try
-       let is = aux e in
-       match !mykid with
-       | None -> None
-       | Some kid -> Some (kid,is)
-    with Not_found -> None
+  let rec sets_from_nc (NC_aux (nc,l)) =
+    match nc with
+    | NC_and (nc1,nc2) -> merge_set_asserts_by_kid (sets_from_nc nc1) (sets_from_nc nc2)
+    | NC_set (kid,is) -> KBindings.singleton kid (l,is)
+    | _ -> KBindings.empty
+  in
+  match e with
+  | E_aux (E_app (Id_aux (Id "and_bool",_),[e1;e2]),_) ->
+     merge_set_asserts_by_kid (sets_from_assert e1) (sets_from_assert e2)
+  | E_aux (E_constraint nc,_) -> sets_from_nc nc
+  | _ -> set_from_or_exps e
 
 (* Find all the easily reached set assertions in a function body, to use as
    case splits *)
@@ -2837,13 +2850,7 @@ let rec find_set_assertions (E_aux (e,_)) =
      let kbound = kids_bound_by_pat p in
      let sets2 = KBindings.filter (fun kid _ -> not (KidSet.mem kid kbound)) sets2 in
      merge_set_asserts_by_kid sets1 sets2
-  | E_assert (E_aux (e1,_) as exp1,_) -> begin
-     match e1 with
-     | E_constraint (NC_aux (NC_set (kid,is),l)) -> KBindings.singleton kid (l,is)
-     | _ -> match set_from_exp exp1 with
-       | None -> KBindings.empty
-       | Some (kid,is) -> KBindings.singleton kid (exp_loc exp1,is)
-     end
+  | E_assert (exp1,_) -> sets_from_assert exp1
   | _ -> KBindings.empty
 
 let print_set_assertions set_assertions =