# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Routines for fully traversing an IR."""

import inspect

from compiler.util import ir_pb2
from compiler.util import simple_memoizer


class _FunctionCaller:
  """Provides a template for setting up a generic call to a function.

  The function parameters are inspected at run-time to build up a set of valid
  and required arguments. When invoking the function unneccessary parameters
  will be trimmed out. If arguments are missing an assertion will be triggered.

  This is currently limited to functions that have at least one positional
  parameter.

  Example usage:
  ```
  def func_1(a, b, c=2): pass
  def func_2(a, d): pass
  caller_1 = _FunctionCaller(func_1)
  caller_2 = _FunctionCaller(func_2)
  generic_params = {"b": 2, "c": 3, "d": 4}

  # Equivalent of: func_1(a, b=2, c=3)
  caller_1.invoke(a, generic_params)

  # Equivalent of: func_2(a, d=4)
  caller_2.invoke(a, generic_params)
  """

  def __init__(self, function):
    self.function = function
    self.needs_filtering = True
    self.valid_arg_names = set()
    self.required_arg_names = set()

    argspec = inspect.getfullargspec(function)
    if argspec.varkw:
      # If the function accepts a kwargs parameter, then it will accept all
      # arguments.
      # Note: this isn't technically true if one of the keyword arguments has the
      # same name as one of the positional arguments.
      self.needs_filtering = False
    else:
      # argspec.args is a list of all parameter names excluding keyword only
      # args. The first element is our required positional_arg and should be
      # ignored.
      args = argspec.args[1:]
      self.valid_arg_names.update(args)

      # args.kwonlyargs gives us the list of keyword only args which are
      # also valid.
      self.valid_arg_names.update(argspec.kwonlyargs)

      # Required args are positional arguments that don't have defaults.
      # Keyword only args are always optional and can be ignored. Args with
      # defaults are the last elements of the argsepec.args list and should
      # be ignored.
      if argspec.defaults:
        # Trim the arguments with defaults.
        args = args[: -len(argspec.defaults)]
      self.required_arg_names.update(args)

  def invoke(self, positional_arg, keyword_args):
    """Invokes the function with the given args."""
    if self.needs_filtering:
      # Trim to just recognized args.
      matched_args = {
          k: v for k, v in keyword_args.items() if k in self.valid_arg_names
      }
      # Check if any required args are missing.
      missing_args = self.required_arg_names.difference(matched_args.keys())
      assert not missing_args, (
          f"Attempting to call '{self.function.__name__}'; "
          f"missing {missing_args} (have {set(keyword_args.keys())})"
      )
      keyword_args = matched_args

    return self.function(positional_arg, **keyword_args)


@simple_memoizer.memoize
def _memoized_caller(function):
  default_lambda_name = (lambda: None).__name__
  assert (
      callable(function) and not function.__name__ == default_lambda_name
  ), "For performance reasons actions must be defined as static functions"
  return _FunctionCaller(function)


def _call_with_optional_args(function, positional_arg, keyword_args):
  """Calls function with whatever keyword_args it will accept."""
  caller = _memoized_caller(function)
  return caller.invoke(positional_arg, keyword_args)


def _fast_traverse_proto_top_down(proto, incidental_actions, pattern,
                                  skip_descendants_of, action, parameters):
  """Traverses an IR, calling `action` on some nodes."""

  # Parameters are scoped to the branch of the tree, so make a copy here, before
  # any action or incidental_action can update them.
  parameters = parameters.copy()

  # If there is an incidental action for this node type, run it.
  if type(proto) in incidental_actions:  # pylint: disable=unidiomatic-typecheck
    for incidental_action in incidental_actions[type(proto)]:
      parameters.update(_call_with_optional_args(
          incidental_action, proto, parameters) or {})

  # If we are at the end of pattern, check to see if we should call action.
  if len(pattern) == 1:
    new_pattern = pattern
    if pattern[0] == type(proto):
      parameters.update(
          _call_with_optional_args(action, proto, parameters) or {})
  else:
    # Otherwise, if this node's type matches the head of pattern, recurse with
    # the tail of the pattern.
    if pattern[0] == type(proto):
      new_pattern = pattern[1:]
    else:
      new_pattern = pattern

  # If the current node's type is one of the types whose branch should be
  # skipped, then bail.  This has to happen after `action` is called, because
  # clients rely on being able to, e.g., get a callback for the "root"
  # Expression without getting callbacks for every sub-Expression.
  # pylint: disable=unidiomatic-typecheck
  if type(proto) in skip_descendants_of:
    return

  # Otherwise, recurse.  _FIELDS_TO_SCAN_BY_CURRENT_AND_TARGET tells us, given
  # the current node's type and the current target type, which fields to check.
  singular_fields, repeated_fields = _FIELDS_TO_SCAN_BY_CURRENT_AND_TARGET[
      type(proto), new_pattern[0]]
  for member_name in singular_fields:
    if proto.HasField(member_name):
      _fast_traverse_proto_top_down(getattr(proto, member_name),
                                    incidental_actions, new_pattern,
                                    skip_descendants_of, action, parameters)
  for member_name in repeated_fields:
    for array_element in getattr(proto, member_name):
      _fast_traverse_proto_top_down(array_element, incidental_actions,
                                    new_pattern, skip_descendants_of, action,
                                    parameters)


