tweaked sandbox generated code

src/Node/CheckSecurityNode.php

@@ -45,12 +45,13 @@ class CheckSecurityNode extends Node
         }
 
         $compiler
+            ->write("\$this->sandbox = \$this->env->getExtension('\Twig\Extension\SandboxExtension');\n")
             ->write('$tags = ')->repr(array_filter($tags))->raw(";\n")
             ->write('$filters = ')->repr(array_filter($filters))->raw(";\n")
             ->write('$functions = ')->repr(array_filter($functions))->raw(";\n\n")
             ->write("try {\n")
             ->indent()
-            ->write("\$this->env->getExtension('\Twig\Extension\SandboxExtension')->checkSecurity(\n")
+            ->write("\$this->sandbox->checkSecurity(\n")
             ->indent()
             ->write(!$tags ? "[],\n" : "['".implode("', '", array_keys($tags))."'],\n")
             ->write(!$filters ? "[],\n" : "['".implode("', '", array_keys($filters))."'],\n")

src/Node/SandboxNode.php

@@ -29,16 +29,15 @@ class SandboxNode extends Node
     {
         $compiler
             ->addDebugInfo($this)
-            ->write("\$sandbox = \$this->env->getExtension('\Twig\Extension\SandboxExtension');\n")
-            ->write("if (!\$alreadySandboxed = \$sandbox->isSandboxed()) {\n")
+            ->write("if (!\$alreadySandboxed = \$this->sandbox->isSandboxed()) {\n")
             ->indent()
-            ->write("\$sandbox->enableSandbox();\n")
+            ->write("\$this->sandbox->enableSandbox();\n")
             ->outdent()
             ->write("}\n")
             ->subcompile($this->getNode('body'))
             ->write("if (!\$alreadySandboxed) {\n")
             ->indent()
-            ->write("\$sandbox->disableSandbox();\n")
+            ->write("\$this->sandbox->disableSandbox();\n")
             ->outdent()
             ->write("}\n")
         ;

src/NodeVisitor/SandboxNodeVisitor.php

@@ -77,7 +77,7 @@ class SandboxNodeVisitor extends AbstractNodeVisitor
         if ($node instanceof ModuleNode) {
             $this->inAModule = false;
 
-            $node->setNode('display_start', new Node([new CheckSecurityNode($this->filters, $this->tags, $this->functions), $node->getNode('display_start')]));
+            $node->setNode('constructor_end', new Node([new CheckSecurityNode($this->filters, $this->tags, $this->functions), $node->getNode('display_start')]));
         }
 
         return $node;

src/Template.php

@@ -39,6 +39,7 @@ abstract class Template implements \Twig_TemplateInterface
     protected $env;
     protected $blocks = [];
     protected $traits = [];
+    protected $sandbox;
 
     public function __construct(Environment $env)
     {

test/Twig/Tests/Node/SandboxTest.php

@@ -32,13 +32,12 @@ class Twig_Tests_Node_SandboxTest extends NodeTestCase
 
         $tests[] = [$node, <<<EOF
 // line 1
-\$sandbox = \$this->env->getExtension('\Twig\Extension\SandboxExtension');
-if (!\$alreadySandboxed = \$sandbox->isSandboxed()) {
-    \$sandbox->enableSandbox();
+if (!\$alreadySandboxed = \$this->sandbox->isSandboxed()) {
+    \$this->sandbox->enableSandbox();
 }
 echo "foo";
 if (!\$alreadySandboxed) {
-    \$sandbox->disableSandbox();
+    \$this->sandbox->disableSandbox();
 }
 EOF
         ];