From 6fa36e853a62d6691d5041f3170098d06d1b732a Mon Sep 17 00:00:00 2001
From: Joshua Behrens <behrens@heptacom.de>
Date: Tue, 5 Dec 2023 03:04:27 +0100
Subject: [PATCH] Allow union type predicted type access

---
 src/Node/Expression/GetAttrExpression.php | 344 +++++++++++++++++-----
 tests/Node/Expression/GetAttrTest.php     |  28 +-
 2 files changed, 292 insertions(+), 80 deletions(-)

diff --git a/src/Node/Expression/GetAttrExpression.php b/src/Node/Expression/GetAttrExpression.php
index 49fe0c84f..dd9c5690f 100644
--- a/src/Node/Expression/GetAttrExpression.php
+++ b/src/Node/Expression/GetAttrExpression.php
@@ -13,10 +13,13 @@
 namespace Twig\Node\Expression;
 
 use Twig\Compiler;
+use Twig\Environment;
 use Twig\Extension\SandboxExtension;
 use Twig\Template;
 use Twig\TypeHint\ArrayType;
 use Twig\TypeHint\ObjectType;
+use Twig\TypeHint\TypeInterface;
+use Twig\TypeHint\UnionType;
 
 class GetAttrExpression extends AbstractExpression
 {
@@ -41,63 +44,23 @@ public function compile(Compiler $compiler): void
                 $type = $this->getNode('node')->getAttribute('typeHint');
             }
 
-            if ($type instanceof ArrayType) {
-                $compiler
-                    ->raw('((')
-                    ->subcompile($this->getNode('node'))
-                    ->raw(')[')
-                    ->subcompile($this->getNode('attribute'))
-                    ->raw('] ?? null)')
-                ;
+            if ($type instanceof TypeInterface) {
+                $sourceCompiler = $this->createNodeSourceCompiler();
+                $accessCompiler = $this->createAccessCompiler($type, $env);
 
-                return;
-            } else if ($type instanceof ObjectType && $this->getNode('attribute') instanceof ConstantExpression) {
-                $attributeName = $this->getNode('attribute')->getAttribute('value');
-
-                if ($type->getPropertyType($attributeName) !== null) {
-                    $compiler
-                        ->raw('((')
-                        ->subcompile($this->getNode('node'))
-                        ->raw(')?->')
-                        ->raw($attributeName)
-                        ->raw(')')
-                    ;
-
-                    return;
+                if (true || $accessCompiler['condition'] === null) {
+                    $accessCompiler['accessor']($compiler, $sourceCompiler);
+                } else {
+                    $compiler->raw('(');
+                    $accessCompiler['condition']($compiler, $sourceCompiler);
+                    $compiler->raw(' ? ');
+                    $accessCompiler['accessor']($compiler, $sourceCompiler);
+                    $compiler->raw(' : ');
+                    $this->createGuessingAccessCompiler($env->hasExtension(SandboxExtension::class))['accessor']($compiler, $sourceCompiler);
+                    $compiler->raw(')');
                 }
 
-                /** Keep similar to @see \Twig\TypeHint\ObjectType::getAttributeType */
-                $methodNames = [
-                    $attributeName,
-                    'get' . $attributeName,
-                    'is' . $attributeName,
-                    'has' . $attributeName,
-                ];
-
-                foreach ($methodNames as $methodName) {
-                    if ($type->getMethodType($methodName) !== null) {
-                        $compiler
-                            ->raw('((')
-                            ->subcompile($this->getNode('node'))
-                            ->raw(')?->')
-                            ->raw($methodName)
-                            ->raw('(')
-                        ;
-
-                        if ($this->hasNode('arguments') && $this->getNode('arguments') instanceof ArrayExpression && $this->getNode('arguments')->count() > 0) {
-                            for ($argIndex = 0; $argIndex < $this->getNode('arguments')->count(); $argIndex += 2) {
-                                if ($argIndex > 0) {
-                                    $compiler->raw(', ');
-                                }
-
-                                $compiler->subcompile($this->getNode('arguments')->getNode($argIndex + 1));
-                            }
-                        }
-
-                        $compiler->raw('))');
-                        return;
-                    }
-                }
+                return;
             }
         }
 
@@ -126,31 +89,264 @@ public function compile(Compiler $compiler): void
             return;
         }
 
