The function definition visiting method for inline functions in tir. Parameters ---------- self : Parser The visiting parser. node : doc.FunctionDef The doc AST function definition node.
(self: Parser, node: doc.FunctionDef)
| 653 | |
| 654 | @dispatch.register(token="tir.inline", type_name="FunctionDef") |
| 655 | def visit_inline_function_def(self: Parser, node: doc.FunctionDef) -> None: |
| 656 | """The function definition visiting method for inline functions in tir. |
| 657 | |
| 658 | Parameters |
| 659 | ---------- |
| 660 | self : Parser |
| 661 | The visiting parser. |
| 662 | |
| 663 | node : doc.FunctionDef |
| 664 | The doc AST function definition node. |
| 665 | """ |
| 666 | # remove the inline decorator |
| 667 | node.decorator_list.pop() |
| 668 | # adjust the node location to the source code location |
| 669 | node.lineno += self.diag.source.start_line - 1 |
| 670 | node.col_offset += self.diag.source.start_column + 1 |
| 671 | node.end_lineno += self.diag.source.start_line - 1 |
| 672 | node.end_col_offset += self.diag.source.start_column + 1 |
| 673 | |
| 674 | # Record definition depth for LEGB late binding |
| 675 | definition_depth = len(self.var_table.frames) |
| 676 | |
| 677 | def get_func(): |
| 678 | func_ast = from_doc(node) |
| 679 | module_ast = ast.Module(body=[func_ast], type_ignores=[]) |
| 680 | ast.fix_missing_locations(module_ast) |
| 681 | # set the filename to the source name, so that the error message can be reported correctly |
| 682 | code_obj = compile(module_ast, filename=self.diag.source.source_name, mode="exec") |
| 683 | namespace = self.var_table.get() |
| 684 | exec(code_obj, namespace) # pylint: disable=exec-used |
| 685 | func_name = func_ast.name |
| 686 | func = namespace[func_name] |
| 687 | return func, func_name |
| 688 | |
| 689 | func, func_name = get_func() |
| 690 | wrapper = inline(func, definition_depth=definition_depth, defining_var_table=self.var_table) |
| 691 | |
| 692 | self.var_table.add(func_name, wrapper, allow_shadowing=False) |
| 693 | return None |
| 694 | |
| 695 | |
| 696 | @dispatch.register(token="tirx", type_name="tvm_annotation") |