text: # Copyright 2017-present Kensho Technologies, LLC.
  """Functions that ensure the IR generated by the front-end satisfies all invariants."""
  from typing import Dict, List
  
  from funcy import pairwise
  
  from .blocks import (
      Backtrack,
      CoerceType,
      ConstructResult,
      Filter,
      Fold,
      MarkLocation,
      OutputSource,
      QueryRoot,
      Recurse,
      Traverse,
      Unfold,
  )
  from .compiler_entities import BasicBlock
  from .helpers import FoldScopeLocation
  from .ir_lowering_common.common import extract_folds_from_ir_blocks
  from .metadata import QueryMetadataTable
  
  
  def self_consistency_check_ir_blocks_from_frontend(
      ir_blocks: List[BasicBlock], query_metadata_table: QueryMetadataTable
  ) -> None:
      """Assert that IR blocks originating from the frontend do not have nonsensical structure.
  
      Args:
          ir_blocks: list of BasicBlocks representing the IR to self-consistency check.
          query_metadata_table: QueryMetadataTable object that captures information about the query.
  
      Raises:
          AssertionError, if the IR has unexpected structure. If the IR produced by the front-end
          cannot be successfully and correctly used to generate MATCH or Gremlin due to a bug,
          this is the method that should catch the problem.
      """
      if not ir_blocks:
          raise AssertionError("Received no ir_blocks: {}".format(ir_blocks))
  
      _assert_fold_scope_locations_are_unique(ir_blocks)
      _assert_no_nested_folds(ir_blocks)
      _assert_query_root_block(ir_blocks)
      _assert_output_source_follower_blocks(ir_blocks)
      _assert_block_pairwise_constraints(ir_blocks)
      _assert_mark_location_preceding_optional_traverse(ir_blocks)
      _assert_every_location_is_marked(ir_blocks)
      _assert_coerce_type_outside_of_fold(ir_blocks)
      _assert_all_marked_locations_are_registered(ir_blocks, query_metadata_table)
      _assert_registered_locations_parent_locations(query_metadata_table)
  
  
  def _assert_registered_locations_parent_locations(
      query_metadata_table: QueryMetadataTable,
  ) -> None:
      """Assert that all registered locations' parent locations are also registered."""
      for location, location_info in query_metadata_table.registered_locations:
          if (
              location != query_metadata_table.root_location
              and not query_metadata_table.root_location.is_revisited_at(location)
          ):
              # If the location is not the root location and is not a revisit of the root,
              # then it must have a parent location.
              if location_info.parent_location is None:
                  raise AssertionError(
                      "Found a location that is not the root location of the query "
                      "or a revisit of the root, but does not have a parent: "
                      "{} {}".format(location, location_info)
                  )
  
          if location_info.parent_location is not None:
              # Make sure the parent_location is also registered.
              # If the location is not registered, the following line will raise an error.
              query_metadata_table.get_location_info(location_info.parent_location)
  
  
  def _assert_all_marked_locations_are_registered(
      ir_blocks: List[BasicBlock], query_metadata_table: QueryMetadataTable
  ) -> None:
      """Assert that all locations in MarkLocation blocks have registered and valid metadata."""
      # Grab all the registered locations, then make sure that:
      # - Any location that appears in a MarkLocation block is also registered.
      # - There are no registered locations that do not appear in a MarkLocation block.
      registered_locations = {location for location, _ in query_metadata_table.registered_locations}
  
      ir_encountered_locations = {
          block.location for block in ir_blocks if isinstance(block, MarkLocation)
      }
  
      unregistered_locations = ir_encountered_locations - registered_locations
      unencountered_locations = registered_locations - ir_encountered_locations
      if unregistered_locations:
          raise AssertionError(
              "IR blocks unexpectedly contain locations not registered in the "
              "QueryMetadataTable: {}".format(unregistered_locations)
          )
      if unencountered_locations:
          raise AssertionError(
              "QueryMetadataTable unexpectedly contains registered locations that "
              "never appear in the IR blocks: {}".format(unencountered_locations)
          )
  
  
  def _assert_fold_scope_locations_are_unique(ir_blocks: List[BasicBlock]) -> None:
      """Assert that every FoldScopeLocation that exists on a Fold block is unique."""
      observed_locations: Dict[FoldScopeLocation, Fold] = dict()
      for block in ir_blocks:
          if isinstance(block, Fold):
              alternate = observed_locations.get(block.fold_scope_location, None)
              if alternate is not None:
                  raise AssertionError(
                      "Found two Fold blocks with identical FoldScopeLocations: "
                      "{} {} {}".format(alternate, block, ir_blocks)
                  )
              observed_locations[block.fold_scope_location] = block
  
  
  def _assert_no_nested_folds(ir_blocks: List[BasicBlock]) -> None:
      """Assert that there are no nested Fold contexts, and that every Fold has a matching Unfold."""
      fold_seen = False
      for block in ir_blocks:
          if isinstance(block, Fold):
              if fold_seen:
                  raise AssertionError("Found a nested Fold contexts: {}".format(ir_blocks))
              else:
                  fold_seen = True
          elif isinstance(block, Unfold):
              if not fold_seen:
                  raise AssertionError(
                      "Found an Unfold block without a matching Fold: {}".format(ir_blocks)
                  )
              else:
                  fold_seen = False
  
  
  def _assert_query_root_block(ir_blocks: List[BasicBlock]) -> None:
      """Assert that QueryRoot is always the first block, and only the first block."""
      if not isinstance(ir_blocks[0], QueryRoot):
          raise AssertionError("The first block was not QueryRoot: {}".format(ir_blocks))
      for block in ir_blocks[1:]:
          if isinstance(block, QueryRoot):
              raise AssertionError("Found QueryRoot after the first block: {}".format(ir_blocks))
  
  
  def _self_consistency_check_construct_result_block(ir_blocks: List[BasicBlock]) -> None:
      """Assert that ConstructResult is always the last block, and only the last block."""
      if not isinstance(ir_blocks[-1], ConstructResult):
          raise AssertionError("The last block was not ConstructResult: {}".format(ir_blocks))
      for block in ir_blocks[:-1]:
          if isinstance(block, ConstructResult):
              raise AssertionError(
                  "Found ConstructResult before the last block: {}".format(ir_blocks)
              )
  
  
  def _assert_output_source_follower_blocks(ir_blocks: List[BasicBlock]) -> None:
      """Ensure there are no Traverse / Backtrack / Recurse blocks after an OutputSource block."""
      seen_output_source = False
      for block in ir_blocks:
          if isinstance(block, OutputSource):
              seen_output_source = True
          elif seen_output_source:
              if isinstance(block, (Backtrack, Traverse, Recurse)):
                  raise AssertionError(
                      "Found Backtrack / Traverse / Recurse "
                      "after OutputSource block: "
                      "{}".format(ir_blocks)
                  )
  
  
  def _assert_block_pairwise_constraints(ir_blocks: List[BasicBlock]) -> None:
      """Assert that adjacent blocks obey all invariants."""
      for first_block, second_block in pairwise(ir_blocks):
          # Always Filter before MarkLocation, never after.
          if isinstance(first_block, MarkLocation) and isinstance(second_block, Filter):
              raise AssertionError("Found Filter after MarkLocation block: {}".format(ir_blocks))
  
          # There's no point in marking the same location twice in a row.
          if isinstance(first_block, MarkLocation) and isinstance(second_block, MarkLocation):
              raise AssertionError("Found consecutive MarkLocation blocks: {}".format(ir_blocks))
  
          # Traverse blocks with optional=True are immediately followed
          # by a MarkLocation, CoerceType or Filter block.
          if isinstance(first_block, Traverse) and first_block.optional:
              if not isinstance(second_block, (MarkLocation, CoerceType, Filter)):
                  raise AssertionError(
                      "Expected MarkLocation, CoerceType or Filter after Traverse "
                      "with optional=True. Found: {}".format(ir_blocks)
                  )
  
          # Backtrack blocks with optional=True are immediately followed by a MarkLocation block.
          if isinstance(first_block, Backtrack) and first_block.optional:
              if not isinstance(second_block, MarkLocation):
                  raise AssertionError(
                      "Expected MarkLocation after Backtrack with optional=True, "
                      "but none was found: {}".format(ir_blocks)
                  )
  
          # Recurse blocks are immediately preceded by a MarkLocation or Backtrack block.
          if isinstance(second_block, Recurse):
              if not (isinstance(first_block, MarkLocation) or isinstance(first_block, Backtrack)):
                  raise AssertionError(
                      "Expected MarkLocation or Backtrack before Recurse, but none "
                      "was found: {}".format(ir_blocks)
                  )
  
  
  def _assert_mark_location_preceding_optional_traverse(
      ir_blocks: List[BasicBlock],
  ) -> None:
      """Assert that optional Traverse blocks are preceded by a MarkLocation."""
      # Once all fold blocks are removed, each optional Traverse must have
      # a MarkLocation block immediately before it.
      _, new_ir_blocks = extract_folds_from_ir_blocks(ir_blocks)
      for first_block, second_block in pairwise(new_ir_blocks):
          # Traverse blocks with optional=True are immediately preceded by a MarkLocation block.
          if isinstance(second_block, Traverse) and second_block.optional:
              if not isinstance(first_block, MarkLocation):
                  raise AssertionError(
                      "Expected MarkLocation before Traverse with optional=True, "
                      "but none was found: {}".format(ir_blocks)
                  )
  
  
  def _assert_every_location_is_marked(ir_blocks: List[BasicBlock]) -> None:
      """Ensure that every new location is marked with a MarkLocation block."""
      # Exactly one MarkLocation block is found between any block that starts an interval of blocks
      # that all affect the same query position, and the first subsequent block that affects a
      # different position in the query. Such intervals include the following examples:
      # - from Fold to Unfold
      # - from QueryRoot to Traverse/Recurse
      # - from one Traverse to the next Traverse
      # - from Traverse to Backtrack
      found_start_block = False
      mark_location_blocks_count = 0
  
      start_interval_types = (QueryRoot, Traverse, Recurse, Fold)
      end_interval_types = (Backtrack, ConstructResult, Recurse, Traverse, Unfold)
  
      for block in ir_blocks:
          # Terminate started intervals before opening new ones.
          if isinstance(block, end_interval_types) and found_start_block:
              found_start_block = False
              if mark_location_blocks_count != 1:
                  raise AssertionError(
                      "Expected 1 MarkLocation block between traversals, found: "
                      "{} {}".format(mark_location_blocks_count, ir_blocks)
                  )
  
          # Now consider opening new intervals or processing MarkLocation blocks.
          if isinstance(block, MarkLocation):
              mark_location_blocks_count += 1
          elif isinstance(block, start_interval_types):
              found_start_block = True
              mark_location_blocks_count = 0
  
  
  def _assert_coerce_type_outside_of_fold(ir_blocks: List[BasicBlock]) -> None:
      """Ensure that CoerceType not in a @fold are followed by a MarkLocation or Filter block."""
      is_in_fold = False
      for first_block, second_block in pairwise(ir_blocks):
          if isinstance(first_block, Fold):
              is_in_fold = True
  
          if not is_in_fold and isinstance(first_block, CoerceType):
              if not isinstance(second_block, (MarkLocation, Filter)):
                  raise AssertionError(
                      "Expected MarkLocation or Filter after CoerceType, "
                      "but none was found: {}".format(ir_blocks)
                  )
  
          if isinstance(second_block, Unfold):
              is_in_fold = False
  