-        $compiler->raw('twig_get_attribute($this->env, $this->source, ');
+        $this->createGuessingAccessCompiler($env->hasExtension(SandboxExtension::class))['accessor']($compiler, $this->createNodeSourceCompiler());
+    }
+
+    /**
+     * @return array{
+     *     condition: \Closure(Compiler, \Closure(Compiler): void): void|null,
+     *     accessor: \Closure(Compiler, \Closure(Compiler): void): void
+     * }
+     */
+    private function createAccessCompiler(TypeInterface $type, Environment $env): array
+    {
+        if ($type instanceof UnionType) {
+            return $this->createUnionAccessCompiler($type, $env);
+        }
+
+        if ($type instanceof ArrayType) {
+            return $this->createArrayAccessCompiler();
+        }
+
+        if ($type instanceof ObjectType && $this->getNode('attribute') instanceof ConstantExpression) {
+            $attributeName = $this->getNode('attribute')->getAttribute('value');
+
+            if ($type->getPropertyType($attributeName) !== null) {
+                return $this->createObjectPropertyAccessCompiler($type, $attributeName);
+            }
+
+            /** Keep similar to @see \Twig\TypeHint\ObjectType::getAttributeType */
+            $methodNames = [
+                $attributeName,
+                'get' . $attributeName,
+                'is' . $attributeName,
+                'has' . $attributeName,
+            ];
+
+            foreach ($methodNames as $methodName) {
+                if ($type->getMethodType($methodName) === null) {
+                    continue;
+                }
 
-        if ($this->getAttribute('ignore_strict_check')) {
-            $this->getNode('node')->setAttribute('ignore_strict_check', true);
+                return $this->createObjectMethodAccessCompiler($type, $methodName);
+            }
         }
 
-        $compiler
-            ->subcompile($this->getNode('node'))
-            ->raw(', ')
-            ->subcompile($this->getNode('attribute'))
-        ;
+        return $this->createGuessingAccessCompiler($env->hasExtension(SandboxExtension::class));
+    }
+
+    /**
+     * @return array{
+     *     condition: null,
+     *     accessor: \Closure(Compiler, \Closure(Compiler): void): void
+     * }
+     */
+    private function createGuessingAccessCompiler(bool $isSandboxed): array
+    {
+        return [
+            'condition' => null,
+            'accessor' => function (Compiler $compiler, \Closure $sourceCompiler) use ($isSandboxed): void {
+                $compiler->raw('twig_get_attribute($this->env, $this->source, ');
+
+                if ($this->getAttribute('ignore_strict_check')) {
+                    $this->getNode('node')->setAttribute('ignore_strict_check', true);
+                }
+
+                $sourceCompiler($compiler);
+
+                $compiler
+                    ->raw(', ')
+                    ->subcompile($this->getNode('attribute'))
+                ;
+
+                if ($this->hasNode('arguments')) {
+                    $compiler->raw(', ')->subcompile($this->getNode('arguments'));
+                } else {
+                    $compiler->raw(', []');
+                }
+
+                $compiler->raw(', ')
+                    ->repr($this->getAttribute('type'))
+                    ->raw(', ')->repr($this->getAttribute('is_defined_test'))
+                    ->raw(', ')->repr($this->getAttribute('ignore_strict_check'))
+                    ->raw(', ')->repr($isSandboxed)
+                    ->raw(', ')->repr($this->getNode('node')->getTemplateLine())
+                    ->raw(')')
+                ;
+            },
+        ];
+    }
+
+    /**
+     * @return array{
+     *     condition: \Closure(Compiler, \Closure(Compiler): void): void,
+     *     accessor: \Closure(Compiler, \Closure(Compiler): void): void
+     * }
+     */
+    private function createObjectMethodAccessCompiler(ObjectType $type, string $attributeName): array
+    {
+        return [
+            'condition' => function (Compiler $compiler, \Closure $sourceCompiler) use ($type): void {
+                $sourceCompiler($compiler);
+                $compiler->raw(' instanceof \\')->raw($type->getType());
+            },
+            'accessor' => function (Compiler $compiler, \Closure $sourceCompiler) use ($attributeName): void {
+                $compiler->raw('(');
+
+                $sourceCompiler($compiler);
 
-        if ($this->hasNode('arguments')) {
-            $compiler->raw(', ')->subcompile($this->getNode('arguments'));
-        } else {
-            $compiler->raw(', []');
+                $compiler->raw('?->')->raw($attributeName)->raw('(');
+
+                if ($this->hasNode('arguments') && $this->getNode('arguments') instanceof ArrayExpression && $this->getNode('arguments')->count() > 0) {
+                    for ($argIndex = 0; $argIndex < $this->getNode('arguments')->count(); $argIndex += 2) {
+                        if ($argIndex > 0) {
+                            $compiler->raw(', ');
+                        }
+
+                        $compiler->subcompile($this->getNode('arguments')->getNode($argIndex + 1));
+                    }
+                }
+
+                $compiler->raw('))');
+            },
+        ];
+    }
+
+    /**
+     * @return array{
+     *     condition: \Closure(Compiler, \Closure(Compiler): void): void,
+     *     accessor: \Closure(Compiler, \Closure(Compiler): void): void
+     * }
+     */
+    private function createObjectPropertyAccessCompiler(ObjectType $type, string $attributeName): array
+    {
+        return [
+            'condition' => function (Compiler $compiler, \Closure $sourceCompiler) use ($type): void {
+                $sourceCompiler($compiler);
+                $compiler->raw(' instanceof \\')->raw($type->getType());
+            },
+            'accessor' => function (Compiler $compiler, \Closure $sourceCompiler) use ($attributeName): void {
+                $sourceCompiler($compiler);
+                $compiler
+                    ->raw('?->')
+                    ->raw($attributeName);
+            },
+        ];
+    }
+
+    /**
+     * @return array{
+     *     condition: null,
+     *     accessor: \Closure(Compiler, \Closure(Compiler): void): void
+     * }
+     */
+    private function createUnionAccessCompiler(UnionType $type, Environment $env): array
+    {
+        $accessors = [];
+
+        foreach ($type->getTypes() as $innerType) {
+            $accessors[] = $this->createAccessCompiler($innerType, $env);
         }
 
-        $compiler->raw(', ')
-            ->repr($this->getAttribute('type'))
-            ->raw(', ')->repr($this->getAttribute('is_defined_test'))
-            ->raw(', ')->repr($this->getAttribute('ignore_strict_check'))
-            ->raw(', ')->repr($env->hasExtension(SandboxExtension::class))
-            ->raw(', ')->repr($this->getNode('node')->getTemplateLine())
-            ->raw(')')
-        ;
+        return [
+            'condition' => null,
+            'accessor' => function (Compiler $compiler, \Closure $sourceCompiler) use ($accessors) {
+                $compiler->raw('match ([');
+                $compiler->indent();
+                $sourceCompiler($compiler);
+                $compiler->raw(", true][1]) {\n");
+
+                foreach ($accessors as $accessor) {
+                    if ($accessor['condition'] === null) {
+                        $compiler->raw('default');
+                    } else {
+                        $accessor['condition']($compiler, $sourceCompiler);
+                    }
+
+                    $compiler->raw(' => ');
+                    $accessor['accessor']($compiler, $sourceCompiler);
+                    $compiler->raw(";\n");
+                }
+
+                $compiler->outdent();
+                $compiler->raw('}');
+            }
+        ];
+    }
+
+    /**
+     * @return array{
+     *     condition: \Closure(Compiler, \Closure(Compiler): void): void,
+     *     accessor: \Closure(Compiler, \Closure(Compiler): void): void
+     * }
+     */
+    private function createArrayAccessCompiler(): array
+    {
+        return [
+            'condition' => function (Compiler $compiler, \Closure $sourceCompiler): void {
+                $compiler->raw('(\is_array(');
+                $sourceCompiler($compiler);
+                $compiler->raw(') || ');
+                $sourceCompiler($compiler);
+                $compiler->raw(' instanceof \\ArrayAccess)');
+            },
+            'accessor' => function (Compiler $compiler, \Closure $sourceCompiler): void {
+                $compiler->raw('(');
+                $sourceCompiler($compiler);
+                $compiler
+                    ->raw('[')
+                    ->subcompile($this->getNode('attribute'))
+                    ->raw('] ?? null)');
+            },
+        ];
+    }
+
+    /**
+     * @return \Closure(Compiler): void
+     */
+    private function createAutoInlineSourceCompiler(): \Closure
+    {
+        $varName = null;
+        $sourceCompiler = $this->createNodeSourceCompiler();
+
+        return function (Compiler $compiler) use (&$varName, &$sourceCompiler): void {
+            if ($varName === null) {
+                $varName = $compiler->getVarName();
+                $newSourceCompiler = $this->createVarNameSourceCompiler($varName);
+
+                $compiler->raw('(');
+                $newSourceCompiler($compiler);
+                $compiler->raw(' = ');
+                $sourceCompiler($compiler);
+                $compiler->raw(')');
+
+                $sourceCompiler = $newSourceCompiler;
+            } else {
+                $sourceCompiler($compiler);
+            }
+        };
+    }
+
+    /**
+     * @return \Closure(Compiler): void
+     */
+    private function createNodeSourceCompiler(): \Closure
+    {
+        return function (Compiler $compiler): void {
+            $compiler->subcompile($this->getNode('node'));
+        };
+    }
+
+    /**
+     * @return \Closure(Compiler): void
+     */
+    private function createVarNameSourceCompiler(string $varName): \Closure
+    {
+        return function (Compiler $compiler) use ($varName): void {
+            $compiler
+                ->raw('$')
+                ->raw($varName)
+            ;
+        };
     }
 }
diff --git a/tests/Node/Expression/GetAttrTest.php b/tests/Node/Expression/GetAttrTest.php
index 580da7459..19942ed2f 100644
--- a/tests/Node/Expression/GetAttrTest.php
+++ b/tests/Node/Expression/GetAttrTest.php
@@ -70,7 +70,7 @@ public function getTests()
             )->getNode('body'),
             <<<'PHP'
 // line 1
