From ab1a95b54a1cd8acacf652b5b75db20cc427f912 Mon Sep 17 00:00:00 2001
From: Vesa Karvonen <vesa.a.j.k@gmail.com>
Date: Sat, 28 Dec 2019 15:38:42 +0200
Subject: [PATCH] Add support for type rec over types with type parameters

---
 elab.ml        | 53 ++++++++++++++++++++++++++++------------
 regression.1ml | 65 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+), 15 deletions(-)

diff --git a/elab.ml b/elab.ml
index d98ae02b..915b540c 100644
--- a/elab.ml
+++ b/elab.ml
@@ -122,6 +122,24 @@ and paths_row ta ps = function
     ts1 @ ts2
 
 
+let rec_from_extyp typ label s =
+  match s with
+  | ExT([], t) ->
+    let rec find_rec = function
+      | AppT(t, ts) ->
+        let rec_t, unroll_t, roll_t, ak = find_rec t in
+        rec_t, AppT(unroll_t, ts), AppT(roll_t, ts), ak
+      | RecT(ak, unroll_t) as rec_t ->
+        rec_t, unroll_t, rec_t, ak
+      | _ ->
+        error typ.at ("non-recursive type for " ^ label ^ ":\n"
+                      ^ "  " ^ Types.string_of_extyp s) in
+    find_rec t
+  | _ ->
+     error typ.at ("non-recursive type for " ^ label ^ ":\n"
+                   ^ "  " ^ Types.string_of_extyp s)
+
+
 (* Instantiation *)
 
 let rec instantiate env t e =
@@ -386,15 +404,18 @@ Trace.debug (lazy ("[FunE] env =" ^ VarSet.fold (fun a s -> s ^ " " ^ a) (domain
 
   | EL.RollE(var, typ) ->
     let s, zs1 = elab_typ env typ l in
-    let t, ak, t' =
-      match s with
-      | ExT([], (RecT(ak, t') as t)) -> t, ak, t'
-      | _ -> error typ.at "non-recursive type for rolling" in
+    let rec_t, unroll_t, roll_t, ak = rec_from_extyp typ "rolling" s in
+    let var_t = lookup_var env var in
+    let unroll_t = subst_typ (subst [ak] [rec_t]) unroll_t in
     let _, zs2, f =
-      try sub_typ env (lookup_var env var) (subst_typ (subst [ak] [t]) t') []
-      with Sub e -> error var.at ("rolled value does not match annotation") in
-    ExT([], t), Pure, zs1 @ zs2,
-    IL.RollE(IL.AppE(f, IL.VarE(var.it)), erase_typ t)
+      try sub_typ env var_t unroll_t []
+      with Sub e ->
+        error var.at ("rolled value does not match annotation:\n"
+                      ^ "  " ^ Types.string_of_typ var_t ^ "\n"
+                      ^ "vs\n"
+                      ^ "  " ^ Types.string_of_typ unroll_t) in
+    ExT([], roll_t), Pure, zs1 @ zs2,
+    IL.RollE(IL.AppE(f, IL.VarE(var.it)), erase_typ roll_t)
 
   | EL.IfE(var, exp1, exp2, typ) ->
     let t0, zs0, ex = elab_instvar env var in
@@ -488,13 +509,15 @@ Trace.debug (lazy ("[UnwrapE] s2 = " ^ string_of_norm_extyp s2));
 
   | EL.UnrollE(var, typ) ->
     let s, zs1 = elab_typ env typ l in
-    let t, ak, t' =
-      match s with
-      | ExT([], (RecT(ak, t') as t)) -> t, ak, t'
-      | _ -> error typ.at "non-recursive type for rolling" in
-    let _, zs2, f = try sub_typ env (lookup_var env var) t [] with Sub e ->
-      error var.at ("unrolled value does not match annotation") in
-    ExT([], subst_typ (subst [ak] [t]) t'), Pure, zs1 @ zs2,
+    let rec_t, unroll_t, roll_t, ak = rec_from_extyp typ "unrolling" s in
+    let var_t = lookup_var env var in
+    let _, zs2, f = try sub_typ env var_t roll_t [] with Sub e ->
+      error var.at ("unrolled value does not match annotation:\n"
+                    ^ "  " ^ Types.string_of_typ var_t ^ "\n"
+                    ^ "vs\n"
+                    ^ "  " ^ Types.string_of_typ roll_t) in
+    let unroll_t = subst_typ (subst [ak] [rec_t]) unroll_t in
+    ExT([], unroll_t), Pure, zs1 @ zs2,
     IL.UnrollE(IL.AppE(f, IL.VarE(var.it)))
 
   | EL.RecE(var, typ, exp1) ->
diff --git a/regression.1ml b/regression.1ml
index a948545e..f2098165 100644
--- a/regression.1ml
+++ b/regression.1ml
@@ -21,3 +21,68 @@ Equivalence: {
   symmetry 'a 'b (eq: t a b) : t b a =
     wr (fun (type p _) => un eq (fun (type b) => type p b -> p a) id);
 };
+
+Hungry = {
+  type eat a = rec eat_a => a -> eat_a;
+
+  eater 'a: eat a = rec (eater: eat a) => @(eat a) (fun a => eater);
+
+  (<+) eater x = eater.@(eat _) x;
+
+  do eater <+ 1 <+ 2;
+};
+
+PolyRec = {
+  type l a = rec (type t) => alt a t;
+  t = rec (type t a) => fun (type a) => alt a (t (type (a, a)));
+
+  t_int = t int;
+
+  hmm (x: t int) = casealt (x.@(t int));
+
+  t0 = @(t int) (right (@(t (type (int, int))) (left (0, 0))));
+};
+
+N :> {
+  type Z;
+  type S _;
+} = {
+  type Z   = {};
+  type S _ = {};
+};
+
+ListN = let
+  type I (type x) (type p _ _) (type t _ _) = {
+    nil     :               p x N.Z;
+    (::) 'n : x -> t x n -> p x (N.S n);
+  };
+  type T x n (type t _ _) = (type p _ _) => I x p t -> p x n;
+in {
+  t = rec (type t _ _) => fun (type x) (type n) => type wrap T x n t;
+
+  case 'x 'n (p: {type t _ _}) (cs: I x p.t t) e =
+    (unwrap e.@(t _ _): wrap T x n t) p.t cs;
+
+  local
+    mk 'x 'n (c: T x n t) = @(t x n) (wrap c: wrap T x n t);
+  in
+    nil  'x                       = mk (fun (type p _ _) (r: I x p t) => r.nil);
+    (::) 'x 'n (v: x) (vs: t x n) = mk (fun (type p _ _) (r: I x p t) => r.:: v vs);
+  end;
+} :> {
+  type t _ _;
+
+  case 'x 'n: (p: {type t _ _}) => I x p.t t => t x n -> p.t x n;
+
+  nil  'x    :               t x N.Z;
+  (::) 'x 'n : x => t x n => t x (N.S n);
+};
+
+ListN = {
+  ...ListN;
+  map 'x 'y 'n (xy: x -> y) = rec (map: 'n => t x n -> t y n) =>
+    case {type t _ n = t y n} {
+      nil = nil;
+      (::) x xs = xy x :: map xs;
+    };
+};