From e4b0309f979d5ed51085f027b834ec3309a90245 Mon Sep 17 00:00:00 2001
From: Jeff Bezanson <jeff.bezanson@gmail.com>
Date: Mon, 16 May 2016 11:37:30 -0400
Subject: [PATCH] improvements to lambda-optimize-vars

This fixes a bug plus optimizes more variables. Should help #16047.
---
 src/julia-syntax.scm | 35 +++++++++++++++++++++++++----------
 test/core.jl         |  7 +++++++
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/src/julia-syntax.scm b/src/julia-syntax.scm
index aea36c9c0c85b..ce53ae5531846 100644
--- a/src/julia-syntax.scm
+++ b/src/julia-syntax.scm
@@ -2620,11 +2620,25 @@ f(x) = yt(x)
                          '())))
                  args)))))
 
+(define (take-statements-while pred body)
+  (let ((acc '()))
+    (define (take expr)
+      ;; returns #t as long as exprs match and we should continue
+      (cond ((and (pair? expr) (memq (car expr) '(block body)))
+             (let loop ((xs (cdr expr)))
+               (cond ((null? xs) #t)
+                     ((take (car xs)) (loop (cdr xs)))
+                     (else #f))))
+            ((pred expr)
+             (set! acc (cons expr acc))
+             #t)
+            (else #f)))
+    (take body)
+    (reverse! acc)))
+
 ;; clear capture bit for vars assigned once at the top, to avoid allocating
 ;; some unnecessary Boxes.
 (define (lambda-optimize-vars! lam)
-  ;; flattening blocks helps us find more dominating statements
-  (set-car! (cdddr lam) (flatten-blocks (lam:body lam)))
   (define (expr-uses-var ex v)
     (cond ((assignment? ex) (expr-contains-eq v (caddr ex)))
           ((eq? (car ex) 'method)
@@ -2638,15 +2652,16 @@ f(x) = yt(x)
         (let* ((leading
                 (filter (lambda (x) (and (pair? x)
                                          (or (and (eq? (car x) 'method)
-                                                  (length> (car x) 2))
+                                                  (length> x 2))
                                              (eq? (car x) '=))))
-                        (take-while (lambda (e)
-                                      (or (atom? e)
-                                          (memq (car e) '(quote top core line inert local
-                                                          meta inbounds boundscheck simdloop
-                                                          implicit-global global globalref
-                                                          const newvar = null method))))
-                                    (lam:body lam))))
+                        (take-statements-while
+                         (lambda (e)
+                           (or (atom? e)
+                               (memq (car e) '(quote top core line inert local
+                                               meta inbounds boundscheck simdloop
+                                               implicit-global global globalref
+                                               const newvar = null method))))
+                         (lam:body lam))))
                (unused (map cadr leading))
                (def (table)))
           ;; TODO: reorder leading statements to put assignments where the RHS is
diff --git a/test/core.jl b/test/core.jl
index eea3f903d608a..bf605b0182f72 100644
--- a/test/core.jl
+++ b/test/core.jl
@@ -437,6 +437,13 @@ Yc(f) = (h->f(x->h(h)(x)))(h->f(x->h(h)(x)))
 yfib = Yc(fib->(n->(n < 2 ? n : fib(n-1) + fib(n-2))))
 @test yfib(20) == 6765
 
+function capt_before_def()
+    f() = y
+    y = 2
+    f
+end
+@test capt_before_def()() == 2
+
 # variable scope, globals
 glob_x = 23
 function glotest()