Fix persistent func typing to compile motoko-core

src/mo_frontend/bi_match.ml

@@ -261,7 +261,7 @@ let bi_match_typs ctx =
     | Tup ts1, Tup ts2 ->
       bi_match_list bi_match_typ rel eq inst any ts1 ts2
     | Func (s1, c1, tbs1, t11, t12), Func (s2, c2, tbs2, t21, t22) ->
-      if s1 = s2 && c1 = c2 then
+      if rel_func_sort ~eq:(rel == eq) s1 s2 tbs1 && c1 = c2 then
       (match bi_match_binds rel eq inst any tbs1 tbs2 with
        | Some (inst, ts) ->
          let any' = List.fold_right

src/mo_frontend/typing.ml

@@ -1697,10 +1697,6 @@ and infer_exp'' env exp : T.typ =
     | _ -> ()
     end;
     let sort, ve = check_shared_pat env shared_pat in
-    let sort = match sort, !Mo_config.Flags.enhanced_orthogonal_persistence with
-    | T.Local T.Stable, false -> T.Local T.Flexible (* named local functions are flexible in classical mode *)
-    | _ -> sort
-    in
     let is_async = match typ_opt with
       | Some { it = AsyncT _; _ } -> true
       | _ -> false
@@ -1769,7 +1765,6 @@ and infer_exp'' env exp : T.typ =
       end
     end;
     let ts1 = match pat.it with TupP _ -> T.seq_of_tup t1 | _ -> [t1] in
-    let sort = if is_flexible && sort = T.Local T.Stable then T.Local T.Flexible else sort in
     T.Func (sort, c, T.close_binds cs tbs, List.map (T.close cs) ts1, List.map (T.close cs) ts2)
   | CallE (par_opt, exp1, inst, exp2) ->
     let t = infer_call env exp1 inst exp2 exp.at None in
@@ -2308,7 +2303,8 @@ and check_func_step in_actor env (shared_pat, pat, typ_opt, exp) (s, c, ts1, ts2
   if sort <> s then
     (match sort, s with
       | T.Local T.Stable, T.Local T.Flexible -> () (* okay *)
-      | T.Local _, T.Local _ -> assert false (* caught by sub-type check *)
+      | T.Local _, T.Local _ ->
+        error env exp.at "M0220" "this function should be declared `persistent`"
       | _, _ ->
         error env exp.at "M0094"
           "%sshared function does not match expected %sshared function type"

src/mo_types/type.ml

@@ -1223,7 +1223,7 @@ let rec rel_typ d rel eq t1 t2 =
   | Tup ts1, Tup ts2 ->
     rel_list "tuple arguments" d rel_typ rel eq ts1 ts2
   | Func (s1, c1, tbs1, t11, t12), Func (s2, c2, tbs2, t21, t22) ->
-    (rel_sort s1 s2 tbs1 || incompatible_func_sorts d t1 t2) &&
+    (rel_func_sort ~eq:(rel == eq) s1 s2 tbs1 || incompatible_func_sorts d t1 t2) &&
     (c1 = c2 || incompatible_bounds d t1 t2) &&
     (match rel_binds d eq eq tbs1 tbs2 with
      | Some ts ->
@@ -1238,7 +1238,10 @@ let rec rel_typ d rel eq t1 t2 =
   | _, _ -> incompatible_types d t1 t2
   end
 
-and rel_sort s1 s2 tbs1 =
+and rel_func_sort ~eq s1 s2 tbs1 =
+  if eq then s1 = s2 else sub_func_sort s1 s2 tbs1
+
+and sub_func_sort s1 s2 tbs1 =
   match s1, s2 with
   | Local Stable, Local Flexible -> 
     not (List.exists (fun b -> b.sort = Type) tbs1)
@@ -1592,19 +1595,22 @@ let rec combine rel lubs glbs t1 t2 =
     | Obj (s1, tf1), Obj (s2, tf2) when s1 = s2 ->
       (try Obj (s1, combine_fields rel lubs glbs tf1 tf2)
       with Mismatch -> assert (rel == glbs); Non)
-    | Func (s1, c1, bs1, ts11, ts12), Func (s2, c2, bs2, ts21, ts22) when
-        s1 = s2 && c1 = c2 && eq_binds bs1 bs2 &&
-        List.(length ts11 = length ts21 && length ts12 = length ts22) ->
-      let ts = open_binds bs1 in
-      let cs = List.map (fun t -> fst (as_con t)) ts in
-      let opened = List.map (open_ ts) in
-      let closed = List.map (close cs) in
-      let rel' = if rel == lubs then glbs else lubs in
-      Func (
-        s1, c1, bs1,
-        closed (List.map2 (combine rel' lubs glbs) (opened ts11) (opened ts21)),
-        closed (List.map2 (combine rel lubs glbs) (opened ts12) (opened ts22))
-      )
+    | Func (s1, c1, bs1, ts11, ts12), Func (s2, c2, bs2, ts21, ts22) ->
+      begin match combine_func_sort rel lubs glbs s1 s2 bs1 bs2 with
+      | Some s when c1 = c2 && eq_binds bs1 bs2 &&
+          List.(length ts11 = length ts21 && length ts12 = length ts22) ->
+        let ts = open_binds bs1 in
+        let cs = List.map (fun t -> fst (as_con t)) ts in
+        let opened = List.map (open_ ts) in
+        let closed = List.map (close cs) in
+        let rel' = if rel == lubs then glbs else lubs in
+        Func (
+          s, c1, bs1,
+          closed (List.map2 (combine rel' lubs glbs) (opened ts11) (opened ts21)),
+          closed (List.map2 (combine rel lubs glbs) (opened ts12) (opened ts22))
+        )
+      | _ -> if rel == lubs then Any else Non
+      end
     | Async (s1, t11, t12), Async (s2, t21, t22) when s1 == s2 && eq t11 t21 ->
       Async (s1, t11, combine rel lubs glbs t12 t22)
     | Con _, _
@@ -1647,6 +1653,13 @@ let rec combine rel lubs glbs t1 t2 =
     | _, _ ->
       if rel == lubs then Any else Non
 
+and combine_func_sort rel lubs glbs s1 s2 bs1 bs2 =
+  if s1 = s2 then Some s1 else
+  if rel == glbs then None else
+  let s = Local Flexible in
+  if sub_func_sort s1 s bs1 && sub_func_sort s2 s bs2 then Some s
+  else None
+
 and cons_if b x xs = if b then x::xs else xs
 
 and combine_fields rel lubs glbs fs1 fs2 =

src/mo_types/type.mli

@@ -277,6 +277,9 @@ val sub : ?src_fields : Field_sources.t -> typ -> typ -> bool
 val sub_explained : ?src_fields : Field_sources.t -> context -> typ -> typ -> compatibility
 val compatible : typ -> typ -> bool
 
+val rel_func_sort : eq:bool -> func_sort -> func_sort -> bind list -> bool
+val sub_func_sort : func_sort -> func_sort -> bind list -> bool
+
 exception PreEncountered
 val lub : ?src_fields : Field_sources.t -> typ -> typ -> typ
 val glb : ?src_fields : Field_sources.t -> typ -> typ -> typ

test/fail/ok/persistent-func-type-argument-and-return.tc.ok

@@ -0,0 +1 @@
+persistent-func-type-argument-and-return.mo:7.40-7.42: type error [M0220], this function should be declared `persistent`

test/fail/ok/persistent-func-type-argument-and-return.tc.ret.ok

@@ -0,0 +1 @@
+Return code 1

test/fail/persistent-func-type-argument-and-return.mo

@@ -0,0 +1,8 @@
+module {
+  persistent func g<T>(x : T) : persistent() -> T {
+    persistent func h() : T { x };
+    h;
+  };
+  // don't allow `func() {}` to typecheck as `persistent() -> ()` for now
+  let h = g<persistent() -> ()>(func() {});
+};

test/run-drun/ok/persistent-actor-with-persistent-funcs-in-type-args-and-return.drun-run.ok

@@ -0,0 +1,2 @@
+ingress Completed: Reply: 0x4449444c016c01b3c4b1f204680100010a00000000000000000101
+ingress Completed: Reply: 0x4449444c0000

test/run-drun/persistent-actor-with-persistent-funcs-in-type-args-and-return.mo

@@ -0,0 +1,9 @@
+//ENHANCED-ORTHOGONAL-PERSISTENCE-ONLY
+persistent actor {
+  persistent func g<T>(x : T) : persistent() -> T {
+    persistent func h() : T { x };
+    h;
+  };
+  let h = g<persistent() -> ()>(persistent func _anon() {});
+  h();
+};

test/run/use-persistent-function.mo

@@ -0,0 +1,23 @@
+type Order = { #less; #equal; #greater };
+
+module Array {
+  public func sort<T>(array : [T], compare : (T, T) -> Order) : [T] {
+    if (compare(array[0], array[1]) == #greater) {
+      [array[1], array[0]];
+    } else {
+      array;
+    };
+  };
+};
+
+module Nat {
+  public persistent func compare(x : Nat, y : Nat) : Order {
+    if (x < y) { #less } else if (x == y) { #equal } else { #greater };
+  };
+};
+
+let a : [Nat] = [2, 1];
+let b : [Nat] = [1, 2];
+
+assert Array.sort(a, Nat.compare) == [1, 2];
+assert Array.sort(b, Nat.compare) == [1, 2];