Support polymorphic recursion

johnynek · Sep 10, 2023 · d1ddbb8 · d1ddbb8
1 parent dd86eea
commit d1ddbb8
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 3 deletions.
diff --git a/core/src/main/scala/org/bykn/bosatsu/Expr.scala b/core/src/main/scala/org/bykn/bosatsu/Expr.scala
@@ -38,6 +38,11 @@ object Expr {
       case None => expr
       case Some(nel) =>
         expr match {
+          case Annotation(expr, tpe, tag) =>
+            val tpeFrees = Type.freeBoundTyVars(tpe :: Nil).toSet
+            // these are the frees that are also in tpeArgs
+            val freeArgs = tpeArgs.filter { case (n, _) => tpeFrees(n) }
+            Annotation(forAll(tpeArgs, expr), Type.forAll(freeArgs, tpe), tag)
           case Generic(typeVars, in) =>
             Generic(nel ::: typeVars, in)
           case notAnn => Generic(nel, notAnn)

diff --git a/core/src/main/scala/org/bykn/bosatsu/SourceConverter.scala b/core/src/main/scala/org/bykn/bosatsu/SourceConverter.scala
@@ -131,7 +131,7 @@ final class SourceConverter(
                 case Some(k) => k
               })
             }
-          val gen = Expr.Generic(bs, lambda)
+          val gen = Expr.forAll(bs.toList, lambda)
           val freeVarsList = Expr.freeBoundTyVars(lambda)
           val freeVars = freeVarsList.toSet
           val notFreeDecl = bs.exists { case (a, _) => !freeVars(a) }

diff --git a/core/src/main/scala/org/bykn/bosatsu/rankn/Infer.scala b/core/src/main/scala/org/bykn/bosatsu/rankn/Infer.scala
@@ -1335,8 +1335,13 @@ object Infer {
 
   private def recursiveTypeCheck[A: HasRegion](name: Bindable, expr: Expr[A]): Infer[TypedExpr[A]] =
     // values are of kind Type
-    newMetaType(Kind.Type).flatMap { tpe =>
-      extendEnv(name, tpe)(typeCheckMeta(expr, Some((name, tpe, region(expr)))))
+    expr match {
+      case Expr.Annotation(e, tpe, _) =>
+        extendEnv(name, tpe)(checkSigma(e, tpe))
+      case _ =>
+        newMetaType(Kind.Type).flatMap { tpe =>
+          extendEnv(name, tpe)(typeCheckMeta(expr, Some((name, tpe, region(expr)))))
+        }
     }
 
 

diff --git a/core/src/test/scala/org/bykn/bosatsu/EvaluationTest.scala b/core/src/test/scala/org/bykn/bosatsu/EvaluationTest.scala
@@ -2973,4 +2973,24 @@ def last[a](nel: NEList[a]) -> a:
 test = Assertion(last(One(True)), "")
 """), "Generic", 1)
   }
+
+  test("support polymorphic recursion") {
+    runBosatsuTest(
+      List("""
+package PolyRec
+
+enum Nat: NZero, NSucc(n: Nat)
+
+def poly_rec(count: Nat, a: a) -> a:
+    recur count:
+        case NZero: a
+        case NSucc(prev):
+          # make a call with a different type
+          (_, b) = poly_rec(prev, ("foo", a))
+          b
+
+test = Assertion(True, "")         
+""")
+    , "PolyRec", 1)
+  }
 }