Fix issue with flat functions having refinement (non-flat) annotations

The underlying issue was that a step with a call to _checkAnn was created for
refinements in flat functions, with no corresponding restart code for the stack
capture. In theory, this regression test should actually work as a flat
function, since the is-predicate is flat itself. However, PPredAnn _always_
does a safeCall when checking the annotation, which can cause a stack capture.

A few things going forward:

- This fix is "obviously" correct, but way more pessimistic than it needs to be
  long-term. It currently makes all functions with not obviously trivial
  annotations be non-flat. It really should use more static information to
  check if any annotation has a non-flat refinement. This is not information we
  have floating around because we cannot track this information through type
  aliases (and hence across modules).
- A slightly less pessimistic, but still naive way to do this would be to check
  if annotations directly contain a refinement, or if they are "unknown"
  aliases other than names like Number, String, or a datatype in the current
  module. Then we could compile those names to a new kind of dynamic annotation
  called PFlatPredAnn (in contrast with PPredAnn), which would count as a
  "cheapAnnotation" that doesn't get a safeCall around it.
- The right thing to do is enrich all static annotation information with
  flatness information, since any compound type (e.g. tuple or object) or
  aliased type could refer to a predicate.

Author: jpolitz

src/arr/compiler/js-of-pyret.arr

@@ -228,6 +228,11 @@ fun flatness-max(a :: Option<Number>, b :: Option<Number>) -> Option<Number>:
   end
 end
 
+fun is-trivial-ann(ann):
+  A.is-a-blank(ann) or A.is-a-any(ann)
+    or (A.is-a-tuple(ann) and ann.fields.all(lam(field-a): A.is-a-blank(field-a) or A.is-a-any(field-a) end))
+end
+
 # Maybe compress Option<Number> into a type like FlatnessInfo or something (maybe something without "Info" in the name)
 fun make-expr-flatness-env(
     aexpr :: AA.AExpr,
@@ -237,7 +242,17 @@ fun make-expr-flatness-env(
       make-expr-flatness-env(body, sd)
     | a-let(_, bind, val, body) =>
       val-flatness = if AA.is-a-lam(val) block:
-        lam-flatness = make-expr-flatness-env(val.body, sd)
+        has-nontrivial-arg-ann = for lists.any(elt from val.args):
+          not(is-trivial-ann(elt))
+        end
+        has-nontrivial-ann = has-nontrivial-arg-ann or not(is-trivial-ann(val.ret))
+
+        lam-flatness = if has-nontrivial-ann:
+          none
+        else:
+          make-expr-flatness-env(val.body, sd)
+        end
+
         sd.set-now(bind.id.key(), lam-flatness)
         # flatness of defining this lambda is 0, since we're not actually
         # doing anything with it

tests/pyret/regression.arr

@@ -21,3 +21,4 @@ import file("./regression/bogus-global-type-name.arr") as _
 import file("./regression/duplicate-check-block-report.arr") as _
 import file("./regression/escaping-module-uris.arr") as _
 import file("./regression/tail-recursion-arg-order.arr") as _
+import file("./regression/flat-refinement.arr") as _

tests/pyret/regression/flat-refinement.arr

@@ -0,0 +1,28 @@
+data D:
+  | d(x)
+end
+
+
+fun flat(my-d :: D%(is-d)):
+  my-d.x
+end
+
+some-d = d(10)
+
+fun recur(upto):
+  if upto == 0 block: nothing
+  else:
+    thing2 = flat(some-d) + 1
+    ans = recur(upto - 1)
+    thing = 1 + 1
+    ans
+  end
+end
+
+check:
+  recur(3000) is nothing
+  recur(300) is nothing
+  recur(500) is nothing
+  recur(391) is nothing
+  recur(5) is nothing
+end