Extended unboxing of let-bound functions to support let loops.

collects/tests/typed-scheme/optimizer/generic/unboxed-let-functions6.rkt

@@ -0,0 +1,10 @@
+#lang typed/scheme #:optimize
+
+(require racket/unsafe/ops racket/flonum)
+
+(let: loop : Inexact-Complex ((z : Inexact-Complex 0.0+0.0i)
+                              (l : (Listof Integer) '(1 2 3)))
+      (if (null? l)
+          (+ z 0.0+1.0i)
+          (loop (+ z (car l))
+                (cdr l))))

collects/tests/typed-scheme/optimizer/generic/unboxed-let-functions7.rkt

@@ -0,0 +1,10 @@
+#lang typed/scheme #:optimize
+
+(require racket/unsafe/ops racket/flonum)
+
+(let: loop : Inexact-Complex ((z : Inexact-Complex 0.0+0.0i)
+                              (l : (Listof Integer) '(1 2 3)))
+  (if (null? l)
+      z ; boxed use. z should be unboxed anyway
+      (loop (+ z (car l))
+            (cdr l))))

collects/tests/typed-scheme/optimizer/hand-optimized/unboxed-let-functions6.rkt

@@ -0,0 +1,25 @@
+#lang racket
+
+(require racket/unsafe/ops racket/flonum)
+
+(let*-values (((unboxed-gensym-1) 0.0+0.0i)
+              ((unboxed-gensym-2) (unsafe-flreal-part unboxed-gensym-1))
+              ((unboxed-gensym-3) (unsafe-flimag-part unboxed-gensym-1)))
+  ((letrec-values
+    (((loop)
+      (lambda (unboxed-real-1 unboxed-imag-2 l)
+        (if (null? l)
+            (let*-values (((unboxed-gensym-3) 0.0+1.0i)
+                          ((unboxed-gensym-4) (unsafe-flreal-part unboxed-gensym-3))
+                          ((unboxed-gensym-5) (unsafe-flimag-part unboxed-gensym-3))
+                          ((unboxed-gensym-6) (unsafe-fl+ unboxed-real-1 unboxed-gensym-4))
+                          ((unboxed-gensym-7) (unsafe-fl+ unboxed-imag-2 unboxed-gensym-5)))
+              (unsafe-make-flrectangular unboxed-gensym-6 unboxed-gensym-7))
+            (let*-values (((unboxed-gensym-1) (->fl (unsafe-car l)))
+                          ((unboxed-gensym-2) (unsafe-fl+ unboxed-real-1 unboxed-gensym-1))
+                          ((unboxed-gensym-3) unboxed-imag-2))
+              (loop unboxed-gensym-2 unboxed-gensym-3
+                    (unsafe-cdr l)))))))
+    loop)
+   unboxed-gensym-2 unboxed-gensym-3 '(1 2 3)))
+(void)

collects/tests/typed-scheme/optimizer/hand-optimized/unboxed-let-functions7.rkt

@@ -0,0 +1,20 @@
+#lang racket
+
+(require racket/unsafe/ops racket/flonum)
+
+(let*-values (((unboxed-gensym-1) 0.0+0.0i)
+              ((unboxed-gensym-2) (unsafe-flreal-part unboxed-gensym-1))
+              ((unboxed-gensym-3) (unsafe-flimag-part unboxed-gensym-1)))
+  ((letrec-values
+    (((loop)
+      (lambda (unboxed-real-1 unboxed-imag-2 l)
+        (if (null? l)
+            (unsafe-make-flrectangular unboxed-real-1 unboxed-imag-2)
+            (let*-values (((unboxed-gensym-3) (->fl (unsafe-car l)))
+                          ((unboxed-gensym-4) (unsafe-fl+ unboxed-real-1 unboxed-gensym-3))
+                          ((unboxed-gensym-5) unboxed-imag-2))
+              (loop unboxed-gensym-4 unboxed-gensym-5
+                    (unsafe-cdr l)))))))
+    loop)
+   unboxed-gensym-2 unboxed-gensym-3 '(1 2 3)))
+(void)

collects/tests/typed-scheme/optimizer/non-optimized/unboxed-let-functions6.rkt

@@ -0,0 +1,10 @@
+#lang typed/scheme
+
+(require racket/unsafe/ops racket/flonum)
+
+(let: loop : Inexact-Complex ((z : Inexact-Complex 0.0+0.0i)
+                              (l : (Listof Integer) '(1 2 3)))
+      (if (null? l)
+          (+ z 0.0+1.0i)
+          (loop (+ z (car l))
+                (cdr l))))

collects/tests/typed-scheme/optimizer/non-optimized/unboxed-let-functions7.rkt

@@ -0,0 +1,10 @@
+#lang typed/scheme
+
+(require racket/unsafe/ops racket/flonum)
+
+(let: loop : Inexact-Complex ((z : Inexact-Complex 0.0+0.0i)
+                              (l : (Listof Integer) '(1 2 3)))
+  (if (null? l)
+      z ; boxed use. z should be unboxed anyway
+      (loop (+ z (car l))
+            (cdr l))))

collects/typed-scheme/optimizer/unboxed-let.rkt

@@ -13,8 +13,37 @@
 
 ;; possibly replace bindings of complex numbers by bindings of their 2 components
 ;; useful for intermediate results used more than once and for loop variables
-
 (define-syntax-class unboxed-let-opt-expr
+  (pattern e:app-of-unboxed-let-opt-expr
+           #:with opt #'e.opt)
+  (pattern (~var e (unboxed-let-opt-expr-internal #f))
+           #:with opt #'e.opt))
+
+;; let loops expand to an application of a letrec-values
+;; thus, the loop function technically escapes from the letrec, but it
+;; escapes in the operator position of a call site we control (here)
+;; we can extend unboxing
+(define-syntax-class app-of-unboxed-let-opt-expr
+  #:literal-sets (kernel-literals)
+  (pattern (~and e ((~literal #%plain-app)
+                    (~and let-e
+                          ((~literal letrec-values)
+                           bindings
+                           loop-fun:id)) ; sole element of the body
+                    args:expr ...))
+           #:with (~var operator (unboxed-let-opt-expr-internal #t)) #'let-e
+           #:with unboxed-info (dict-ref unboxed-funs-table #'loop-fun #f)
+           #:when (syntax->datum #'unboxed-info)
+           #:with (~var e* (inexact-complex-call-site-opt-expr
+                            #'unboxed-info #'operator.opt))
+           #'e
+           #:with opt
+           #'e*.opt))
+
+;; does the bulk of the work
+;; detects which let bindings can be unboxed, same for arguments of let-bound
+;; functions
+(define-syntax-class (unboxed-let-opt-expr-internal let-loop?)
   #:literal-sets (kernel-literals)
   (pattern (~and exp (letk:let-like-keyword
                       (clause:expr ...) body:expr ...))
@@ -40,8 +69,8 @@
                     (let ((fun-name (car (syntax-e (car p)))))
                       (and
                        ;; if the function escapes, we can't change it's interface
-                       (and (not (is-var-mutated? fun-name))
+                       (not (is-var-mutated? fun-name))
-                            (not (escapes? fun-name #'(begin body ...))))
+                       (not (escapes? fun-name #'(begin body ...) let-loop?))
                        (match (type-of (cadr p)) ; rhs, we want a lambda
                          [(tc-result1: (Function: (list (arr: doms rngs
                                                               (and rests #f)
@@ -158,7 +187,10 @@
 ;; very simple escape analysis for functions
 ;; if a function is ever used in a non-operator position, we consider it escapes
 ;; if it doesn't escape, we may be able to pass its inexact complex args unboxed
-(define (escapes? v exp)
+;; if we are in a let loop, don't consider functions that escape by being the
+;; sole thing in the let's body as escaping, since they would only escape to
+;; a call site that we control, which is fine
+(define (escapes? v exp let-loop?)
 
   (define (look-at exp)
     (or (direct-child-of? v exp)
@@ -193,7 +225,17 @@
 
       ;; does not escape
       [_ #f]))
-  (rec exp))
+
+  ;; if the given var is the _only_ element of the body and we're in a
+  ;; let loop, we let it slide
+  (and (not (and let-loop?
+                 (syntax-parse exp
+                   #:literal-sets (kernel-literals)
+                   ;; the body gets wrapped in a begin before it's sent here
+                   [(begin i:identifier)
+                    (free-identifier=? #'i v)]
+                   [_ #f])))
+       (rec exp)))
 
 ;; let clause whose rhs is going to be unboxed (turned into multiple bindings)
 (define-syntax-class unboxed-let-clause