-echo ((((["bar" => ["baz" => 42]])["bar"] ?? null))["baz"] ?? null);
+echo ((["bar" => ["baz" => 42]]["bar"] ?? null)["baz"] ?? null);
 PHP,
             $optimizedEnv,
         ];
@@ -85,7 +85,7 @@ public function getTests()
 // line 1
 $context["foo"] = ["bar" => ["baz" => 42]];
 // line 2
-echo ((((($context["foo"] ?? null))["bar"] ?? null))["baz"] ?? null);
+echo ((($context["foo"] ?? null)["bar"] ?? null)["baz"] ?? null);
 PHP,
             $optimizedEnv,
         ];
@@ -98,7 +98,7 @@ public function getTests()
             <<<'PHP'
 // line 1
 // line 2
-echo ((($context["obj"] ?? null))?->name);
+echo ($context["obj"] ?? null)?->name;
 PHP,
             $optimizedEnv,
         ];
@@ -111,7 +111,7 @@ public function getTests()
             <<<'PHP'
 // line 1
 // line 2
-echo ((($context["obj"] ?? null))?->getname());
+echo (($context["obj"] ?? null)?->getname());
 PHP,
             $optimizedEnv,
         ];
@@ -124,7 +124,7 @@ public function getTests()
             <<<'PHP'
 // line 1
 // line 2
