From 75208f43288c90ee82b7fa24748bc793fbd06da1 Mon Sep 17 00:00:00 2001
From: Sam Tobin-Hochstadt <samth@ccs.neu.edu>
Date: Thu, 13 Jul 2017 15:03:58 -0400
Subject: [PATCH] Improve loop inference heuristics.

Closes #579
---
 .../typed-racket/typecheck/tc-app/tc-app-lambda.rkt         | 3 ++-
 .../typed-racket/typecheck/tc-app/tc-app-list.rkt           | 2 ++
 typed-racket-test/unit-tests/typecheck-tests.rkt            | 6 ++++++
 3 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/typed-racket-lib/typed-racket/typecheck/tc-app/tc-app-lambda.rkt b/typed-racket-lib/typed-racket/typecheck/tc-app/tc-app-lambda.rkt
index 59d2b745..68a45d21 100644
--- a/typed-racket-lib/typed-racket/typecheck/tc-app/tc-app-lambda.rkt
+++ b/typed-racket-lib/typed-racket/typecheck/tc-app/tc-app-lambda.rkt
@@ -88,7 +88,8 @@
      #:when (free-identifier=? #'val #'e3)
      (let ([ts (for/list ([ac (in-syntax #'(actuals ...))]
                           [f (in-syntax #'(acc ...))])
-                 (let ([type (type-annotation f #:infer #t)])
+                 (let ([type (or (type-annotation f #:infer #t)
+                                 (find-annotation #'inner-body f))])
                    (if type
                        (tc-expr/check/t ac (ret type))
                        (generalize (tc-expr/t ac)))))]
diff --git a/typed-racket-lib/typed-racket/typecheck/tc-app/tc-app-list.rkt b/typed-racket-lib/typed-racket/typecheck/tc-app/tc-app-list.rkt
index 85936634..8fc941e0 100644
--- a/typed-racket-lib/typed-racket/typecheck/tc-app/tc-app-list.rkt
+++ b/typed-racket-lib/typed-racket/typecheck/tc-app/tc-app-list.rkt
@@ -130,5 +130,7 @@
        (match (single-value #'arg)
          [(tc-result1: (List: ts))
           (ret (-Tuple (reverse ts)))]
+         [(tc-result1: (and r (Listof: t)))
+          (ret r)]
          [arg-ty
           (tc/funapp #'fun #'(arg) (tc-expr/t #'fun) (list arg-ty) expected)])])))
diff --git a/typed-racket-test/unit-tests/typecheck-tests.rkt b/typed-racket-test/unit-tests/typecheck-tests.rkt
index 182f6f0e..c3b6171e 100644
--- a/typed-racket-test/unit-tests/typecheck-tests.rkt
+++ b/typed-racket-test/unit-tests/typecheck-tests.rkt
@@ -4015,6 +4015,12 @@
                    (hfun h)))
                (void))
              -Void]
+       ;; check that in-naturals works
+       [tc-e (for/list : (Listof Integer) ([i (in-naturals)]) i) (-lst -Int)]
+       [tc-e (for/list : (Listof Natural) ([s '(foo bar)]
+                                           [i : Natural (in-naturals)])
+               i)
+             (-lst -Nat)]
        )
 
   (test-suite