From 536889f2b340489d328f5ffb7b02bb6b183ddedc Mon Sep 17 00:00:00 2001
From: Richard van Velzen <richard@goreply.com>
Date: Mon, 6 May 2024 14:04:23 +0200
Subject: [PATCH] Support array and offset access on const types

---
 src/Ast/Type/OffsetAccessTypeNode.php   |  1 -
 src/Parser/TypeParser.php               | 26 +++++++++++++++++--------
 src/Printer/Printer.php                 |  2 --
 tests/PHPStan/Parser/TypeParserTest.php |  7 +++++++
 tests/PHPStan/Printer/PrinterTest.php   | 24 +++++++++++++++++++++++
 5 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/src/Ast/Type/OffsetAccessTypeNode.php b/src/Ast/Type/OffsetAccessTypeNode.php
index 39e83dfe..c27ec0a3 100644
--- a/src/Ast/Type/OffsetAccessTypeNode.php
+++ b/src/Ast/Type/OffsetAccessTypeNode.php
@@ -25,7 +25,6 @@ public function __toString(): string
 	{
 		if (
 			$this->type instanceof CallableTypeNode
-			|| $this->type instanceof ConstTypeNode
 			|| $this->type instanceof NullableTypeNode
 		) {
 			return '(' . $this->type . ')[' . $this->offset . ']';
diff --git a/src/Parser/TypeParser.php b/src/Parser/TypeParser.php
index ebc2fbab..5669fe45 100644
--- a/src/Parser/TypeParser.php
+++ b/src/Parser/TypeParser.php
@@ -232,7 +232,17 @@ private function parseAtomic(TokenIterator $tokens): Ast\Type\TypeNode
 				);
 			}
 
-			return $this->enrichWithAttributes($tokens, new Ast\Type\ConstTypeNode($constExpr), $startLine, $startIndex);
+			$type = $this->enrichWithAttributes(
+				$tokens,
+				new Ast\Type\ConstTypeNode($constExpr),
+				$startLine,
+				$startIndex
+			);
+			if ($tokens->isCurrentTokenType(Lexer::TOKEN_OPEN_SQUARE_BRACKET)) {
+				$type = $this->tryParseArrayOrOffsetAccess($tokens, $type);
+			}
+
+			return $type;
 		} catch (LogicException $e) {
 			throw new ParserException(
 				$currentTokenValue,
@@ -733,14 +743,14 @@ private function parseCallableReturnType(TokenIterator $tokens): Ast\Type\TypeNo
 				);
 			}
 
-			$type = new Ast\Type\ConstTypeNode($constExpr);
+			$type = $this->enrichWithAttributes(
+				$tokens,
+				new Ast\Type\ConstTypeNode($constExpr),
+				$startLine,
+				$startIndex
+			);
 			if ($tokens->isCurrentTokenType(Lexer::TOKEN_OPEN_SQUARE_BRACKET)) {
-				$type = $this->tryParseArrayOrOffsetAccess($tokens, $this->enrichWithAttributes(
-					$tokens,
-					$type,
-					$startLine,
-					$startIndex
-				));
+				$type = $this->tryParseArrayOrOffsetAccess($tokens, $type);
 			}
 
 			return $type;
diff --git a/src/Printer/Printer.php b/src/Printer/Printer.php
index b1769932..044d07f8 100644
--- a/src/Printer/Printer.php
+++ b/src/Printer/Printer.php
@@ -141,7 +141,6 @@ final class Printer
 			CallableTypeNode::class,
 			UnionTypeNode::class,
 			IntersectionTypeNode::class,
-			ConstTypeNode::class,
 			NullableTypeNode::class,
 		],
 	];
@@ -512,7 +511,6 @@ private function printOffsetAccessType(TypeNode $type): string
 			$type instanceof CallableTypeNode
 			|| $type instanceof UnionTypeNode
 			|| $type instanceof IntersectionTypeNode
-			|| $type instanceof ConstTypeNode
 			|| $type instanceof NullableTypeNode
 		) {
 			return $this->wrapInParentheses($type);
diff --git a/tests/PHPStan/Parser/TypeParserTest.php b/tests/PHPStan/Parser/TypeParserTest.php
index babaf9d1..d6c66bb8 100644
--- a/tests/PHPStan/Parser/TypeParserTest.php
+++ b/tests/PHPStan/Parser/TypeParserTest.php
@@ -1067,6 +1067,13 @@ public function provideParseData(): array
 					new IdentifierTypeNode('int')
 				),
 			],
+			[
+				'self::TYPES[ int ]',
+				new OffsetAccessTypeNode(
+					new ConstTypeNode(new ConstFetchNode('self', 'TYPES')),
+					new IdentifierTypeNode('int')
+				),
+			],
 			[
 				"?\t\xA009", // edge-case with \h
 				new NullableTypeNode(
diff --git a/tests/PHPStan/Printer/PrinterTest.php b/tests/PHPStan/Printer/PrinterTest.php
index 25a268ee..d73481e2 100644
--- a/tests/PHPStan/Printer/PrinterTest.php
+++ b/tests/PHPStan/Printer/PrinterTest.php
@@ -1723,6 +1723,23 @@ public function enterNode(Node $node)
 
 			},
 		];
+
+		yield [
+			'/** @return Foo[abc] */',
+			'/** @return self::FOO[abc] */',
+			new class extends AbstractNodeVisitor {
+
+				public function enterNode(Node $node)
+				{
+					if ($node instanceof ReturnTagValueNode && $node->type instanceof OffsetAccessTypeNode) {
+						$node->type->type = new ConstTypeNode(new ConstFetchNode('self', 'FOO'));
+					}
+
+					return $node;
+				}
+
+			},
+		];
 	}
 
 	/**
@@ -1862,6 +1879,13 @@ public function dataPrintType(): iterable
 			]),
 			'Foo|Bar|(Baz|Lorem)',
 		];
+		yield [
+			new OffsetAccessTypeNode(
+				new ConstTypeNode(new ConstFetchNode('self', 'TYPES')),
+				new IdentifierTypeNode('int')
+			),
+			'self::TYPES[int]',
+		];
 	}
 
 	/**