-echo ((($context["obj"] ?? null))?->byName("foobar"));
+echo (($context["obj"] ?? null)?->byName("foobar"));
 PHP,
             $optimizedEnv,
         ];
@@ -137,7 +137,23 @@ public function getTests()
             <<<'PHP'
 // line 1
 // line 2
-echo ((((($context["obj"] ?? null))?->getinstance()))?->getname());
+echo ((($context["obj"] ?? null)?->getinstance())?->getname());
+PHP,
+            $optimizedEnv,
+        ];
+
+        $tests[] = [
+            $optimizedEnv->parse($optimizedEnv->tokenize(new Source(<<<'TWIG'
+{% type obj "\\Twig\\Tests\\Node\\Expression\\ClassWithPublicProperty|\\Twig\\Tests\\Node\\Expression\\ClassWithPublicGetter" %}
+{{ obj.name|raw }}
+TWIG, 'index.twig')))->getNode('body'),
+            <<<'PHP'
+// line 1
+// line 2
+echo match ([($context["obj"] ?? null), true][1]) {
+($context["obj"] ?? null) instanceof \Twig\Tests\Node\Expression\ClassWithPublicProperty => ($context["obj"] ?? null)?->name;
+($context["obj"] ?? null) instanceof \Twig\Tests\Node\Expression\ClassWithPublicGetter => (($context["obj"] ?? null)?->getname());
+};
 PHP,
             $optimizedEnv,
         ];