refactored cleaning parser logic

phpstan · Jan 30, 2025 · 32cf692 · 32cf692
1 parent 13d7837
commit 32cf692
Show file tree

Hide file tree

Showing 3 changed files with 128 additions and 65 deletions.
diff --git a/src/Parser/CleaningVisitor.php b/src/Parser/CleaningVisitor.php
@@ -3,102 +3,152 @@
 namespace PHPStan\Parser;
 
 use PhpParser\Node;
-use PhpParser\NodeFinder;
+use PhpParser\NodeTraverser;
 use PhpParser\NodeVisitorAbstract;
 use PHPStan\Reflection\ParametersAcceptor;
+use PHPStan\ShouldNotHappenException;
+use function array_filter;
+use function array_map;
 use function in_array;
 use function is_array;
 
 final class CleaningVisitor extends NodeVisitorAbstract
 {
 
-	private NodeFinder $nodeFinder;
+	private const CONTEXT_DEFAULT = 0;
 
-	public function __construct()
+	private const CONTEXT_FUNCTION_OR_METHOD = 1;
+
+	private const CONTEXT_PROPERTY_HOOK = 2;
+
+	/** @var self::CONTEXT_* */
+	private int $context = self::CONTEXT_DEFAULT;
+
+	private string|null $propertyName = null;
+
+	/**
+	 * @return int|Node[]|null
+	 */
+	public function enterNode(Node $node): int|array|null
 	{
-		$this->nodeFinder = new NodeFinder();
+		return match ($this->context) {
+			self::CONTEXT_DEFAULT => $this->clean($node),
+			self::CONTEXT_FUNCTION_OR_METHOD => $this->cleanFunctionOrMethod($node),
+			self::CONTEXT_PROPERTY_HOOK => $this->cleanPropertyHook($node),
+		};
 	}
 
-	public function enterNode(Node $node): ?Node
+	private function clean(Node $node): int|null
 	{
-		if ($node instanceof Node\Stmt\Function_) {
-			$node->stmts = $this->keepVariadicsAndYields($node->stmts, null);
-			return $node;
-		}
+		if (($node instanceof Node\Stmt\Function_ || $node instanceof Node\Stmt\ClassMethod) && $node->stmts !== null) {
+			$params = [];
+			foreach ($this->traverse($node->params, self::CONTEXT_DEFAULT) as $param) {
+				$params[] = $param instanceof Node\Param ? $param : throw new ShouldNotHappenException();
+			}
+			$node->params = $params;
 
-		if ($node instanceof Node\Stmt\ClassMethod && $node->stmts !== null) {
-			$node->stmts = $this->keepVariadicsAndYields($node->stmts, null);
-			return $node;
-		}
+			$stmts = [];
+			foreach ($this->traverse($node->stmts, self::CONTEXT_FUNCTION_OR_METHOD) as $stmt) {
+				$stmts[] = $stmt instanceof Node\Stmt ? $stmt : throw new ShouldNotHappenException();
+			}
+			$node->stmts = $stmts;
 
-		if ($node instanceof Node\Expr\Closure) {
-			$node->stmts = $this->keepVariadicsAndYields($node->stmts, null);
-			return $node;
+			return self::DONT_TRAVERSE_CHILDREN;
 		}
 
 		if ($node instanceof Node\PropertyHook && is_array($node->body)) {
 			$propertyName = $node->getAttribute('propertyName');
 			if ($propertyName !== null) {
-				$node->body = $this->keepVariadicsAndYields($node->body, $propertyName);
-				return $node;
+				$body = [];
+				foreach ($this->traverse($node->body, self::CONTEXT_PROPERTY_HOOK, $propertyName) as $stmt) {
+					$body[] = $stmt instanceof Node\Stmt ? $stmt : throw new ShouldNotHappenException();
+				}
+				$node->body = $body;
+
+				return self::DONT_TRAVERSE_CHILDREN;
 			}
 		}
 
 		return null;
 	}
 
 	/**
-	 * @param Node\Stmt[] $stmts
-	 * @return Node\Stmt[]
+	 * @return int|Node[]
 	 */
-	private function keepVariadicsAndYields(array $stmts, ?string $hookedPropertyName): array
+	private function cleanFunctionOrMethod(Node $node): int|array
 	{
-		$results = $this->nodeFinder->find($stmts, static function (Node $node) use ($hookedPropertyName): bool {
-			if ($node instanceof Node\Expr\YieldFrom || $node instanceof Node\Expr\Yield_) {
-				return true;
-			}
-			if ($node instanceof Node\Expr\FuncCall && $node->name instanceof Node\Name) {
-				return in_array($node->name->toLowerString(), ParametersAcceptor::VARIADIC_FUNCTIONS, true);
-			}
+		if ($node instanceof Node\Expr\YieldFrom || $node instanceof Node\Expr\Yield_) {
+			return self::DONT_TRAVERSE_CHILDREN;
+		}
 
-			if ($node instanceof Node\Expr\Closure || $node instanceof Node\Expr\ArrowFunction) {
-				return true;
-			}
+		if ($node instanceof Node\Expr\FuncCall && $node->name instanceof Node\Name
+			&& in_array($node->name->toLowerString(), ParametersAcceptor::VARIADIC_FUNCTIONS, true)
+		) {
+			$node->name = new Node\Name\FullyQualified('func_get_args');
+			return self::DONT_TRAVERSE_CHILDREN;
+		}
 
-			if ($hookedPropertyName !== null) {
-				if (
-					$node instanceof Node\Expr\PropertyFetch
-					&& $node->var instanceof Node\Expr\Variable
-					&& $node->var->name === 'this'
-					&& $node->name instanceof Node\Identifier
-					&& $node->name->toString() === $hookedPropertyName
-				) {
-					return true;
-				}
-			}
+		if ($node instanceof Node\Expr\Closure || $node instanceof Node\Expr\ArrowFunction) {
+			return self::REMOVE_NODE;
+		}
 
-			return false;
-		});
-		$newStmts = [];
-		foreach ($results as $result) {
-			if (
-				$result instanceof Node\Expr\Yield_
-				|| $result instanceof Node\Expr\YieldFrom
-				|| $result instanceof Node\Expr\Closure
-				|| $result instanceof Node\Expr\ArrowFunction
-				|| $result instanceof Node\Expr\PropertyFetch
-			) {
-				$newStmts[] = new Node\Stmt\Expression($result);
-				continue;
-			}
-			if (!$result instanceof Node\Expr\FuncCall) {
-				continue;
-			}
+		return $this->cleanSubnodes($node);
+	}
+
+	/**
+	 * @param Node[] $nodes
+	 * @param self::CONTEXT_* $context
+	 * @return Node[]
+	 */
+	private function traverse(
+		array $nodes,
+		int $context = self::CONTEXT_DEFAULT,
+		string|null $propertyName = null,
+	): array
+	{
+		$visitor = new self();
+		$visitor->context = $context;
+		$visitor->propertyName = $propertyName;
 
-			$newStmts[] = new Node\Stmt\Expression(new Node\Expr\FuncCall(new Node\Name\FullyQualified('func_get_args')));
+		return (new NodeTraverser($visitor))->traverse($nodes);
+	}
+
+	/**
+	 * @return Node[]
+	 */
+	private function cleanPropertyHook(Node $node): int|array
+	{
+		if (
+			$node instanceof Node\Expr\PropertyFetch
+			&& $node->var instanceof Node\Expr\Variable
+			&& $node->var->name === 'this'
+			&& $node->name instanceof Node\Identifier
+			&& $node->name->toString() === $this->propertyName
+		) {
+			return self::DONT_TRAVERSE_CHILDREN;
+		}
+
+		return $this->cleanSubnodes($node);
+	}
+
+	/**
+	 * @return Node[]
+	 */
+	private function cleanSubnodes(Node $node): array
+	{
+		$subnodes = [];
+		foreach ($node->getSubNodeNames() as $subnodeName) {
+			$subnodes = [...$subnodes, ...array_filter(
+				is_array($node->$subnodeName) ? $node->$subnodeName : [$node->$subnodeName],
+				static fn ($subnode) => $subnode instanceof Node,
+			)];
 		}
 
-		return $newStmts;
+		return array_map(static fn ($node) => match (true) {
+				$node instanceof Node\Stmt => $node,
+				$node instanceof Node\Expr => new Node\Stmt\Expression($node),
+				default => throw new ShouldNotHappenException()
+		}, $this->traverse($subnodes, $this->context, $this->propertyName));
 	}
 
 }
diff --git a/tests/PHPStan/Parser/data/cleaning-1-after.php b/tests/PHPStan/Parser/data/cleaning-1-after.php
@@ -21,6 +21,10 @@ public function someGenerator2()
     {
         yield from [1, 2, 3];
     }
+    public function someGenerator3()
+    {
+        yield;
+    }
     public function someVariadics()
     {
         \func_get_args();
@@ -43,9 +47,8 @@ class ContainsClosure
 {
     public function doFoo()
     {
-        static function () {
-            yield;
-        };
-        yield;
+    }
+    public function doBar()
+    {
     }
 }
diff --git a/tests/PHPStan/Parser/data/cleaning-1-before.php b/tests/PHPStan/Parser/data/cleaning-1-before.php
@@ -36,6 +36,11 @@ public function someGenerator2()
 		}
 	}
 
+	public function someGenerator3()
+	{
+		echo yield;
+	}
+
 	public function someVariadics()
 	{
 		if (rand(0, 1)) {
@@ -82,4 +87,9 @@ public function doFoo()
 		};
 	}
 
+	public function doBar()
+	{
+		$fn = fn() => yield;
+	}
+
 }