Chris@16: <?php
Chris@16: 
Chris@16: /*
Chris@16:  * This file is part of Psy Shell.
Chris@16:  *
Chris@16:  * (c) 2012-2018 Justin Hileman
Chris@16:  *
Chris@16:  * For the full copyright and license information, please view the LICENSE
Chris@16:  * file that was distributed with this source code.
Chris@16:  */
Chris@16: 
Chris@16: namespace Psy\Command\TimeitCommand;
Chris@16: 
Chris@16: use PhpParser\Node;
Chris@16: use PhpParser\Node\Arg;
Chris@16: use PhpParser\Node\Expr;
Chris@16: use PhpParser\Node\Expr\StaticCall;
Chris@16: use PhpParser\Node\FunctionLike;
Chris@16: use PhpParser\Node\Name\FullyQualified as FullyQualifiedName;
Chris@16: use PhpParser\Node\Stmt\Expression;
Chris@16: use PhpParser\Node\Stmt\Return_;
Chris@16: use PhpParser\NodeVisitorAbstract;
Chris@16: use Psy\CodeCleaner\NoReturnValue;
Chris@16: 
Chris@16: /**
Chris@16:  * A node visitor for instrumenting code to be executed by the `timeit` command.
Chris@16:  *
Chris@16:  * Injects `TimeitCommand::markStart()` at the start of code to be executed, and
Chris@16:  * `TimeitCommand::markEnd()` at the end, and on top-level return statements.
Chris@16:  */
Chris@16: class TimeitVisitor extends NodeVisitorAbstract
Chris@16: {
Chris@16:     private $functionDepth;
Chris@16: 
Chris@16:     /**
Chris@16:      * {@inheritdoc}
Chris@16:      */
Chris@16:     public function beforeTraverse(array $nodes)
Chris@16:     {
Chris@16:         $this->functionDepth = 0;
Chris@16:     }
Chris@16: 
Chris@16:     /**
Chris@16:      * {@inheritdoc}
Chris@16:      */
Chris@16:     public function enterNode(Node $node)
Chris@16:     {
Chris@16:         // keep track of nested function-like nodes, because they can have
Chris@16:         // returns statements... and we don't want to call markEnd for those.
Chris@16:         if ($node instanceof FunctionLike) {
Chris@16:             $this->functionDepth++;
Chris@16: 
Chris@16:             return;
Chris@16:         }
Chris@16: 
Chris@16:         // replace any top-level `return` statements with a `markEnd` call
Chris@16:         if ($this->functionDepth === 0 && $node instanceof Return_) {
Chris@16:             return new Return_($this->getEndCall($node->expr), $node->getAttributes());
Chris@16:         }
Chris@16:     }
Chris@16: 
Chris@16:     /**
Chris@16:      * {@inheritdoc}
Chris@16:      */
Chris@16:     public function leaveNode(Node $node)
Chris@16:     {
Chris@16:         if ($node instanceof FunctionLike) {
Chris@16:             $this->functionDepth--;
Chris@16:         }
Chris@16:     }
Chris@16: 
Chris@16:     /**
Chris@16:      * {@inheritdoc}
Chris@16:      */
Chris@16:     public function afterTraverse(array $nodes)
Chris@16:     {
Chris@16:         // prepend a `markStart` call
Chris@17:         \array_unshift($nodes, $this->maybeExpression($this->getStartCall()));
Chris@16: 
Chris@16:         // append a `markEnd` call (wrapping the final node, if it's an expression)
Chris@17:         $last = $nodes[\count($nodes) - 1];
Chris@16:         if ($last instanceof Expr) {
Chris@17:             \array_pop($nodes);
Chris@16:             $nodes[] = $this->getEndCall($last);
Chris@16:         } elseif ($last instanceof Expression) {
Chris@17:             \array_pop($nodes);
Chris@16:             $nodes[] = new Expression($this->getEndCall($last->expr), $last->getAttributes());
Chris@16:         } elseif ($last instanceof Return_) {
Chris@16:             // nothing to do here, we're already ending with a return call
Chris@16:         } else {
Chris@16:             $nodes[] = $this->maybeExpression($this->getEndCall());
Chris@16:         }
Chris@16: 
Chris@16:         return $nodes;
Chris@16:     }
Chris@16: 
Chris@16:     /**
Chris@16:      * Get PhpParser AST nodes for a `markStart` call.
Chris@16:      *
Chris@16:      * @return PhpParser\Node\Expr\StaticCall
Chris@16:      */
Chris@16:     private function getStartCall()
Chris@16:     {
Chris@16:         return new StaticCall(new FullyQualifiedName('Psy\Command\TimeitCommand'), 'markStart');
Chris@16:     }
Chris@16: 
Chris@16:     /**
Chris@16:      * Get PhpParser AST nodes for a `markEnd` call.
Chris@16:      *
Chris@16:      * Optionally pass in a return value.
Chris@16:      *
Chris@16:      * @param Expr|null $arg
Chris@16:      *
Chris@16:      * @return PhpParser\Node\Expr\StaticCall
Chris@16:      */
Chris@16:     private function getEndCall(Expr $arg = null)
Chris@16:     {
Chris@16:         if ($arg === null) {
Chris@16:             $arg = NoReturnValue::create();
Chris@16:         }
Chris@16: 
Chris@16:         return new StaticCall(new FullyQualifiedName('Psy\Command\TimeitCommand'), 'markEnd', [new Arg($arg)]);
Chris@16:     }
Chris@16: 
Chris@16:     /**
Chris@16:      * Compatibility shim for PHP Parser 3.x.
Chris@16:      *
Chris@16:      * Wrap $expr in a PhpParser\Node\Stmt\Expression if the class exists.
Chris@16:      *
Chris@16:      * @param PhpParser\Node $expr
Chris@16:      * @param array          $attrs
Chris@16:      *
Chris@16:      * @return PhpParser\Node\Expr|PhpParser\Node\Stmt\Expression
Chris@16:      */
Chris@16:     private function maybeExpression($expr, $attrs = [])
Chris@16:     {
Chris@17:         return \class_exists('PhpParser\Node\Stmt\Expression') ? new Expression($expr, $attrs) : $expr;
Chris@16:     }
Chris@16: }