From da574a47d06a8c8f2ed4d88b9c916b1b7075b3ef Mon Sep 17 00:00:00 2001
From: Sam Tobin-Hochstadt <samth@cs.indiana.edu>
Date: Sun, 8 Nov 2015 14:27:08 -0500
Subject: [PATCH] Handle `in-vector` with range arguments.

* Add `normalise-inputs` to special function env.
* Treat eta-expansion specially. Now
    `(lambda (x ...) (f x ...))`
  will typecheck like `f` but with a type restricted to
  the size of `x ...`.

  Currently, this special case only works for non-polymorphic
  functions.
---
 .../base-env/base-special-env.rkt             |  4 ++
 .../typed-racket/typecheck/tc-lambda-unit.rkt | 53 ++++++++++++++++---
 typed-racket-test/succeed/in-vector-range.rkt |  7 +++
 3 files changed, 57 insertions(+), 7 deletions(-)
 create mode 100644 typed-racket-test/succeed/in-vector-range.rkt

diff --git a/typed-racket-lib/typed-racket/base-env/base-special-env.rkt b/typed-racket-lib/typed-racket/base-env/base-special-env.rkt
index 032984e9..18e5280a 100644
--- a/typed-racket-lib/typed-racket/base-env/base-special-env.rkt
+++ b/typed-racket-lib/typed-racket/base-env/base-special-env.rkt
@@ -46,6 +46,10 @@
          (cl->*
           (-> (-lst a) -Null (-lst a))
           (-> (-lst a) (-lst b) (-lst (Un a b)))))]
+  ;; normalise-inputs
+  [(make-template-identifier 'normalise-inputs 'racket/private/for)
+   (-poly (a)
+          (-> -Symbol -String (-> a -Boolean) (-> a -Nat) a -Nat (Un (-val #f) -Nat) -Nat (-values (list a -Index -Index -Index))))]
   ;; make-sequence
   [(make-template-identifier 'make-sequence 'racket/private/for)
    (-poly (a b)
diff --git a/typed-racket-lib/typed-racket/typecheck/tc-lambda-unit.rkt b/typed-racket-lib/typed-racket/typecheck/tc-lambda-unit.rkt
index 1696cef3..c2d4cfbe 100644
--- a/typed-racket-lib/typed-racket/typecheck/tc-lambda-unit.rkt
+++ b/typed-racket-lib/typed-racket/typecheck/tc-lambda-unit.rkt
@@ -185,7 +185,16 @@
   (for/list ([arg-types (in-list new-arg-types)])
     (tc-lambda-body arg-list arg-types body)))
 
-
+;; restrict-to-arity : arr? nat -> (or/c #f arr?)
+;; either produces a new arity which is a subtype of arr with arity n,
+;; or #f is that is not possible
+(define (restrict-to-arity arr n)
+  (match arr
+    ;; currently does not handle rest arguments
+    [(arr: args ret #f #f '())
+     #:when (= n (length args))
+     arr]
+    [_ #f]))
 
 ;; formals syntax -> listof[arr?]
 (define (tc/lambda-clause formals body)
@@ -196,11 +205,21 @@
 
   (define arg-list (formals-positional formals))
   (define rest-id (formals-rest formals))
+
+  (define eta-expanded?
+    (syntax-parse body
+      [((~literal #%plain-app) fun:expr j:id ...)
+       #:when (equal? (length arg-list)
+                      (length (syntax->list #'(j ...))))
+       #:when (andmap free-identifier=? arg-list (syntax->list #'(j ...)))
+       #'fun]
+      [_ #f]))
+  
   (cond
     [(and (> (dict-count aux-table) 0) (not rest-id))
      (tc/opt-lambda-clause arg-list body aux-table flag-table)]
     [else
-     (define arg-types (get-types arg-list #:default Univ))
+     (define arg-types (get-types arg-list #:default (lambda () #f)))
      (define rest-type
        (cond
          ;; Lambda with poly dotted rest argument
@@ -217,11 +236,31 @@
           (get-type rest-id #:default Univ)]
          ;; Lambda with no rest argument
          [else #f]))
-
-     (list
-       (tc-lambda-body arg-list arg-types
-         #:rest (and rest-type (list rest-id rest-type))
-         body))]))
+     (cond 
+      ;; special case for un-annotated eta-expansions
+      [(and eta-expanded? (not rest-id) (andmap not arg-types)
+            ;; FIXME: should also handle polymorphic types
+            ;; but we can't return anything but a (listof arr?) here 
+            ;; FIXME: sometimes will typecheck the relevant code twice if it doesn't match
+            ;; FIXME: misses optimization opportunities of this code
+            (match (tc-expr eta-expanded?)
+              [(tc-result1: (Function: arrs))
+               (define possibles (for*/list ([arr (in-list arrs)]
+                                             [restricted (in-value (restrict-to-arity arr (length arg-list)))]
+                                             #:when restricted)
+                                   restricted))
+               (if (null? possibles)
+                   #f
+                   possibles)]))
+       =>
+       (lambda (x)
+         (register-ignored! (car (syntax-e body)))
+         x)]
+      [else
+       (list
+        (tc-lambda-body arg-list (map (lambda (v) (or v Univ)) arg-types)
+                        #:rest (and rest-type (list rest-id rest-type))
+                        body))])]))
 
 
 ;; positional: natural? - the number of positional arguments
diff --git a/typed-racket-test/succeed/in-vector-range.rkt b/typed-racket-test/succeed/in-vector-range.rkt
new file mode 100644
index 00000000..ad23b77e
--- /dev/null
+++ b/typed-racket-test/succeed/in-vector-range.rkt
@@ -0,0 +1,7 @@
+
+#lang typed/racket/base
+
+;(: v : (Vectorof Integer))
+(define v (vector 1 2 3))
+(for ([t (in-vector v 0 2)])
+    (displayln t))