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

1import copy 

2import typing 

3from collections import defaultdict 

4from enum import Enum 

5from functools import singledispatchmethod 

6 

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) 

31 

32 

33class AssignSide(Enum): 

34 NEITHER = 0 

35 LEFT = 1 

36 RIGHT = 2 

37 

38 

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] 

43 

44 def __init__(self): 

45 self.modified_vars = {} 

46 self.vars_accessed_before_mod = {} 

47 self.used_vars = {} 

48 

49 

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() 

62 

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 

71 

72 # add cached container data to container nodes 

73 self.add_container_data_to_nodes() 

74 

75 # save the dict mapping container scope to AnnCastNode 

76 self.pipeline_state.con_scopestr_to_node = self.con_str_to_node 

77 

78 self.propagate_globals_through_calls() 

79 

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}"] 

85 

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}"] 

91 

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 ) 

97 

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 

109 

110 if ( 

111 scopestr == MODULE_SCOPE 

112 or not self.pipeline_state.is_container(scopestr) 

113 ): 

114 continue 

115 

116 container_node = self.pipeline_state.con_node_from_scopestr( 

117 scopestr 

118 ) 

119 

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 ) 

133 

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 

141 

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) 

159 

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 

171 

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 

177 

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 

184 

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 

195 

196 # DEBUG printing 

197 if self.pipeline_state.PRINT_DEBUGGING_INFO: 

198 print(container.grfn_con_src_ref) 

199 

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() 

208 

209 # map con_scopestr to passed in node 

210 self.con_str_to_node[con_scopestr] = node 

211 

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}") 

223 

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) 

229 

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 ) 

241 

242 return combine_grfn_con_src_refs([children_src_ref, grfn_src_ref]) 

243 

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)}") 

261 

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) 

280 

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 ) 

305 

306 return combine_grfn_con_src_refs([right_src_ref, left_src_ref]) 

307 

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 ) 

320 

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 

337 

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 ) 

346 

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 

358 

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) 

362 

363 if isinstance(node.func, AnnCastAttribute): 

364 self.visit( 

365 node.func, base_func_scopestr, enclosing_con_scope, assign_side 

366 ) 

367 

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 ) 

375 

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) 

384 

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 ) 

392 

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 ) 

409 

410 return args_src_ref 

411 

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 

430 

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 ) 

439 

440 return combine_grfn_con_src_refs([funcs_src_ref, fields_src_ref]) 

441 

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 

462 

463 # Modify scope to include the function name 

464 funcscope = enclosing_con_scope + [func_def_container_name(node)] 

465 

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) 

470 

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 

475 

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 ) 

481 

482 body_src_ref = self.visit_node_list( 

483 node.body, base_scopestr, funcscope, assign_side 

484 ) 

485 

486 # return children GrfnContainerSrcRef 

487 return combine_grfn_con_src_refs([args_src_ref, body_src_ref]) 

488 

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) 

500 

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 

511 

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 

526 

527 # visit size's anncast name node 

528 self.visit( 

529 val.size, base_func_scopestr, enclosing_con_scope, assign_side 

530 ) 

531 

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 

541 

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 

558 

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 

563 

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 ) 

571 

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 ) 

579 

580 loopbodyscope = loopscope + [LOOPBODY] 

581 body_src_ref = self.visit_node_list( 

582 node.body, base_func_scopestr, loopbodyscope, assign_side 

583 ) 

584 

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 ) 

592 

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 

612 

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 

622 

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 

632 

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 

642 

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 

663 

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 ) 

671 

672 ifbodyscope = ifscope + [IFBODY] 

673 body_src_ref = self.visit_node_list( 

674 node.body, base_func_scopestr, ifbodyscope, assign_side 

675 ) 

676 

677 orelsebodyscope = ifscope + [ELSEBODY] 

678 orelse_src_ref = self.visit_node_list( 

679 node.orelse, base_func_scopestr, orelsebodyscope, assign_side 

680 ) 

681 

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 

688 

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 

703 

704 return self.visit( 

705 node.value, base_func_scopestr, enclosing_con_scope, assign_side 

706 ) 

707 

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 ) 

727 

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 

732 

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 

743 

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]) 

751 

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 

764 

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 

776 

777 # for any type of use, add to containers used_vars 

778 con_data.used_vars[node.id] = node.name 

779 

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 = [] 

787 

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 ) 

795 

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]) 

804 

805 @_visit.register 

806 def visit_set(self, node: AnnCastSet, assign_side): 

807 pass 

808 

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 ) 

820 

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 )