def prepare_for_rebinding(comp): """Prepares `comp` for extracting rebound variables. Currently, this means replacing all called lambdas and inlining all blocks. This does not necessarly guarantee that the resulting computation has no called lambdas, it merely reduces a level of indirection here. This reduction has proved sufficient for identifying variables which are about to be rebound in the top-level lambda, necessarily when compiler components factor work out from a single function into multiple functions. Since this function makes no guarantees about sufficiency, it is the responsibility of the caller to ensure that no unbound variables are introduced during the rebinding. Args: comp: Instance of `building_blocks.ComputationBuildingBlock` from which all occurrences of a given variable need to be extracted and rebound. Returns: Another instance of `building_blocks.ComputationBuildingBlock` which has had all called lambdas replaced by blocks, all blocks inlined and all selections from tuples collapsed. """ # TODO(b/146430051): Follow up here and consider removing or enforcing more # strict output invariants when `remove_lambdas_and_blocks` is moved in here. py_typecheck.check_type(comp, building_blocks.ComputationBuildingBlock) comp, _ = transformations.uniquify_reference_names(comp) comp, _ = transformations.replace_called_lambda_with_block(comp) block_inliner = transformations.InlineBlock(comp) selection_replacer = transformations.ReplaceSelectionFromTuple() transforms = [block_inliner, selection_replacer] symbol_tree = transformation_utils.SymbolTree( transformation_utils.ReferenceCounter) def _transform_fn(comp, symbol_tree): """Transform function chaining inlining and collapsing selections.""" modified = False for transform in transforms: if transform.global_transform: comp, transform_modified = transform.transform( comp, symbol_tree) else: comp, transform_modified = transform.transform(comp) modified = modified or transform_modified return comp, modified return transformation_utils.transform_postorder_with_symbol_bindings( comp, _transform_fn, symbol_tree)
def remove_lambdas_and_blocks(comp): """Removes any called lambdas and blocks from `comp`. This function will rename all the variables in `comp` in a single walk of the AST, then replace called lambdas with blocks in another walk, since this transformation interacts with scope in delicate ways. It will chain inlining the blocks and collapsing the selection-from-tuple pattern together into a final pass. Args: comp: Instance of `building_blocks.ComputationBuildingBlock` from which we want to remove called lambdas and blocks. Returns: A transformed version of `comp` which has no called lambdas or blocks, and no extraneous selections from tuples. """ py_typecheck.check_type(comp, building_blocks.ComputationBuildingBlock) # TODO(b/146904968): In general, any bounded number of passes of these # transforms as currently implemented is insufficient in order to satisfy # the purpose of this function. Filing a new bug to followup if this becomes a # pressing issue. modified = False for fn in [ transformations.remove_unused_block_locals, transformations.inline_selections_from_tuple, transformations.replace_called_lambda_with_block, ] * 2: comp, inner_modified = fn(comp) modified = inner_modified or modified for fn in [ transformations.remove_unused_block_locals, transformations.uniquify_reference_names, ]: comp, inner_modified = fn(comp) modified = inner_modified or modified block_inliner = transformations.InlineBlock(comp) selection_replacer = transformations.ReplaceSelectionFromTuple() transforms = [block_inliner, selection_replacer] def _transform_fn(comp, symbol_tree): """Transform function chaining inlining and collapsing selections. This function is inlined here as opposed to factored out and parameterized by the transforms to apply, due to the delicacy of chaining transformations which rely on state. These transformations should be safe if they appear first in the list of transforms, but due to the difficulty of reasoning about the invariants the transforms can rely on in this setting, there is no function exposed which hoists out the internal logic. Args: comp: Instance of `building_blocks.ComputationBuildingBlock` we wish to check for inlining and collapsing of selections. symbol_tree: Instance of `building_blocks.SymbolTree` defining the bindings available to `comp`. Returns: A transformed version of `comp`. """ modified = False for transform in transforms: if transform.global_transform: comp, transform_modified = transform.transform(comp, symbol_tree) else: comp, transform_modified = transform.transform(comp) modified = modified or transform_modified return comp, modified symbol_tree = transformation_utils.SymbolTree( transformation_utils.ReferenceCounter) transformed_comp, inner_modified = transformation_utils.transform_postorder_with_symbol_bindings( comp, _transform_fn, symbol_tree) modified = modified or inner_modified return transformed_comp, modified