1from .visitor import NodeVisitor
2
3VAR_LOAD_PARAMETER = "param"
4VAR_LOAD_RESOLVE = "resolve"
5VAR_LOAD_ALIAS = "alias"
6VAR_LOAD_UNDEFINED = "undefined"
7
8
9def find_symbols(nodes, parent_symbols=None):
10    sym = Symbols(parent=parent_symbols)
11    visitor = FrameSymbolVisitor(sym)
12    for node in nodes:
13        visitor.visit(node)
14    return sym
15
16
17def symbols_for_node(node, parent_symbols=None):
18    sym = Symbols(parent=parent_symbols)
19    sym.analyze_node(node)
20    return sym
21
22
23class Symbols:
24    def __init__(self, parent=None, level=None):
25        if level is None:
26            if parent is None:
27                level = 0
28            else:
29                level = parent.level + 1
30        self.level = level
31        self.parent = parent
32        self.refs = {}
33        self.loads = {}
34        self.stores = set()
35
36    def analyze_node(self, node, **kwargs):
37        visitor = RootVisitor(self)
38        visitor.visit(node, **kwargs)
39
40    def _define_ref(self, name, load=None):
41        ident = f"l_{self.level}_{name}"
42        self.refs[name] = ident
43        if load is not None:
44            self.loads[ident] = load
45        return ident
46
47    def find_load(self, target):
48        if target in self.loads:
49            return self.loads[target]
50        if self.parent is not None:
51            return self.parent.find_load(target)
52
53    def find_ref(self, name):
54        if name in self.refs:
55            return self.refs[name]
56        if self.parent is not None:
57            return self.parent.find_ref(name)
58
59    def ref(self, name):
60        rv = self.find_ref(name)
61        if rv is None:
62            raise AssertionError(
63                "Tried to resolve a name to a reference that was"
64                f" unknown to the frame ({name!r})"
65            )
66        return rv
67
68    def copy(self):
69        rv = object.__new__(self.__class__)
70        rv.__dict__.update(self.__dict__)
71        rv.refs = self.refs.copy()
72        rv.loads = self.loads.copy()
73        rv.stores = self.stores.copy()
74        return rv
75
76    def store(self, name):
77        self.stores.add(name)
78
79        # If we have not see the name referenced yet, we need to figure
80        # out what to set it to.
81        if name not in self.refs:
82            # If there is a parent scope we check if the name has a
83            # reference there.  If it does it means we might have to alias
84            # to a variable there.
85            if self.parent is not None:
86                outer_ref = self.parent.find_ref(name)
87                if outer_ref is not None:
88                    self._define_ref(name, load=(VAR_LOAD_ALIAS, outer_ref))
89                    return
90
91            # Otherwise we can just set it to undefined.
92            self._define_ref(name, load=(VAR_LOAD_UNDEFINED, None))
93
94    def declare_parameter(self, name):
95        self.stores.add(name)
96        return self._define_ref(name, load=(VAR_LOAD_PARAMETER, None))
97
98    def load(self, name):
99        target = self.find_ref(name)
100        if target is None:
101            self._define_ref(name, load=(VAR_LOAD_RESOLVE, name))
102
103    def branch_update(self, branch_symbols):
104        stores = {}
105        for branch in branch_symbols:
106            for target in branch.stores:
107                if target in self.stores:
108                    continue
109                stores[target] = stores.get(target, 0) + 1
110
111        for sym in branch_symbols:
112            self.refs.update(sym.refs)
113            self.loads.update(sym.loads)
114            self.stores.update(sym.stores)
115
116        for name, branch_count in stores.items():
117            if branch_count == len(branch_symbols):
118                continue
119            target = self.find_ref(name)
120            assert target is not None, "should not happen"
121
122            if self.parent is not None:
123                outer_target = self.parent.find_ref(name)
124                if outer_target is not None:
125                    self.loads[target] = (VAR_LOAD_ALIAS, outer_target)
126                    continue
127            self.loads[target] = (VAR_LOAD_RESOLVE, name)
128
129    def dump_stores(self):
130        rv = {}
131        node = self
132        while node is not None:
133            for name in node.stores:
134                if name not in rv:
135                    rv[name] = self.find_ref(name)
136            node = node.parent
137        return rv
138
139    def dump_param_targets(self):
140        rv = set()
141        node = self
142        while node is not None:
143            for target, (instr, _) in self.loads.items():
144                if instr == VAR_LOAD_PARAMETER:
145                    rv.add(target)
146            node = node.parent
147        return rv
148
149
150class RootVisitor(NodeVisitor):
151    def __init__(self, symbols):
152        self.sym_visitor = FrameSymbolVisitor(symbols)
153
154    def _simple_visit(self, node, **kwargs):
155        for child in node.iter_child_nodes():
156            self.sym_visitor.visit(child)
157
158    visit_Template = (
159        visit_Block
160    ) = (
161        visit_Macro
162    ) = (
163        visit_FilterBlock
164    ) = visit_Scope = visit_If = visit_ScopedEvalContextModifier = _simple_visit
165
166    def visit_AssignBlock(self, node, **kwargs):
167        for child in node.body:
168            self.sym_visitor.visit(child)
169
170    def visit_CallBlock(self, node, **kwargs):
171        for child in node.iter_child_nodes(exclude=("call",)):
172            self.sym_visitor.visit(child)
173
174    def visit_OverlayScope(self, node, **kwargs):
175        for child in node.body:
176            self.sym_visitor.visit(child)
177
178    def visit_For(self, node, for_branch="body", **kwargs):
179        if for_branch == "body":
180            self.sym_visitor.visit(node.target, store_as_param=True)
181            branch = node.body
182        elif for_branch == "else":
183            branch = node.else_
184        elif for_branch == "test":
185            self.sym_visitor.visit(node.target, store_as_param=True)
186            if node.test is not None:
187                self.sym_visitor.visit(node.test)
188            return
189        else:
190            raise RuntimeError("Unknown for branch")
191        for item in branch or ():
192            self.sym_visitor.visit(item)
193
194    def visit_With(self, node, **kwargs):
195        for target in node.targets:
196            self.sym_visitor.visit(target)
197        for child in node.body:
198            self.sym_visitor.visit(child)
199
200    def generic_visit(self, node, *args, **kwargs):
201        raise NotImplementedError(
202            f"Cannot find symbols for {node.__class__.__name__!r}"
203        )
204
205
206class FrameSymbolVisitor(NodeVisitor):
207    """A visitor for `Frame.inspect`."""
208
209    def __init__(self, symbols):
210        self.symbols = symbols
211
212    def visit_Name(self, node, store_as_param=False, **kwargs):
213        """All assignments to names go through this function."""
214        if store_as_param or node.ctx == "param":
215            self.symbols.declare_parameter(node.name)
216        elif node.ctx == "store":
217            self.symbols.store(node.name)
218        elif node.ctx == "load":
219            self.symbols.load(node.name)
220
221    def visit_NSRef(self, node, **kwargs):
222        self.symbols.load(node.name)
223
224    def visit_If(self, node, **kwargs):
225        self.visit(node.test, **kwargs)
226
227        original_symbols = self.symbols
228
229        def inner_visit(nodes):
230            self.symbols = rv = original_symbols.copy()
231            for subnode in nodes:
232                self.visit(subnode, **kwargs)
233            self.symbols = original_symbols
234            return rv
235
236        body_symbols = inner_visit(node.body)
237        elif_symbols = inner_visit(node.elif_)
238        else_symbols = inner_visit(node.else_ or ())
239
240        self.symbols.branch_update([body_symbols, elif_symbols, else_symbols])
241
242    def visit_Macro(self, node, **kwargs):
243        self.symbols.store(node.name)
244
245    def visit_Import(self, node, **kwargs):
246        self.generic_visit(node, **kwargs)
247        self.symbols.store(node.target)
248
249    def visit_FromImport(self, node, **kwargs):
250        self.generic_visit(node, **kwargs)
251        for name in node.names:
252            if isinstance(name, tuple):
253                self.symbols.store(name[1])
254            else:
255                self.symbols.store(name)
256
257    def visit_Assign(self, node, **kwargs):
258        """Visit assignments in the correct order."""
259        self.visit(node.node, **kwargs)
260        self.visit(node.target, **kwargs)
261
262    def visit_For(self, node, **kwargs):
263        """Visiting stops at for blocks.  However the block sequence
264        is visited as part of the outer scope.
265        """
266        self.visit(node.iter, **kwargs)
267
268    def visit_CallBlock(self, node, **kwargs):
269        self.visit(node.call, **kwargs)
270
271    def visit_FilterBlock(self, node, **kwargs):
272        self.visit(node.filter, **kwargs)
273
274    def visit_With(self, node, **kwargs):
275        for target in node.values:
276            self.visit(target)
277
278    def visit_AssignBlock(self, node, **kwargs):
279        """Stop visiting at block assigns."""
280        self.visit(node.target, **kwargs)
281
282    def visit_Scope(self, node, **kwargs):
283        """Stop visiting at scopes."""
284
285    def visit_Block(self, node, **kwargs):
286        """Stop visiting at blocks."""
287
288    def visit_OverlayScope(self, node, **kwargs):
289        """Do not visit into overlay scopes."""
290