From 8cd9bcb775657e49da54964114fdaed10b3e75cb Mon Sep 17 00:00:00 2001
From: Fabien Potencier <fabien.potencier@gmail.com>
Date: Mon, 13 Dec 2010 15:46:09 +0100
Subject: [PATCH] enhanced the for-loop optimizer pass 2 (patch from nikic)

---
 lib/Twig/NodeVisitor/Optimizer.php            |   52 +++++++++++++++++++------
 test/Twig/Tests/NodeVisitor/OptimizerTest.php |    9 ++--
 2 files changed, 45 insertions(+), 16 deletions(-)

diff --git a/lib/Twig/NodeVisitor/Optimizer.php b/lib/Twig/NodeVisitor/Optimizer.php
index 005dc51..151ba7b 100644
--- a/lib/Twig/NodeVisitor/Optimizer.php
+++ b/lib/Twig/NodeVisitor/Optimizer.php
@@ -86,19 +86,35 @@ class Twig_NodeVisitor_Optimizer implements Twig_NodeVisitorInterface
             // disable the loop variable by default
             $node->setAttribute('with_loop', false);
             array_unshift($this->loops, $node);
-        } elseif ($this->loops
-            // when do we need to add the loop variable back?
-            && (
-                // the loop variable is referenced here
-                ($node instanceof Twig_Node_Expression_Name && 'loop' === $node->getAttribute('name'))
-                ||
-                // include without the only attribute
-                ($this->loops && $node instanceof Twig_Node_Include && !$node->getAttribute('only'))
-            )
+        } elseif (!$this->loops) {
+            // we are outside a loop
+            return;
+        }
+
+        // when do we need to add the loop variable back?
+
+        // the loop variable is referenced for the current loop
+        elseif ($node instanceof Twig_Node_Expression_Name && 'loop' === $node->getAttribute('name')) {
+            $this->addLoopToCurrent();
+        }
+
+        // include without the only attribute
+        elseif ($node instanceof Twig_Node_Include && !$node->getAttribute('only')) {
+            $this->addLoopToAll();
+        }
+
+        // the loop variable is referenced via an attribute
+        elseif ($node instanceof Twig_Node_Expression_GetAttr
+            && (!$node->getNode('attribute') instanceof Twig_Node_Expression_Constant
+                || 'parent' === $node->getNode('attribute')->getAttribute('value')
+               )
+            && (true === $this->loops[0]->getAttribute('with_loop')
+                || ($node->getNode('node') instanceof Twig_Node_Expression_Name
+                    && 'loop' === $node->getNode('node')->getAttribute('name')
+                   )
+               )
         ) {
-            foreach ($this->loops as $loop) {
-                $loop->setAttribute('with_loop', true);
-            }
+            $this->addLoopToAll();
         }
     }
 
@@ -109,6 +125,18 @@ class Twig_NodeVisitor_Optimizer implements Twig_NodeVisitorInterface
         }
     }
 
+    protected function addLoopToCurrent()
+    {
+        $this->loops[0]->setAttribute('with_loop', true);
+    }
+
+    protected function addLoopToAll()
+    {
+        foreach ($this->loops as $loop) {
+            $loop->setAttribute('with_loop', true);
+        }
+    }
+
     /**
      * {@inheritdoc}
      */
diff --git a/test/Twig/Tests/NodeVisitor/OptimizerTest.php b/test/Twig/Tests/NodeVisitor/OptimizerTest.php
index 1501b57..96ccfc7 100644
--- a/test/Twig/Tests/NodeVisitor/OptimizerTest.php
+++ b/test/Twig/Tests/NodeVisitor/OptimizerTest.php
@@ -21,7 +21,7 @@ class Twig_Tests_NodeVisitor_OptimizerTest extends PHPUnit_Framework_TestCase
         $stream = $env->parse($env->tokenize($template, 'index'));
 
         foreach ($expected as $target => $withLoop) {
-            $this->assertTrue($this->checkForConfiguration($stream, $target, $withLoop));
+            $this->assertTrue($this->checkForConfiguration($stream, $target, $withLoop), sprintf('variable %s is %soptimized', $target, $withLoop ? 'not ' : ''));
         }
     }
 
@@ -38,12 +38,13 @@ class Twig_Tests_NodeVisitor_OptimizerTest extends PHPUnit_Framework_TestCase
 
             array('{% for i in foo %}{% include "foo" only %}{% endfor %}', array('i' => false)),
 
-            // "i" should be false
-            array('{% for i in foo %}{% for j in foo %}{{ loop.index }}{% endfor %}{% endfor %}', array('i' => true, 'j' => true)),
+            array('{% for i in foo %}{% for j in foo %}{{ loop.index }}{% endfor %}{% endfor %}', array('i' => false, 'j' => true)),
 
             array('{% for i in foo %}{% for j in foo %}{{ loop.parent.loop.index }}{% endfor %}{% endfor %}', array('i' => true, 'j' => true)),
 
-            //array('{% for i in foo %}{% for j in foo %}{{ foo.parent.loop.index }}{% endfor %}{% endfor %}', array('i' => false, 'j' => false)),
+            array('{% for i in foo %}{% set l = loop %}{% for j in foo %}{{ l.index }}{% endfor %}{% endfor %}', array('i' => true, 'j' => false)),
+
+            array('{% for i in foo %}{% for j in foo %}{{ foo.parent.loop.index }}{% endfor %}{% endfor %}', array('i' => false, 'j' => false)),
 
             array('{% for i in foo %}{% for j in foo %}{{ loop["parent"].loop.index }}{% endfor %}{% endfor %}', array('i' => true, 'j' => true)),
         );
-- 
1.7.2.5