use find-free-Xs, propagate expected type more

tapl/mlish.rkt

@@ -156,11 +156,11 @@
       [(_ . rst) (lookup x #'rst)]
       [() #f]))
 
-  ;; find-unsolved-Xs : (Stx-Listof Id) Constraints -> (Listof Id)
-  ;; finds the free Xs that aren't constrained by cs
-  (define (find-unsolved-Xs Xs cs)
+  ;; find-free-Xs : (Stx-Listof Id) Type -> (Listof Id)
+  ;; finds the free Xs in the type
+  (define (find-free-Xs Xs ty)
     (for/list ([X (in-list (stx->list Xs))]
-               #:when (not (lookup X cs)))
+               #:when (stx-contains-id? ty X))
       X))
 
   ;; lookup-Xs/keep-unsolved : (Stx-Listof Id) Constraints -> (Listof Type-Stx)
@@ -177,7 +177,7 @@
   ;; and it expands and returns all of the arguments.
   ;; It returns list of 3 values if successful, else throws a type error
   ;;  - a list of all the arguments, expanded
-  ;;  - a list of the the un-constrained type variables
+  ;;  - a list of all the type variables
   ;;  - the constraints for substituting the types
   (define (solve Xs tyXs stx)
     (syntax-parse tyXs
@@ -197,7 +197,7 @@
                          [tyXin (in-list (syntax->list #'(τ_inX ...)))])
                 (define ty_in (inst-type/cs Xs cs tyXin))
                 (define/with-syntax [a- ty_a]
-                  (infer+erase (if (empty? (find-unsolved-Xs Xs cs))
+                  (infer+erase (if (empty? (find-free-Xs Xs ty_in))
                                    (add-expected-ty a ty_in)
                                    a)))
                 (values 
@@ -210,7 +210,7 @@
                                                          (syntax->datum id2))))
                                               #'ty_a))))))
 
-         (list (reverse as-) (find-unsolved-Xs Xs cs) cs)])]))
+         (list (reverse as-) Xs cs)])]))
 
   (define (raise-app-poly-infer-error stx expected-tys given-tys e_fn)
     (type-error #:src stx
@@ -830,10 +830,11 @@
         #'(ext-stlc:#%app e_fn/ty (add-expected e_arg τ_inX) ...)])]
     [else
      ;; ) solve for type variables Xs
-     (define/with-syntax ((e_arg- ...) (unsolved-X ...) cs) (solve #'Xs #'tyX_args stx))
+     (define/with-syntax ((e_arg- ...) Xs* cs) (solve #'Xs #'tyX_args stx))
      ;; ) instantiate polymorphic function type
-     (syntax-parse (inst-types/cs #'Xs #'cs #'tyX_args)
+     (syntax-parse (inst-types/cs #'Xs* #'cs #'tyX_args)
       [(τ_in ... τ_out) ; concrete types
+       #:with (unsolved-X ...) (find-free-Xs #'Xs* #'τ_out)
        ;; ) arity check
        #:fail-unless (stx-length=? #'(τ_in ...) #'e_args)
                      (mk-app-err-msg stx #:expected #'(τ_in ...)
@@ -851,7 +852,9 @@
                           (define new-orig
                             (and old-orig
                                  (substs 
-                                   (stx-map get-orig (lookup-Xs/keep-unsolved #'Xs #'cs)) #'Xs old-orig
+                                   (stx-map get-orig (lookup-Xs/keep-unsolved #'Xs* #'cs))
+                                   #'Xs*
+                                   old-orig
                                    (lambda (x y) 
                                     (equal? (syntax->datum x) (syntax->datum y))))))
                           (set-stx-prop/preserved tyin 'orig (list new-orig)))

tapl/stx-utils.rkt

@@ -1,5 +1,5 @@
 #lang racket/base
-(require syntax/stx racket/list version/utils)
+(require syntax/stx syntax/parse racket/list version/utils)
 (provide (all-defined-out))
 
 (define (stx-cadr stx) (stx-car (stx-cdr stx)))
@@ -70,6 +70,9 @@
 (define (generate-temporariesss stx)
   (stx-map generate-temporariess stx))
 
+;; set-stx-prop/preserved : Stx Any Any -> Stx
+;; Returns a new syntax object with the prop property set to val. If preserved
+;; syntax properties are supported, this also marks the property as preserved.
 (define REQUIRED-VERSION "6.5.0.4")
 (define VERSION (version))
 (define PRESERVED-STX-PROP-SUPPORTED? (version<=? REQUIRED-VERSION VERSION))
@@ -78,6 +81,16 @@
       (syntax-property stx prop val #t)
       (syntax-property stx prop val)))
 
+;; stx-contains-id? : Stx Id -> Boolean
+;; Returns true if stx contains the identifier x, false otherwise.
+(define (stx-contains-id? stx x)
+  (syntax-parse stx
+    [a:id (free-identifier=? #'a x)]
+    [(a . b)
+     (or (stx-contains-id? #'a x)
+         (stx-contains-id? #'b x))]
+    [_ #false]))
+
 ;; based on make-variable-like-transformer from syntax/transformer,
 ;; but using (#%app id ...) instead of ((#%expression id) ...)
 (define (make-variable-like-transformer ref-stx)

tapl/tests/mlish-tests.rkt

@@ -179,6 +179,24 @@
 (check-type (build-list 5 (λ (x) (add1 (add1 x))))
   : (List Int) ⇒ (Cons 6 (Cons 5 (Cons 4 (Cons 3 (Cons 2 Nil))))))
 
+(define (build-list/comp [i : Int] [n : Int] [nf : (→ Int Int)] [f : (→ Int X)] → (List X))
+  (if (= i n)
+      Nil
+      (Cons (f (nf i)) (build-list/comp (add1 i) n nf f))))
+
+(define built-list-1 (build-list/comp 0 3 (λ (x) (* 2 x)) add1))
+(define built-list-2 (build-list/comp 0 3 (λ (x) (* 2 x)) number->string))
+(check-type built-list-1 : (List Int) -> (Cons 1 (Cons 3 (Cons 5 Nil))))
+(check-type built-list-2 : (List String) -> (Cons "0" (Cons "2" (Cons "4" Nil))))
+
+(define (~>2 [a : A] [f : (→ A A)] [g : (→ A B)] → B)
+  (g (f a)))
+
+(define ~>2-result-1 (~>2 1 (λ (x) (* 2 x)) add1))
+(define ~>2-result-2 (~>2 1 (λ (x) (* 2 x)) number->string))
+(check-type ~>2-result-1 : Int -> 3)
+(check-type ~>2-result-2 : String -> "2")
+
 (define (append [lst1 : (List X)] [lst2 : (List X)] → (List X))
   (match lst1 with
    [Nil -> lst2]