1<?php
2
3declare(strict_types=1);
4
5namespace Antlr\Antlr4\Runtime\Tree;
6
7use Antlr\Antlr4\Runtime\Atn\ATN;
8use Antlr\Antlr4\Runtime\ParserRuleContext;
9use Antlr\Antlr4\Runtime\RuleContext;
10use Antlr\Antlr4\Runtime\Token;
11use Antlr\Antlr4\Runtime\Utils\StringUtils;
12
13/**
14 * A set of utility routines useful for all kinds of ANTLR trees.
15 */
16class Trees
17{
18    /**
19     * Print out a whole tree in LISP form. {@see Trees::getNodeText()} is used
20     * on the node payloads to get the text for the nodes. Detect parse trees
21     * and extract data appropriately.
22     *
23     * @param array<string>|null $ruleNames
24     */
25    public static function toStringTree(Tree $tree, ?array $ruleNames = null) : string
26    {
27        $string = self::getNodeText($tree, $ruleNames);
28        $string = StringUtils::escapeWhitespace($string, false);
29
30        $childCount = $tree->getChildCount();
31
32        if ($childCount === 0) {
33            return $string;
34        }
35
36        $result = '(' . $string . ' ';
37
38        for ($i = 0; $i < $childCount; $i++) {
39            $child = $tree->getChild($i);
40
41            if ($child !== null) {
42                $result .= ($i > 0 ? ' ' : '') . self::toStringTree($child, $ruleNames);
43            }
44        }
45
46        return $result . ')';
47    }
48
49    /**
50     * @param array<string>|null $ruleNames
51     */
52    public static function getNodeText(Tree $tree, ?array $ruleNames) : string
53    {
54        if ($ruleNames !== null) {
55            if ($tree instanceof RuleContext) {
56                $ruleIndex = $tree->getRuleContext()->getRuleIndex();
57                $ruleName = $ruleNames[$ruleIndex];
58                $altNumber = $tree->getAltNumber();
59
60                if ($altNumber !== ATN::INVALID_ALT_NUMBER) {
61                    return \sprintf('%s:%s', $ruleName, $altNumber);
62                }
63
64                return $ruleName;
65            }
66
67            if ($tree instanceof ErrorNode) {
68                return (string) $tree;
69            }
70
71            if ($tree instanceof TerminalNode && $tree->getSymbol() !== null) {
72                return $tree->getSymbol()->getText() ?? '';
73            }
74        }
75
76        // no recog for rule names
77        $payload = $tree->getPayload();
78
79        if ($payload instanceof Token) {
80            return $payload->getText() ?? '';
81        }
82
83        return (string) $tree->getPayload();
84    }
85
86    /**
87     * Return ordered list of all children of this node
88     *
89     * @return array<Tree|null>
90     */
91    public static function getChildren(Tree $tree) : array
92    {
93        $list = [];
94        for ($i=0; $i < $tree->getChildCount(); $i++) {
95            $list[] = $tree->getChild($i);
96        }
97
98        return $list;
99    }
100
101    /**
102     * Return a list of all ancestors of this node. The first node of list
103     * is the root and the last is the parent of this node.
104     *
105     * @return array<Tree>
106     */
107    public static function getAncestors(Tree $tree) : array
108    {
109        $ancestors = [];
110        $tree = $tree->getParent();
111
112        while ($tree !== null) {
113            \array_unshift($ancestors, $tree);
114
115            $tree = $tree->getParent();
116        }
117
118        return $ancestors;
119    }
120
121    /**
122     * @return array<ParseTree>
123     */
124    public static function findAllTokenNodes(ParseTree $tree, int $ttype) : array
125    {
126        return self::findAllNodes($tree, $ttype, true);
127    }
128
129    /**
130     * @return array<ParseTree>
131     */
132    public static function findAllRuleNodes(ParseTree $tree, int $ruleIndex) : array
133    {
134        return self::findAllNodes($tree, $ruleIndex, false);
135    }
136
137    /**
138     * @return array<ParseTree>
139     */
140    public static function findAllNodes(ParseTree $tree, int $index, bool $findTokens) : array
141    {
142        return self::findNodesInTree($tree, $index, $findTokens, []);
143    }
144
145    /**
146     * @param array<ParseTree> $nodes
147     *
148     * @return array<ParseTree>
149     */
150    private static function findNodesInTree(ParseTree $tree, int $index, bool $findTokens, array $nodes) : array
151    {
152        // check this node (the root) first
153        if ($findTokens && $tree instanceof TerminalNode && $tree->getSymbol()->getType() === $index) {
154            $nodes[] = $tree;
155        } elseif (!$findTokens && $tree instanceof ParserRuleContext && $tree->getRuleIndex() === $index) {
156            $nodes[] = $tree;
157        }
158
159        // check children
160        for ($i = 0; $i < $tree->getChildCount(); $i++) {
161            $child = $tree->getChild($i);
162
163            if ($child !== null) {
164                $nodes = self::findNodesInTree($child, $index, $findTokens, $nodes);
165            }
166        }
167
168        return $nodes;
169    }
170
171    /**
172     * @return array<ParseTree>
173     */
174    public static function descendants(ParseTree $tree) : array
175    {
176        $nodes = [$tree];
177        for ($i = 0; $i < $tree->getChildCount(); $i++) {
178            $child = $tree->getChild($i);
179
180            if ($child !== null) {
181                $nodes[] = self::descendants($child);
182            }
183        }
184
185        return \array_merge(...$nodes);
186    }
187}
188