1<?php declare(strict_types=1);
2
3namespace PhpParser\NodeVisitor;
4
5use PhpParser\Comment;
6use PhpParser\Node;
7use PhpParser\NodeVisitorAbstract;
8use PhpParser\Token;
9
10class CommentAnnotatingVisitor extends NodeVisitorAbstract {
11    /** @var int Last seen token start position */
12    private int $pos = 0;
13    /** @var Token[] Token array */
14    private array $tokens;
15    /** @var list<int> Token positions of comments */
16    private array $commentPositions = [];
17
18    /**
19     * Create a comment annotation visitor.
20     *
21     * @param Token[] $tokens Token array
22     */
23    public function __construct(array $tokens) {
24        $this->tokens = $tokens;
25
26        // Collect positions of comments. We use this to avoid traversing parts of the AST where
27        // there are no comments.
28        foreach ($tokens as $i => $token) {
29            if ($token->id === \T_COMMENT || $token->id === \T_DOC_COMMENT) {
30                $this->commentPositions[] = $i;
31            }
32        }
33    }
34
35    public function enterNode(Node $node) {
36        $nextCommentPos = current($this->commentPositions);
37        if ($nextCommentPos === false) {
38            // No more comments.
39            return self::STOP_TRAVERSAL;
40        }
41
42        $oldPos = $this->pos;
43        $this->pos = $pos = $node->getStartTokenPos();
44        if ($nextCommentPos > $oldPos && $nextCommentPos < $pos) {
45            $comments = [];
46            while (--$pos >= $oldPos) {
47                $token = $this->tokens[$pos];
48                if ($token->id === \T_DOC_COMMENT) {
49                    $comments[] = new Comment\Doc(
50                        $token->text, $token->line, $token->pos, $pos,
51                        $token->getEndLine(), $token->getEndPos() - 1, $pos);
52                    continue;
53                }
54                if ($token->id === \T_COMMENT) {
55                    $comments[] = new Comment(
56                        $token->text, $token->line, $token->pos, $pos,
57                        $token->getEndLine(), $token->getEndPos() - 1, $pos);
58                    continue;
59                }
60                if ($token->id !== \T_WHITESPACE) {
61                    break;
62                }
63            }
64            if (!empty($comments)) {
65                $node->setAttribute('comments', array_reverse($comments));
66            }
67
68            do {
69                $nextCommentPos = next($this->commentPositions);
70            } while ($nextCommentPos !== false && $nextCommentPos < $this->pos);
71        }
72
73        $endPos = $node->getEndTokenPos();
74        if ($nextCommentPos > $endPos) {
75            // Skip children if there are no comments located inside this node.
76            $this->pos = $endPos;
77            return self::DONT_TRAVERSE_CHILDREN;
78        }
79
80        return null;
81    }
82}
83