Improve type-checker for optional arg functions

original commit: eb7a0f87974f9bdeba8819cfbd7899773eae019b

pkgs/typed-racket-pkgs/typed-racket-lib/typed-racket/typecheck/tc-expr-unit.rkt

@@ -289,12 +289,9 @@
                           (attribute opt.value))
             (opt-convert fun-type required-pos optional-pos)]
            [_ #f]))
-       (match-define (tc-result1: returned-fun-type)
-         (if conv-type
-             (tc-expr/check/type #'fun conv-type)
-             (tc-expr #'fun)))
-       (with-lexical-env/extend (list #'f) (list returned-fun-type)
-         (tc-body/check #'body expected))]
+       (if conv-type
+           (begin (tc-expr/check/type #'fun conv-type) expected)
+           (tc-expr (remove-ascription form)))]
       ;; let
       [(let-values ([(name ...) expr] ...) . body)
        (tc/let-values #'((name ...) ...) #'(expr ...) #'body form expected)]
@@ -388,6 +385,12 @@
                           (syntax->list #'(formals ...))
                           (syntax->datum #'(mand-kw ...))
                           (syntax->datum #'(all-kw ...))))]
+      ;; opt function def
+      [(~and opt:opt-lambda^
+             (let-values ([(f) fun])
+               (case-lambda (formals . cl-body) ...)))
+       (ret (opt-unconvert (tc-expr/t #'fun)
+                           (syntax->list #'(formals ...))))]
       ;; let
       [(let-values ([(name ...) expr] ...) . body)
        (tc/let-values #'((name ...) ...) #'(expr ...) #'body form)]

pkgs/typed-racket-pkgs/typed-racket-lib/typed-racket/types/kw-types.rkt

@@ -200,13 +200,7 @@
                    (if (equal? (last opt-and-rest-args) Univ)
                        Univ
                        -Bottom)))
-            (define opt-types
-              (let loop ([opt-args opt-and-rest-args]
-                         [opt-types '()])
-                (if (or (null? opt-args)
-                        (null? (cdr opt-args)))
-                    (reverse opt-types)
-                    (loop (cddr opt-args) (cons (car opt-args) opt-types)))))
+            (define opt-types (take opt-and-rest-args opt-non-kw-argc))
             (make-Function
              (for/list ([to-take (in-range (add1 (length opt-types)))])
                (make-arr* (append mand-args (take opt-types to-take))
@@ -275,9 +269,54 @@
          (make-PolyDots names (loop f))]
         [t t]))))
 
+;; opt-unconvert : Type (Listof Syntax) -> Type
+;; Given a type for a core optional arg function, unconvert it to a
+;; normal function type. See `kw-unconvert` above.
+(define (opt-unconvert ft formalss)
+  (define (lengthish formals)
+    (define lst (syntax->list formals))
+    (if lst (length lst) +inf.0))
+  (define max-formals (argmax lengthish formalss))
+  (define min-formals (argmin lengthish formalss))
+  (define-values (raw-argc rest?)
+    (syntax-parse max-formals
+      [(arg:id ...)
+       (values (length (syntax->list #'(arg ...))) #f)]
+      [(arg:id ... . rst:id)
+       (values (length (syntax->list #'(arg ...))) #t)]))
+  (define opt-argc
+    (syntax-parse min-formals
+      [(arg:id ...)
+       (- raw-argc (length (syntax->list #'(arg ...))))]
+      ;; if min and max both have rest args, then there cannot
+      ;; have been any optional arguments
+      [(arg:id ... . rst:id) 0]))
+  ;; counted twice since optionals expand to two arguments
+  (define argc (+ raw-argc opt-argc))
+  (define mand-argc (- argc (* 2 opt-argc)))
+  (match ft
+    [(Function: arrs)
+     (cond [(and (even? (length arrs)) (>= (length arrs) 2))
+            (match-define (arr: doms rng _ _ _) (car arrs))
+            (define-values (mand-args opt-and-rest-args)
+              (split-at doms mand-argc))
+            (define rest-type
+              (and rest?
+                   (if (equal? (last opt-and-rest-args) Univ)
+                       Univ
+                       -Bottom)))
+            (define opt-types (take opt-and-rest-args opt-argc))
+            (make-Function
+             (for/list ([to-take (in-range (add1 (length opt-types)))])
+               (make-arr* (append mand-args (take opt-types to-take))
+                          rng
+                          #:rest rest-type)))]
+           [else (int-err "unsupported arrs in keyword function type")])]
+    [_ (int-err "unsupported keyword function type")]))
+
 ;; partition-kws : (Listof Keyword) -> (values (Listof Keyword) (Listof Keyword))
 ;; Partition keywords by whether they are mandatory or not
 (define (partition-kws kws)
   (partition (match-lambda [(Keyword: _ _ mand?) mand?]) kws))
 
-(provide kw-convert kw-unconvert opt-convert partition-kws)
+(provide kw-convert kw-unconvert opt-convert opt-unconvert partition-kws)

pkgs/typed-racket-pkgs/typed-racket-test/tests/typed-racket/unit-tests/typecheck-tests.rkt

@@ -2075,10 +2075,15 @@
                #:msg "expected: Symbol.*given: String"]
        [tc-err (tr:lambda (x [y : String "a"] z) (string-append y "b"))
                #:msg "expected optional lambda argument"]
-       #| FIXME: requires improvement in opt-lambda checker
        [tc-e (tr:lambda (x [y : String "a"]) (string-append y "b"))
              (->opt Univ [-String] -String)]
-       |#
+       [tc-e (tr:lambda (x y [z : String "a"]) (string-append z "b"))
+             (->opt Univ Univ [-String] -String)]
+       [tc-e (tr:lambda (w x [y : String "y"] [z : String "z"]) (string-append y z))
+             (->opt Univ Univ [-String -String] -String)]
+       [tc-e (tr:lambda (w [x : String] [y : String "y"] [z : String "z"])
+               (string-append x z))
+             (->opt Univ -String [-String -String] -String)]
        [tc-e (tr:lambda (x #:y [y : String]) (string-append y "b"))
              (->key Univ #:y -String #t -String)]
        [tc-e (tr:lambda (x #:y [y : String]) : String (string-append y "b"))