Correctly specialise type annotation in polymorphic functions

4 files changed, 43 insertions, 2 deletions

diff --git a/src/rewriter.ml b/src/rewriter.ml
index f1de2f47..82e1ec95 100644
--- a/src/rewriter.ml
+++ b/src/rewriter.ml
@@ -388,7 +388,7 @@ let rewrite_lexp rewriters (LEXP_aux(lexp,(l,annot))) =
 
 let rewrite_fun rewriters (FD_aux (FD_function(recopt,tannotopt,effectopt,funcls),(l,fdannot))) = 
   let rewrite_funcl (FCL_aux (FCL_Funcl(id,pexp),(l,annot))) =
-    (FCL_aux (FCL_Funcl (id,rewrite_pexp rewriters 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/specialize.ml b/src/specialize.ml
index b1f1f4b3..9723689f 100644
--- a/src/specialize.ml
+++ b/src/specialize.ml
@@ -325,6 +325,27 @@ let specialize_id_valspec instantiations id ast =
 
      append_ast pre_ast (append_ast (Defs (vs :: specializations)) post_ast)
 
+(* When we specialize a function definition we also need to specialize
+   all the types that appear as annotations within the function
+   body. *)
+let specialize_annotations instantiation =
+  let open Type_check in
+  let rw_pat = {
+      id_pat_alg with
+      p_typ = (fun (typ, pat) -> P_typ (subst_unifiers instantiation typ, pat))
+    } in
+  let rw_exp = {
+      id_exp_alg with
+      e_cast = (fun (typ, exp) -> E_cast (subst_unifiers instantiation typ, exp));
+      lEXP_cast = (fun (typ, lexp) -> LEXP_cast (subst_unifiers instantiation typ, lexp));
+      pat_alg = rw_pat
+    } in
+  rewrite_fun {
+      rewriters_base with
+      rewrite_exp = (fun _ -> fold_exp rw_exp);
+      rewrite_pat = (fun _ -> fold_pat rw_pat)
+    }
+     
 let specialize_id_fundef instantiations id ast =
   match split_defs (is_fundef id) ast with
   | None -> ast
@@ -335,7 +356,7 @@ let specialize_id_fundef instantiations id ast =
        if IdSet.mem spec_id !spec_ids then [] else
          begin
            spec_ids := IdSet.add spec_id !spec_ids;
-           [DEF_fundef (rename_fundef spec_id fundef)]
+           [DEF_fundef (specialize_annotations instantiation (rename_fundef spec_id fundef))]
          end
      in
      let fundefs = List.map specialize_fundef instantiations |> List.concat in
diff --git a/test/c/special_annot.expect b/test/c/special_annot.expect
new file mode 100644
index 00000000..ee8a39c3
--- /dev/null
+++ b/test/c/special_annot.expect
@@ -0,0 +1 @@
+string
diff --git a/test/c/special_annot.sail b/test/c/special_annot.sail
new file mode 100644
index 00000000..ca2f5f4c
--- /dev/null
+++ b/test/c/special_annot.sail
@@ -0,0 +1,19 @@
+
+union option ('a : Type) = {
+  None : unit,
+  Some : 'a
+}
+
+val test : forall ('a : Type). 'a -> option('a)
+
+function test(x) =
+  let y : 'a = x in Some(y)
+
+val "print_endline" : string -> unit
+
+val main : unit -> unit
+
+function main() = {
+  let Some(y) = test("string");
+  print_endline(y)
+}
\ No newline at end of file