def _fields_to_scan_by_current_and_target():
  """Generates _FIELDS_TO_SCAN_BY_CURRENT_AND_TARGET."""
  # In order to avoid spending a *lot* of time just walking the IR, this
  # function sets up a dict that allows `_fast_traverse_proto_top_down()` to
  # skip traversing large portions of the IR, depending on what node types it is
  # targeting.
  #
  # Without this branch culling scheme, the Emboss front end (at time of
  # writing) spends roughly 70% (19s out of 31s) of its time just walking the
  # IR.  With branch culling, that goes down to 6% (0.7s out of 12.2s).

  # type_to_fields is a map of types to maps of field names to field types.
  # That is, type_to_fields[ir_pb2.Module]["type"] == ir_pb2.TypeDefinition.
  type_to_fields = {}

  # Later, we need to know which fields are singular and which are repeated,
  # because the access methods are not uniform.  This maps (type, field_name)
  # tuples to descriptor labels: type_fields_to_cardinality[ir_pb2.Module,
  # "type"] == ir_pb2.Repeated.
  type_fields_to_cardinality = {}

  # Fill out the above maps by recursively walking the IR type tree, starting
  # from the root.
  types_to_check = [ir_pb2.EmbossIr]
  while types_to_check:
    type_to_check = types_to_check.pop()
    if type_to_check in type_to_fields:
      continue
    fields = {}
    for field_name, field_type in type_to_check.field_specs.items():
      if issubclass(field_type.type, ir_pb2.Message):
        fields[field_name] = field_type.type
        types_to_check.append(field_type.type)
        type_fields_to_cardinality[type_to_check, field_name] = (
            field_type.__class__)
    type_to_fields[type_to_check] = fields

  # type_to_descendant_types is a map of all types that can be reached from a
  # particular type.  After the setup, type_to_descendant_types[ir_pb2.EmbossIr]
  # == set(<all types>) and type_to_descendant_types[ir_pb2.Reference] ==
  # {ir_pb2.CanonicalName, ir_pb2.Word, ir_pb2.Location} and
  # type_to_descendant_types[ir_pb2.Word] == set().
  #
  # The while loop basically ors in the known descendants of each known
  # descendant of each type until the dict stops changing, which is a bit
  # brute-force, but in practice only iterates a few times.
  type_to_descendant_types = {}
  for parent_type, field_map in type_to_fields.items():
    type_to_descendant_types[parent_type] = set(field_map.values())
  previous_map = {}
  while type_to_descendant_types != previous_map:
    # In order to check the previous iteration against the current iteration, it
    # is necessary to make a two-level copy.  Otherwise, the updates to the
    # values will also update previous_map's values, which causes the loop to
    # exit prematurely.
    previous_map = {k: set(v) for k, v in type_to_descendant_types.items()}
    for ancestor_type, descendents in previous_map.items():
      for descendent in descendents:
        type_to_descendant_types[ancestor_type] |= previous_map[descendent]

  # Finally, we have all of the information we need to make the map we really
  # want: given a current node type and a target node type, which fields should
  # be checked?  (This implicitly skips fields that *can't* contain the target
  # type.)
  fields_to_scan_by_current_and_target = {}
  for current_node_type in type_to_fields:
    for target_node_type in type_to_fields:
      singular_fields_to_scan = []
      repeated_fields_to_scan = []
      for field_name, field_type in type_to_fields[current_node_type].items():
        # If the target node type cannot contain another instance of itself, it
        # is still necessary to scan fields that have the actual target type.
        if (target_node_type == field_type or
            target_node_type in type_to_descendant_types[field_type]):
          # Singular and repeated fields go to different lists, so that they can
          # be handled separately.
          if (type_fields_to_cardinality[current_node_type, field_name] ==
              ir_pb2.Optional):
            singular_fields_to_scan.append(field_name)
          else:
            repeated_fields_to_scan.append(field_name)
      fields_to_scan_by_current_and_target[
          current_node_type, target_node_type] = (
              singular_fields_to_scan, repeated_fields_to_scan)
  return fields_to_scan_by_current_and_target


_FIELDS_TO_SCAN_BY_CURRENT_AND_TARGET = _fields_to_scan_by_current_and_target()

def _emboss_ir_action(ir):
  return {"ir": ir}

def _module_action(m):
  return {"source_file_name": m.source_file_name}

def _type_definition_action(t):
  return {"type_definition": t}

def _field_action(f):
  return {"field": f}

