From 79176c7ac58900c95db5db9cf75978e91e9a453c Mon Sep 17 00:00:00 2001
From: Alasdair
Date: Sat, 18 Aug 2018 00:47:10 +0100
Subject: Correctly specialise type annotation in polymorphic functions

---
 src/rewriter.ml   |  2 +-
 src/specialize.ml | 23 ++++++++++++++++++++++-
 2 files changed, 23 insertions(+), 2 deletions(-)

(limited to 'src')

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
-- 
cgit v1.2.3