Close #10

1 files changed, 55 insertions, 38 deletions

diff --git a/src/trace/smtAtom.ml b/src/trace/smtAtom.ml
index ce28930..a794def 100644
--- a/src/trace/smtAtom.ml
+++ b/src/trace/smtAtom.ml
@@ -744,41 +744,6 @@ module Atom =
     let to_smt (fmt:Format.formatter) h = to_smt_named fmt h
 
 
-    exception NotWellTyped of atom
-
-    let check a =
-      (* Format.eprintf "Checking %a @." to_smt_atom a; *)
-      match a with
-      | Acop _ -> ()
-      | Auop(op,h) ->
-	if not (SmtBtype.equal (Op.u_type_arg op) (type_of h)) then
-	  raise (NotWellTyped a)
-      | Abop(op,h1,h2) ->
-	let (t1,t2) = Op.b_type_args op in
- 	if not (SmtBtype.equal t1 (type_of h1) && SmtBtype.equal t2 (type_of h2))
-        then (Format.eprintf "1. Wanted %a, got %a@.2. Wanted %a, got %a@."
-                SmtBtype.to_smt t1 SmtBtype.to_smt (type_of h1)
-                SmtBtype.to_smt t2 SmtBtype.to_smt (type_of h2);
-              raise (NotWellTyped a))
-      | Atop(op,h1,h2,h3) ->
-	let (t1,t2,t3) = Op.t_type_args op in
-        if not (SmtBtype.equal t1 (type_of h1) &&
-                SmtBtype.equal t2 (type_of h2) &&
-                SmtBtype.equal t3 (type_of h3))
-	then raise (NotWellTyped a)
-      | Anop(op,ha) ->
-        let ty = Op.n_type_args op in
-        Array.iter
-          (fun h -> if not (SmtBtype.equal ty (type_of h)) then
-              raise (NotWellTyped a)) ha
-      | Aapp(op,args) ->
-	let tparams = Op.i_type_args op in
-	Array.iteri (fun i t ->
-	    if not (SmtBtype.equal t (type_of args.(i))) then
-	      (Format.eprintf "Wanted %a, got %a@."
-                SmtBtype.to_smt t SmtBtype.to_smt (type_of args.(i));
-                      raise (NotWellTyped a))) tparams
-
     type reify_tbl =
       { mutable count : int;
 	tbl : hatom HashAtom.t
@@ -795,19 +760,71 @@ module Atom =
       reify.count <- 0;
       HashAtom.clear reify.tbl
 
+
+    exception NotWellTyped of atom
+
     let declare reify a =
-      check a;
       let res = {index = reify.count; hval = a} in
       HashAtom.add reify.tbl a res;
       reify.count <- reify.count + 1;
       res
 
-    let get ?declare:(decl=true) reify a =
+    let rec get ?declare:(decl=true) reify a =
       if decl
       then try HashAtom.find reify.tbl a
-           with Not_found -> declare reify a
+           with Not_found ->
+             (let a = check reify a in
+              try HashAtom.find reify.tbl a
+              with Not_found ->
+                declare reify a)
       else {index = -1; hval = a}
 
+    and check reify a =
+      (* Format.eprintf "Checking %a @." to_smt_atom a; *)
+
+      let check_one t h =
+        let th = type_of h in
+        if SmtBtype.equal t th then
+          h
+        else if t == TZ && th == Tpositive then
+          (* Special case: the SMT solver cannot distinguish Z from
+             positive, we have to add the injection back *)
+          get reify (Auop(UO_Zpos, h))
+        else (
+          Format.eprintf "Incorrect type: wanted %a, got %a@."
+            SmtBtype.to_smt t SmtBtype.to_smt th;
+	  raise (NotWellTyped a)
+        )
+      in
+
+      let a =
+        match a with
+          | Acop _ -> a
+          | Auop(op,h) ->
+             let t = Op.u_type_arg op in
+             Auop(op, check_one t h)
+          | Abop(op,h1,h2) ->
+	     let (t1,t2) = Op.b_type_args op in
+             let h1 = check_one t1 h1 in
+             let h2 = check_one t2 h2 in
+             Abop(op, h1, h2)
+          | Atop(op,h1,h2,h3) ->
+	     let (t1,t2,t3) = Op.t_type_args op in
+             let h1 = check_one t1 h1 in
+             let h2 = check_one t2 h2 in
+             let h3 = check_one t3 h3 in
+             Atop(op, h1, h2, h3)
+          | Anop(op,ha) ->
+             let ty = Op.n_type_args op in
+             Anop(op, Array.map (check_one ty) ha)
+          | Aapp(op,args) ->
+             let tparams = Op.i_type_args op in
+             Aapp(op, Array.mapi (fun i -> check_one tparams.(i)) args)
+      in
+      close_out outc;
+      a
+
+
     let mk_eq reify ?declare:(decl=true) ty h1 h2 =
       let op = BO_eq ty in
       try