def fast_traverse_ir_top_down(ir, pattern, action, incidental_actions=None,
                              skip_descendants_of=(), parameters=None):
  """Traverses an IR from the top down, executing the given actions.

  `fast_traverse_ir_top_down` walks the given IR in preorder traversal,
  specifically looking for nodes whose path from the root of the tree matches
  `pattern`.  For every node which matches `pattern`, `action` will be called.

  `pattern` is just a list of node types.  For example, to execute `print` on
  every `ir_pb2.Word` in the IR:

      fast_traverse_ir_top_down(ir, [ir_pb2.Word], print)

  If more than one type is specified, then each one must be found inside the
  previous.  For example, to print only the Words inside of import statements:

      fast_traverse_ir_top_down(ir, [ir_pb2.Import, ir_pb2.Word], print)

  The optional arguments provide additional control.

  `skip_descendants_of` is a list of types that should be treated as if they are
  leaf nodes when they are encountered.  That is, traversal will skip any
  nodes with any ancestor node whose type is in `skip_descendants_of`.  For
  example, to `do_something` only on outermost `Expression`s:

      fast_traverse_ir_top_down(ir, [ir_pb2.Expression], do_something,
                                skip_descendants_of={ir_pb2.Expression})

  `parameters` specifies a dictionary of initial parameters which can be passed
  as arguments to `action` and `incidental_actions`.  Note that the parameters
  can be overridden for parts of the tree by `action` and `incidental_actions`.
  Parameters can be used to set an object which may be updated by `action`, such
  as a list of errors generated by some check in `action`:

      def check_structure(structure, errors):
        if structure_is_bad(structure):
          errors.append(error_for_structure(structure))

      errors = []
      fast_traverse_ir_top_down(ir, [ir_pb2.Structure], check_structure,
                                parameters={"errors": errors})
      if errors:
        print("Errors: {}".format(errors))
        sys.exit(1)

  `incidental_actions` is a map from node types to functions (or tuples of
  functions or lists of functions) which should be called on those nodes.
  Because `fast_traverse_ir_top_down` may skip branches that can't contain
  `pattern`, functions in `incidental_actions` should generally not have any
  side effects: instead, they may return a dictionary, which will be used to
  override `parameters` for any children of the node they were called on.  For
  example:

      def do_something(expression, field_name=None):
        if field_name:
          print("Found {} inside {}".format(expression, field_name))
        else:
          print("Found {} not in any field".format(expression))

      fast_traverse_ir_top_down(
          ir, [ir_pb2.Expression], do_something,
          incidental_actions={ir_pb2.Field: lambda f: {"field_name": f.name}})

  (The `action` may also return a dict in the same way.)

  A few `incidental_actions` are built into `fast_traverse_ir_top_down`, so
  that certain parameters are contextually available with well-known names:

      ir: The complete IR (the root ir_pb2.EmbossIr node).
      source_file_name: The file name from which the current node was sourced.
      type_definition: The most-immediate ancestor type definition.
      field: The field containing the current node, if any.

  Arguments:
    ir: An ir_pb2.Ir object to walk.
    pattern: A list of node types to match.
    action: A callable, which will be called on nodes matching `pattern`.
    incidental_actions: A dict of node types to callables, which can be used to
        set new parameters for `action` for part of the IR tree.
    skip_descendants_of: A list of types whose children should be skipped when
        traversing `ir`.
    parameters: A list of top-level parameters.

  Returns:
    None
  """
  all_incidental_actions = {
      ir_pb2.EmbossIr: [_emboss_ir_action],
      ir_pb2.Module: [_module_action],
      ir_pb2.TypeDefinition: [_type_definition_action],
      ir_pb2.Field: [_field_action],
  }
  if incidental_actions:
    for key, incidental_action in incidental_actions.items():
      if not isinstance(incidental_action, (list, tuple)):
        incidental_action = [incidental_action]
      all_incidental_actions.setdefault(key, []).extend(incidental_action)
  _fast_traverse_proto_top_down(ir, all_incidental_actions, pattern,
                                skip_descendants_of, action, parameters or {})


def fast_traverse_node_top_down(node, pattern, action, incidental_actions=None,
                                skip_descendants_of=(), parameters=None):
  """Traverse a subtree of an IR, executing the given actions.

  fast_traverse_node_top_down is like fast_traverse_ir_top_down, except that:

  It may be called on a subtree, instead of the top of the IR.

  It does not have any built-in incidental actions.

  Arguments:
    node: An ir_pb2.Ir object to walk.
    pattern: A list of node types to match.
    action: A callable, which will be called on nodes matching `pattern`.
    incidental_actions: A dict of node types to callables, which can be used to
        set new parameters for `action` for part of the IR tree.
    skip_descendants_of: A list of types whose children should be skipped when
        traversing `node`.
    parameters: A list of top-level parameters.

  Returns:
    None
  """
  _fast_traverse_proto_top_down(node, incidental_actions or {}, pattern,
                                skip_descendants_of or {}, action,
                                parameters or {})