Coverage for skema/program_analysis/CAST2FN/ann_cast/container_scope_pass.py: 91%
304 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-30 17:15 +0000
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-30 17:15 +0000
1import copy
2import typing
3from collections import defaultdict
4from enum import Enum
5from functools import singledispatchmethod
7from skema.program_analysis.CAST2FN.ann_cast.ann_cast_helpers import (
8 CON_STR_SEP,
9 ELSEBODY,
10 IFBODY,
11 IFEXPR,
12 LOOPPRE,
13 LOOPBODY,
14 LOOPEXPR,
15 LOOPPOST,
16 MODULE_SCOPE,
17 GrfnContainerSrcRef,
18 call_container_name,
19 combine_grfn_con_src_refs,
20 combine_source_refs,
21 con_scope_to_str,
22 func_def_container_name,
23 var_dict_to_str,
24)
25from skema.program_analysis.CAST2FN.ann_cast.annotated_cast import *
26from skema.program_analysis.CAST2FN.model.cast import (
27 ScalarType,
28 StructureType,
29 ValueConstructor,
30)
33class AssignSide(Enum):
34 NEITHER = 0
35 LEFT = 1
36 RIGHT = 2
39class ContainerData:
40 modified_vars: typing.Dict[int, str]
41 vars_accessed_before_mod: typing.Dict[int, str]
42 used_vars: typing.Dict[int, str]
44 def __init__(self):
45 self.modified_vars = {}
46 self.vars_accessed_before_mod = {}
47 self.used_vars = {}
50class ContainerScopePass:
51 def __init__(self, pipeline_state: PipelineState):
52 self.pipeline_state = pipeline_state
53 # dicts mapping container scope strs to the if/loop count inside
54 # the container
55 self.if_count = defaultdict(int)
56 self.loop_count = defaultdict(int)
57 # dict mapping container scope str to AnnCastNode
58 self.con_str_to_node = {}
59 # dict mapping container scope str to cached Container Data
60 self.con_str_to_con_data = {}
61 self.calls_to_process = list()
63 for node in self.pipeline_state.nodes:
64 # assign_side is False at the start of our visitor
65 base_scopestr = ""
66 enclosing_con_scope = []
67 self.visit(
68 node, base_scopestr, enclosing_con_scope, AssignSide.NEITHER
69 )
70 self.nodes = self.pipeline_state.nodes
72 # add cached container data to container nodes
73 self.add_container_data_to_nodes()
75 # save the dict mapping container scope to AnnCastNode
76 self.pipeline_state.con_scopestr_to_node = self.con_str_to_node
78 self.propagate_globals_through_calls()
80 def next_if_scope(self, enclosing_con_scope):
81 scopestr = con_scope_to_str(enclosing_con_scope)
82 count = self.if_count[scopestr]
83 self.if_count[scopestr] += 1
84 return enclosing_con_scope + [f"if{count}"]
86 def next_loop_scope(self, enclosing_con_scope):
87 scopestr = con_scope_to_str(enclosing_con_scope)
88 count = self.loop_count[scopestr]
89 self.loop_count[scopestr] += 1
90 return enclosing_con_scope + [f"loop{count}"]
92 def propagate_globals_through_calls(self):
93 for call_node in self.calls_to_process:
94 func_def = self.pipeline_state.func_def_node_from_id(
95 call_node.func.id
96 )
98 # propagate up used variables to enclosing container scopes
99 scopestr = ""
100 for index, name in enumerate(call_node.func.con_scope):
101 scopestr2 = CON_STR_SEP.join(
102 call_node.func.con_scope[: index + 1]
103 )
104 # add separator between container scope component names
105 if index != 0:
106 scopestr += f"{CON_STR_SEP}"
107 scopestr += f"{name}"
108 assert scopestr == scopestr2
110 if (
111 scopestr == MODULE_SCOPE
112 or not self.pipeline_state.is_container(scopestr)
113 ):
114 continue
116 container_node = self.pipeline_state.con_node_from_scopestr(
117 scopestr
118 )
120 if self.pipeline_state.is_con_scopestr_func_def(scopestr):
121 container_node.used_globals.update(func_def.used_globals)
122 container_node.modified_globals.update(
123 func_def.modified_globals
124 )
125 container_node.globals_accessed_before_mod.update(
126 func_def.globals_accessed_before_mod
127 )
128 container_node.used_vars.update(func_def.used_globals)
129 container_node.modified_vars.update(func_def.modified_globals)
130 container_node.vars_accessed_before_mod.update(
131 func_def.globals_accessed_before_mod
132 )
134 def add_container_data_to_expr(self, container, data):
135 """
136 Adds container data to `expr_*_vars` attributes of ModelIf and Loop nodes
137 """
138 container.expr_vars_accessed_before_mod = data.vars_accessed_before_mod
139 container.expr_modified_vars = data.modified_vars
140 container.expr_used_vars = data.used_vars
142 def add_container_data_to_nodes(self):
143 for scopestr, data in self.con_str_to_con_data.items():
144 # DEBUG printing
145 if self.pipeline_state.PRINT_DEBUGGING_INFO:
146 print(f"For scopestr: {scopestr} found data with")
147 modified_vars = var_dict_to_str(
148 " Modified: ", data.modified_vars
149 )
150 print(modified_vars)
151 vars_accessed_before_mod = var_dict_to_str(
152 " Accessed: ", data.vars_accessed_before_mod
153 )
154 print(vars_accessed_before_mod)
155 used_vars = var_dict_to_str(
156 " Used: ", data.vars_accessed_before_mod
157 )
158 print(used_vars)
160 # Note: for the ModelIf.Expr and Loop.Expr nodes,
161 # we put the ModelIf and Loop nodes respectively in
162 # `con_str_to_node`.
163 # We need to put the container data for the Expr nodes in
164 # the expr_*_vars attributes of their associated container nodes
165 # so we call `add_container_data_to_expr()`
166 if_expr_suffix = CON_STR_SEP + IFEXPR
167 if scopestr.endswith(if_expr_suffix):
168 if_container = self.con_str_to_node[scopestr]
169 self.add_container_data_to_expr(if_container, data)
170 continue
172 loop_expr_suffix = CON_STR_SEP + LOOPEXPR
173 if scopestr.endswith(loop_expr_suffix):
174 loop_container = self.con_str_to_node[scopestr]
175 self.add_container_data_to_expr(loop_container, data)
176 continue
178 # otherwise, store container data, in the container nodes
179 # *_vars attributes
180 container = self.con_str_to_node[scopestr]
181 container.vars_accessed_before_mod = data.vars_accessed_before_mod
182 container.modified_vars = data.modified_vars
183 container.used_vars = data.used_vars
185 # if the container is a FunctionDef, we want to store how globals are used
186 if isinstance(container, AnnCastFunctionDef):
187 all_globals = self.pipeline_state.all_globals_dict()
188 for id, name in all_globals.items():
189 if id in container.vars_accessed_before_mod:
190 container.globals_accessed_before_mod[id] = name
191 if id in container.modified_vars:
192 container.modified_globals[id] = name
193 if id in container.used_vars:
194 container.used_globals[id] = name
196 # DEBUG printing
197 if self.pipeline_state.PRINT_DEBUGGING_INFO:
198 print(container.grfn_con_src_ref)
200 def initialize_con_scope_data(self, con_scope: typing.List, node):
201 """
202 Create an empty `ContainterData` in `self.con_str_to_con_data`
203 and cache the container `node` in `self.con_str_to_node`
204 """
205 con_scopestr = con_scope_to_str(con_scope)
206 # initialize container data for this node
207 self.con_str_to_con_data[con_scopestr] = ContainerData()
209 # map con_scopestr to passed in node
210 self.con_str_to_node[con_scopestr] = node
212 def visit(
213 self,
214 node: AnnCastNode,
215 base_func_scopestr: str,
216 enclosing_con_scope: typing.List,
217 assign_side: AssignSide,
218 ):
219 # print current node being visited.
220 # this can be useful for debugging
221 # class_name = node.__class__.__name__
222 # print(f"\nProcessing node type {class_name}")
224 children_src_ref = self._visit(
225 node, base_func_scopestr, enclosing_con_scope, assign_side
226 )
227 if children_src_ref is None:
228 children_src_ref = GrfnContainerSrcRef(None, None, None)
230 # to keep determine GrfnContainerSrcRef for enclosing containers
231 # each node we visit returns None or a GrfnContainerSrcRef with data copied from the nodes
232 # source_refs attribute
233 grfn_src_ref = GrfnContainerSrcRef(None, None, None)
234 if node.source_refs is not None:
235 src_ref = combine_source_refs(node.source_refs)
236 grfn_src_ref = GrfnContainerSrcRef(
237 line_begin=src_ref.row_start,
238 line_end=src_ref.row_end,
239 source_file_name=src_ref.source_file_name,
240 )
242 return combine_grfn_con_src_refs([children_src_ref, grfn_src_ref])
244 @singledispatchmethod
245 def _visit(
246 self,
247 node: AnnCastNode,
248 base_func_scopestr: str,
249 enclosing_con_scope: typing.List,
250 assign_side: AssignSide,
251 ):
252 """
253 Visit each AnnCastNode
254 Parameters:
255 - `assign_side`: this denotes whether we are visiting the LHS or RHS of an AnnCastAssignment
256 or if we are not under an AnnCastAssignment
257 This is used to determine whether a variable (AnnCastName node) is
258 accessed or modified in that context
259 """
260 raise Exception(f"Unimplemented AST node of type: {type(node)}")
262 def visit_node_list(
263 self,
264 node_list: typing.List[AnnCastNode],
265 base_func_scopestr,
266 enclosing_con_scope,
267 assign_side,
268 ):
269 grfn_src_refs = [
270 self.visit(
271 node, base_func_scopestr, enclosing_con_scope, assign_side
272 )
273 for node in node_list
274 ]
275 #print("---------")
276 #print(node_list)
277 #print(grfn_src_refs)
278 #print("---------")
279 return combine_grfn_con_src_refs(grfn_src_refs)
281 @_visit.register
282 def visit_assignment(
283 self,
284 node: AnnCastAssignment,
285 base_func_scopestr,
286 enclosing_con_scope,
287 assign_side,
288 ):
289 right_src_ref = self.visit(
290 node.right,
291 base_func_scopestr,
292 enclosing_con_scope,
293 AssignSide.RIGHT,
294 )
295 # The AnnCastTuple is added to handle scenarios where an assignment
296 # is made by assigning to a tuple of values, as opposed to one singular value
297 assert (
298 isinstance(node.left, AnnCastVar)
299 or (isinstance(node.left, AnnCastLiteralValue) and (node.left.value_type == StructureType.TUPLE))
300 or isinstance(node.left, AnnCastAttribute) or isinstance(node.left, AnnCastCall)
301 ), f"container_scope: visit_assigment: node.left is {type(node.left)}"
302 left_src_ref = self.visit(
303 node.left, base_func_scopestr, enclosing_con_scope, AssignSide.LEFT
304 )
306 return combine_grfn_con_src_refs([right_src_ref, left_src_ref])
308 @_visit.register
309 def visit_attribute(
310 self,
311 node: AnnCastAttribute,
312 base_func_scopestr,
313 enclosing_con_scope,
314 assign_side,
315 ):
316 # TODO: what to do with the attr?
317 self.visit(
318 node.value, base_func_scopestr, enclosing_con_scope, assign_side
319 )
321 @_visit.register
322 def visit_call(
323 self,
324 node: AnnCastCall,
325 base_func_scopestr,
326 enclosing_con_scope,
327 assign_side,
328 ):
329 assert isinstance(node.func, AnnCastName) or isinstance(
330 node.func, AnnCastAttribute
331 )
332 # if this call is on the RHS of an assignment, then it should have a ret val
333 # FUTURE: this logic is not sufficient to determine
334 # all cases that a Call node should have a ret val
335 if assign_side == AssignSide.RIGHT:
336 node.has_ret_val = True
338 node.func.con_scope = enclosing_con_scope
339 # if we are trying to generate GrFN 2.2 and this call has an associated
340 # FunctionDef, make a GrFN 2.2 container for it
341 if self.pipeline_state.GENERATE_GRFN_2_2 and node.has_func_def:
342 node.is_grfn_2_2 = True
343 return self.visit_call_grfn_2_2(
344 node, base_func_scopestr, enclosing_con_scope, assign_side
345 )
347 # otherwise, this Call should not be treated as a GrFN 2.2 call,
348 # so we store a GrfnContainerSrcRef for it
349 grfn_src_ref = GrfnContainerSrcRef(None, None, None)
350 if node.source_refs is not None:
351 src_ref = combine_source_refs(node.source_refs)
352 grfn_src_ref = GrfnContainerSrcRef(
353 line_begin=src_ref.row_start,
354 line_end=src_ref.row_end,
355 source_file_name=src_ref.source_file_name,
356 )
357 node.grfn_con_src_ref = grfn_src_ref
359 # queue node to process globals through interfaces later if we have the associated FunctionDef
360 if node.has_func_def:
361 self.calls_to_process.append(node)
363 if isinstance(node.func, AnnCastAttribute):
364 self.visit(
365 node.func, base_func_scopestr, enclosing_con_scope, assign_side
366 )
368 # For a call, we do not care about the arguments source refs
369 return self.visit_node_list(
370 node.arguments,
371 base_func_scopestr,
372 enclosing_con_scope,
373 assign_side,
374 )
376 def visit_call_grfn_2_2(
377 self,
378 node: AnnCastCall,
379 base_func_scopestr,
380 enclosing_con_scope,
381 assign_side,
382 ):
383 assert isinstance(node.func, AnnCastName)
385 # the children GrFN source ref for the call node is the src ref of the call's arguments
386 args_src_ref = self.visit_node_list(
387 node.arguments,
388 base_func_scopestr,
389 enclosing_con_scope,
390 assign_side,
391 )
393 node.func_def_copy = copy.deepcopy(
394 self.pipeline_state.func_id_to_def[node.func.id]
395 )
396 # make a new id for the copy's Name node, and store in func_id_to_def
397 node.func_def_copy.name.id = self.pipeline_state.next_collapsed_id()
398 self.pipeline_state.func_id_to_def[
399 node.func_def_copy.name.id
400 ] = node.func_def_copy
401 calling_scope = enclosing_con_scope + [call_container_name(node)]
402 call_assign_side = AssignSide.NEITHER
403 self.visit_function_def(
404 node.func_def_copy,
405 base_func_scopestr,
406 calling_scope,
407 call_assign_side,
408 )
410 return args_src_ref
412 # FUTURE: decide how to handle a ClassDef's accessed, modified, and used variables
413 @_visit.register
414 def visit_record_def(
415 self,
416 node: AnnCastRecordDef,
417 base_func_scopestr,
418 enclosing_con_scope,
419 assign_side,
420 ):
421 # we believe the start of the container should not be on either side of an assignment
422 assert assign_side == AssignSide.NEITHER
423 # we do not visit the name because it is a string
424 assert isinstance(node.name, str)
425 classscope = enclosing_con_scope # + [node.name]
426 # NOTE:
427 # node.bases is a list of strings
428 # node.funcs is a list of Vars
429 # node.fields is a list of Vars
431 # ClassDef's reset the `base_func_scopestr`
432 base_scopestr = con_scope_to_str(classscope)
433 funcs_src_ref = self.visit_node_list(
434 node.funcs, base_scopestr, classscope, assign_side
435 )
436 fields_src_ref = self.visit_node_list(
437 node.fields, base_scopestr, classscope, assign_side
438 )
440 return combine_grfn_con_src_refs([funcs_src_ref, fields_src_ref])
442 @_visit.register
443 def visit_function_def(
444 self,
445 node: AnnCastFunctionDef,
446 base_func_scopestr,
447 enclosing_con_scope,
448 assign_side,
449 ):
450 # we believe the start of the container should not be on either side of an assignment
451 assert assign_side == AssignSide.NEITHER
452 # store GrfnContainerSrcRef for this function def
453 grfn_src_ref = GrfnContainerSrcRef(None, None, None)
454 if node.source_refs is not None:
455 src_ref = combine_source_refs(node.source_refs)
456 grfn_src_ref = GrfnContainerSrcRef(
457 line_begin=src_ref.row_start,
458 line_end=src_ref.row_end,
459 source_file_name=src_ref.source_file_name,
460 )
461 node.grfn_con_src_ref = grfn_src_ref
463 # Modify scope to include the function name
464 funcscope = enclosing_con_scope + [func_def_container_name(node)]
466 self.initialize_con_scope_data(funcscope, node)
467 node.con_scope = funcscope
468 # FunctionDef's reset the `base_func_scopestr`
469 base_scopestr = con_scope_to_str(funcscope)
471 # Cache function container scopestr for use during Variable Version pass
472 self.pipeline_state.func_con_scopestr_to_id[
473 base_scopestr
474 ] = node.name.id
476 # Each argument is a AnnCastVar node
477 # Initialize each Name and visit to modify its scope
478 args_src_ref = self.visit_node_list(
479 node.func_args, base_scopestr, funcscope, assign_side
480 )
482 body_src_ref = self.visit_node_list(
483 node.body, base_scopestr, funcscope, assign_side
484 )
486 # return children GrfnContainerSrcRef
487 return combine_grfn_con_src_refs([args_src_ref, body_src_ref])
489 @_visit.register
490 def visit_goto(
491 self,
492 node: AnnCastGoto,
493 base_func_scope_str,
494 enclosing_con_scope,
495 assign_side
496 ):
497 if node.expr != None:
498 self.visit(node.expr, base_func_scope_str, enclosing_con_scope, assign_side)
499 # self.visit(node.label, base_func_scope_str, enclosing_con_scope, assign_side)
501 @_visit.register
502 def visit_label(
503 self,
504 node: AnnCastLabel,
505 base_func_scope_str,
506 enclosing_con_scope,
507 assign_side
508 ):
509 # self.visit(node.label, base_func_scope_str, enclosing_con_scope, assign_side)
510 pass
512 @_visit.register
513 def visit_literal_value(
514 self,
515 node: AnnCastLiteralValue,
516 base_func_scopestr,
517 enclosing_con_scope,
518 assign_side,
519 ):
520 if node.value_type == "List[Any]":
521 # val has
522 # operator - string
523 # size - Var node or a LiteralValue node (for number)
524 # initial_value - LiteralValue node
525 val = node.value
527 # visit size's anncast name node
528 self.visit(
529 val.size, base_func_scopestr, enclosing_con_scope, assign_side
530 )
532 # List literal doesn't need to add any other changes
533 # to the anncast at this pass
534 elif node.value_type == StructureType.TUPLE: # or node.value_type == StructureType.LIST:
535 self.visit_node_list(node.value, base_func_scopestr, enclosing_con_scope, assign_side)
536 elif node.value_type == ScalarType.INTEGER:
537 pass
538 elif node.value_type == ScalarType.ABSTRACTFLOAT:
539 pass
540 pass
542 @_visit.register
543 def visit_loop(
544 self,
545 node: AnnCastLoop,
546 base_func_scopestr,
547 enclosing_con_scope,
548 assign_side,
549 ):
550 # we believe the start of the container should not be on either side of an assignment
551 # (NOTE: In the case of list/dict comprehensions, they could be on the right hand side)
552 assert (
553 assign_side == AssignSide.NEITHER
554 or assign_side == AssignSide.RIGHT
555 )
556 # store the base_func_scopestr for this container
557 node.base_func_scopestr = base_func_scopestr
559 loopscope = self.next_loop_scope(enclosing_con_scope)
560 # print(f"-----------{loopscope}-----------")
561 self.initialize_con_scope_data(loopscope, node)
562 node.con_scope = loopscope
564 if len(node.pre) > 0:
565 # Additional modifications to support the loop post body
566 loopprescope = loopscope + [LOOPPRE]
567 # self.initialize_con_scope_data(loopinitscope, node)
568 init_src_ref = self.visit_node_list(
569 node.pre, base_func_scopestr, loopprescope, assign_side
570 )
572 # we store an additional ContainerData for the loop expression, but
573 # we store the Loop node in `self.con_str_to_node`
574 loopexprscope = loopscope + [LOOPEXPR]
575 self.initialize_con_scope_data(loopexprscope, node)
576 expr_src_ref = self.visit(
577 node.expr, base_func_scopestr, loopexprscope, assign_side
578 )
580 loopbodyscope = loopscope + [LOOPBODY]
581 body_src_ref = self.visit_node_list(
582 node.body, base_func_scopestr, loopbodyscope, assign_side
583 )
585 if len(node.post) > 0:
586 # Additional modifications to support the loop post body
587 looppostscope = loopscope + [LOOPPOST]
588 # self.initialize_con_scope_data(loopinitscope, node)
589 post_src_ref = self.visit_node_list(
590 node.post, base_func_scopestr, looppostscope, assign_side
591 )
593 # store GrfnContainerSrcRef for this loop
594 if len(node.pre) > 0 and len(node.post) > 0:
595 node.grfn_con_src_ref = combine_grfn_con_src_refs(
596 [init_src_ref, expr_src_ref, body_src_ref, post_src_ref]
597 )
598 elif len(node.pre) > 0:
599 node.grfn_con_src_ref = combine_grfn_con_src_refs(
600 [init_src_ref, expr_src_ref, body_src_ref]
601 )
602 elif len(node.post) > 0:
603 node.grfn_con_src_ref = combine_grfn_con_src_refs(
604 [expr_src_ref, body_src_ref, post_src_ref]
605 )
606 else:
607 node.grfn_con_src_ref = combine_grfn_con_src_refs(
608 [expr_src_ref, body_src_ref]
609 )
610 # return the children GrfnContainerSrcRef
611 return node.grfn_con_src_ref
613 @_visit.register
614 def visit_model_break(
615 self,
616 node: AnnCastModelBreak,
617 base_func_scopestr,
618 enclosing_con_scope,
619 assign_side,
620 ):
621 pass
623 @_visit.register
624 def visit_model_continue(
625 self,
626 node: AnnCastModelContinue,
627 base_func_scopestr,
628 enclosing_con_scope,
629 assign_side,
630 ):
631 pass
633 @_visit.register
634 def visit_model_import(
635 self,
636 node: AnnCastModelImport,
637 base_func_scopestr,
638 enclosing_con_scope,
639 assign_side,
640 ):
641 pass
643 @_visit.register
644 def visit_model_if(
645 self,
646 node: AnnCastModelIf,
647 base_func_scopestr,
648 enclosing_con_scope,
649 assign_side,
650 ):
651 # we believe the start of the container should not be on either side of an assignment
652 # (NOTE: In the case of list/dict comprehensions, they could be on the right hand side)
653 assert (
654 assign_side == AssignSide.NEITHER
655 or assign_side == AssignSide.RIGHT
656 )
657 # store the base_func_scopestr for this container
658 node.base_func_scopestr = base_func_scopestr
659 # want orig enclosing
660 ifscope = self.next_if_scope(enclosing_con_scope)
661 self.initialize_con_scope_data(ifscope, node)
662 node.con_scope = ifscope
664 # we store an additional ContainerData for the if expression, but
665 # we store the ModelIf node in `self.con_str_to_node`
666 ifexprscope = ifscope + [IFEXPR]
667 self.initialize_con_scope_data(ifexprscope, node)
668 expr_src_ref = self.visit(
669 node.expr, base_func_scopestr, ifexprscope, assign_side
670 )
672 ifbodyscope = ifscope + [IFBODY]
673 body_src_ref = self.visit_node_list(
674 node.body, base_func_scopestr, ifbodyscope, assign_side
675 )
677 orelsebodyscope = ifscope + [ELSEBODY]
678 orelse_src_ref = self.visit_node_list(
679 node.orelse, base_func_scopestr, orelsebodyscope, assign_side
680 )
682 # store GrfnContainerSrcRef for this loop
683 node.grfn_con_src_ref = combine_grfn_con_src_refs(
684 [expr_src_ref, body_src_ref, orelse_src_ref]
685 )
686 # return the children GrfnContainerSrcRef
687 return node.grfn_con_src_ref
689 @_visit.register
690 def visit_return(
691 self,
692 node: AnnCastModelReturn,
693 base_func_scopestr,
694 enclosing_con_scope,
695 assign_side,
696 ):
697 # store the owning FunctionDef, and mark it as having a return value
698 function_def = self.pipeline_state.func_def_node_from_scopestr(
699 base_func_scopestr
700 )
701 node.owning_func_def = function_def
702 node.owning_func_def.has_ret_val = True
704 return self.visit(
705 node.value, base_func_scopestr, enclosing_con_scope, assign_side
706 )
708 @_visit.register
709 def visit_module(
710 self,
711 node: AnnCastModule,
712 base_func_scopestr,
713 enclosing_con_scope,
714 assign_side,
715 ):
716 # we believe the start of the container should not be on either side of an assignment
717 assert assign_side == AssignSide.NEITHER
718 module_con_scope = [MODULE_SCOPE]
719 node.con_scope = module_con_scope
720 # module resets the `base_func_scopestr`
721 base_scopestr = con_scope_to_str(module_con_scope)
722 # initialize container data for module which will store global variables
723 self.initialize_con_scope_data(module_con_scope, node)
724 body_src_ref = self.visit_node_list(
725 node.body, base_scopestr, module_con_scope, assign_side
726 )
728 # store GrfnContainerSrcRef for the module
729 node.grfn_con_src_ref = body_src_ref
730 # return the children GrfnContainerSrcRef
731 return node.grfn_con_src_ref
733 @_visit.register
734 def visit_name(
735 self,
736 node: AnnCastName,
737 base_func_scopestr,
738 enclosing_con_scope,
739 assign_side,
740 ):
741 node.con_scope = enclosing_con_scope
742 node.base_func_scopestr = base_func_scopestr
744 # check every prefix of enclosing_con_scope and add this Name node
745 # to the associated container data if either
746 # 1. the container scopestr extends base_func_scopestr
747 # 2. this Name node is a global variable
748 for index, name in enumerate(enclosing_con_scope):
749 # add separator between container scope component names
750 scopestr = CON_STR_SEP.join(enclosing_con_scope[: index + 1])
752 # if this Name node is a global, or if the scopestr extends base_func_scopestr
753 # we will add the node to scopestr's container data
754 # otherwise, we skip it
755 # we must do a compound check to propagate globals correctly
756 # we would like to stop propagation of variable use at base_func_scopestr, but
757 # this would only be correct for function locals. global use must be propagated above
758 # base_func_scopestr
759 if not (
760 self.pipeline_state.is_global_var(node.id)
761 or scopestr.startswith(base_func_scopestr)
762 ):
763 continue
765 # fill in container data if this is a cached container str
766 if scopestr in self.con_str_to_con_data:
767 con_data = self.con_str_to_con_data[scopestr]
768 # if we are on LHS of assignment, this Name should be
769 # added to modified vars
770 if assign_side == AssignSide.LEFT:
771 con_data.modified_vars[node.id] = node.name
772 # if this is the first time visiting the variable id in this scope,
773 # then it is accessed before modified
774 elif node.id not in con_data.used_vars:
775 con_data.vars_accessed_before_mod[node.id] = node.name
777 # for any type of use, add to containers used_vars
778 con_data.used_vars[node.id] = node.name
780 @_visit.register
781 def visit_operator(self, node: AnnCastOperator,
782 base_func_scopestr,
783 enclosing_con_scope,
784 assign_side,
785 ):
786 src_refs = []
788 # Visit each operand
789 for operand in node.operands:
790 src_refs.append(
791 self.visit(
792 operand, base_func_scopestr, enclosing_con_scope, assign_side
793 )
794 )
796 # NOTE: does this list need to be reversed?
797 if len(src_refs) > 1:
798 return combine_grfn_con_src_refs(src_refs)
799 # elif len(src_refs) == 1:
800 # return src_refs[0]
801 else:
802 return src_refs[0]
803 # return combine_grfn_con_src_refs([right_src_ref, left_src_ref])
805 @_visit.register
806 def visit_set(self, node: AnnCastSet, assign_side):
807 pass
809 @_visit.register
810 def visit_tuple(
811 self,
812 node: AnnCastTuple,
813 base_func_scopestr,
814 enclosing_con_scope,
815 assign_side,
816 ):
817 self.visit_node_list(
818 node.values, base_func_scopestr, enclosing_con_scope, assign_side
819 )
821 @_visit.register
822 def visit_var(
823 self,
824 node: AnnCastVar,
825 base_func_scopestr,
826 enclosing_con_scope,
827 assign_side,
828 ):
829 return self.visit(
830 node.val, base_func_scopestr, enclosing_con_scope, assign_side
831 )