# 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.

"""Utility functions for reading and manipulating Emboss IR."""

import operator

from compiler.util import ir_pb2


_FIXED_SIZE_ATTRIBUTE = "fixed_size_in_bits"


def get_attribute(attribute_list, name):
  """Finds name in attribute_list and returns a AttributeValue or None."""
  attribute_value = None
  for attr in attribute_list:
    if attr.name.text == name and not attr.is_default:
      assert attribute_value is None, 'Duplicate attribute "{}".'.format(name)
      attribute_value = attr.value
  return attribute_value


def get_boolean_attribute(attribute_list, name, default_value=None):
  """Returns the boolean value of an attribute, if any, or default_value.

  Arguments:
      attribute_list: A list of attributes to search.
      name: The name of the desired attribute.
      default_value: A value to return if name is not found in attribute_list,
          or the attribute does not have a boolean value.

  Returns:
      The boolean value of the requested attribute, or default_value if the
      requested attribute is not found or has a non-boolean value.
  """
  attribute_value = get_attribute(attribute_list, name)
  if (not attribute_value or
      not attribute_value.expression.HasField("boolean_constant")):
    return default_value
  return attribute_value.expression.boolean_constant.value


def get_integer_attribute(attribute_list, name, default_value=None):
  """Returns the integer value of an attribute, if any, or default_value.

  Arguments:
      attribute_list: A list of attributes to search.
      name: The name of the desired attribute.
      default_value: A value to return if name is not found in attribute_list,
          or the attribute does not have an integer value.

  Returns:
      The integer value of the requested attribute, or default_value if the
      requested attribute is not found or has a non-integer value.
  """
  attribute_value = get_attribute(attribute_list, name)
  if (not attribute_value or
      attribute_value.expression.type.WhichOneof("type") != "integer" or
      not is_constant(attribute_value.expression)):
    return default_value
  return constant_value(attribute_value.expression)


def is_constant(expression, bindings=None):
  return constant_value(expression, bindings) is not None


def is_constant_type(expression_type):
  """Returns True if expression_type is inhabited by a single value."""
  return (expression_type.integer.modulus == "infinity" or
          expression_type.boolean.HasField("value") or
          expression_type.enumeration.HasField("value"))


def constant_value(expression, bindings=None):
  """Evaluates expression with the given bindings."""
  if expression.WhichOneof("expression") == "constant":
    return int(expression.constant.value)
  elif expression.WhichOneof("expression") == "constant_reference":
    # We can't look up the constant reference without the IR, but by the time
    # constant_value is called, the actual values should have been propagated to
    # the type information.
    if expression.type.WhichOneof("type") == "integer":
      assert expression.type.integer.modulus == "infinity"
      return int(expression.type.integer.modular_value)
    elif expression.type.WhichOneof("type") == "boolean":
      assert expression.type.boolean.HasField("value")
      return expression.type.boolean.value
    elif expression.type.WhichOneof("type") == "enumeration":
      assert expression.type.enumeration.HasField("value")
      return int(expression.type.enumeration.value)
    else:
      assert False, "Unexpected expression type {}".format(
          expression.type.WhichOneof("type"))
  elif expression.WhichOneof("expression") == "function":
    return _constant_value_of_function(expression.function, bindings)
  elif expression.WhichOneof("expression") == "field_reference":
    return None
  elif expression.WhichOneof("expression") == "boolean_constant":
    return expression.boolean_constant.value
  elif expression.WhichOneof("expression") == "builtin_reference":
    name = expression.builtin_reference.canonical_name.object_path[0]
    if bindings and name in bindings:
      return bindings[name]
    else:
      return None
  elif expression.WhichOneof("expression") is None:
    return None
  else:
    assert False, "Unexpected expression kind {}".format(
        expression.WhichOneof("expression"))


def _constant_value_of_function(function, bindings):
  """Returns the constant value of evaluating `function`, or None."""
  values = [constant_value(arg, bindings) for arg in function.args]
  # Expressions like `$is_statically_sized && 1 <= $static_size_in_bits <= 64`
  # should return False, not None, if `$is_statically_sized` is false, even
  # though `$static_size_in_bits` is unknown.
  #
  # The easiest way to allow this is to use a three-way logic chart for each;
  # specifically:
  #
  # AND:      True     False    Unknown
  #         +--------------------------
  # True    | True     False    Unknown
  # False   | False    False    False
  # Unknown | Unknown  False    Unknown
  #
  # OR:       True     False    Unknown
  #         +--------------------------
  # True    | True     True     True
  # False   | True     False    Unknown
  # Unknown | True     Unknown  Unknown
  #
  # This raises the question of just how many constant-from-nonconstant
  # expressions Emboss should support.  There are, after all, a vast number of
  # constant expression patterns built from non-constant subexpressions, such as
  # `0 * X` or `X == X` or `3 * X == X + X + X`.  I (bolms@) am not implementing
  # any further special cases because I do not see any practical use for them.
  if function.function == ir_pb2.Function.AND:
    if any(value is False for value in values):
      return False
    elif any(value is None for value in values):
      return None
    else:
      return True
  elif function.function == ir_pb2.Function.OR:
    if any(value is True for value in values):
      return True
    elif any(value is None for value in values):
      return None
    else:
      return False
  elif function.function == ir_pb2.Function.CHOICE:
    if values[0] is None:
      return None
    else:
      return values[1] if values[0] else values[2]
  # Other than the logical operators and choice operator, the result of any
  # function on an unknown value is, itself, considered unknown.
  if any(value is None for value in values):
    return None
  functions = {
      ir_pb2.Function.ADDITION: operator.add,
      ir_pb2.Function.SUBTRACTION: operator.sub,
      ir_pb2.Function.MULTIPLICATION: operator.mul,
      ir_pb2.Function.EQUALITY: operator.eq,
      ir_pb2.Function.INEQUALITY: operator.ne,
      ir_pb2.Function.LESS: operator.lt,
      ir_pb2.Function.LESS_OR_EQUAL: operator.le,
      ir_pb2.Function.GREATER: operator.gt,
      ir_pb2.Function.GREATER_OR_EQUAL: operator.ge,
      # Python's max([1, 2]) == 2; max(1, 2) == 2; max([1]) == 1; but max(1)
      # throws a TypeError ("'int' object is not iterable").
      ir_pb2.Function.MAXIMUM: lambda *x: max(x),
  }
  return functions[function.function](*values)


def _hashable_form_of_name(name):
  return (name.module_file,) + tuple(name.object_path)


def hashable_form_of_reference(reference):
  """Returns a representation of reference that can be used as a dict key.

  Arguments:
    reference: An ir_pb2.Reference or ir_pb2.NameDefinition.

  Returns:
    A tuple of the module_file and object_path.
  """
  return _hashable_form_of_name(reference.canonical_name)


def hashable_form_of_field_reference(field_reference):
  """Returns a representation of field_reference that can be used as a dict key.

  Arguments:
    field_reference: An ir_pb2.FieldReference

  Returns:
    A tuple of tuples of the module_files and object_paths.
  """
  return tuple(_hashable_form_of_name(reference.canonical_name)
               for reference in field_reference.path)


def is_array(type_ir):
  """Returns true if type_ir is an array type."""
  return type_ir.HasField("array_type")


def _find_path_in_structure_field(path, field):
  if not path:
    return field
  return None


def _find_path_in_structure(path, type_definition):
  for field in type_definition.structure.field:
    if field.name.name.text == path[0]:
      return _find_path_in_structure_field(path[1:], field)
  return None


def _find_path_in_enumeration(path, type_definition):
  if len(path) != 1:
    return None
  for value in type_definition.enumeration.value:
    if value.name.name.text == path[0]:
      return value
  return None


def _find_path_in_parameters(path, type_definition):
  if len(path) > 1:
    return None
  for parameter in type_definition.runtime_parameter:
    if parameter.name.name.text == path[0]:
      return parameter
  return None


def _find_path_in_type_definition(path, type_definition):
  """Finds the object with the given path in the given type_definition."""
  if not path:
    return type_definition
  obj = _find_path_in_parameters(path, type_definition)
  if obj:
    return obj
  if type_definition.HasField("structure"):
    obj = _find_path_in_structure(path, type_definition)
  elif type_definition.HasField("enumeration"):
    obj = _find_path_in_enumeration(path, type_definition)
  if obj:
    return obj
  else:
    return _find_path_in_type_list(path, type_definition.subtype)


def _find_path_in_type_list(path, type_list):
  for type_definition in type_list:
    if type_definition.name.name.text == path[0]:
      return _find_path_in_type_definition(path[1:], type_definition)
  return None


def _find_path_in_module(path, module_ir):
  if not path:
    return module_ir
  return _find_path_in_type_list(path, module_ir.type)


def find_object_or_none(name, ir):
  """Finds the object with the given canonical name, if it exists.."""
  if (isinstance(name, ir_pb2.Reference) or
      isinstance(name, ir_pb2.NameDefinition)):
    path = _hashable_form_of_name(name.canonical_name)
  elif isinstance(name, ir_pb2.CanonicalName):
    path = _hashable_form_of_name(name)
  else:
    path = name

  for module in ir.module:
    if module.source_file_name == path[0]:
      return _find_path_in_module(path[1:], module)

  return None


def find_object(name, ir):
  """Finds the IR of the type, field, or value with the given canonical name."""
  result = find_object_or_none(name, ir)
  assert result is not None, "Bad reference {}".format(name)
  return result


def find_parent_object(name, ir):
  """Finds the parent object of the object with the given canonical name."""
  if (isinstance(name, ir_pb2.Reference) or
      isinstance(name, ir_pb2.NameDefinition)):
    path = _hashable_form_of_name(name.canonical_name)
  elif isinstance(name, ir_pb2.CanonicalName):
    path = _hashable_form_of_name(name)
  else:
    path = name
  return find_object(path[:-1], ir)


def get_base_type(type_ir):
  """Returns the base type of the given type.

  Arguments:
    type_ir: IR of a type reference.

  Returns:
    If type_ir corresponds to an atomic type (like "UInt"), returns type_ir.  If
    type_ir corresponds to an array type (like "UInt:8[12]" or "Square[8][8]"),
    returns the type after stripping off the array types ("UInt" or "Square").
  """
  while type_ir.HasField("array_type"):
    type_ir = type_ir.array_type.base_type
  assert type_ir.HasField("atomic_type"), (
      "Unknown kind of type {}".format(type_ir))
  return type_ir


def fixed_size_of_type_in_bits(type_ir, ir):
  """Returns the fixed, known size for the given type, in bits, or None.

  Arguments:
    type_ir: The IR of a type.
    ir: A complete IR, used to resolve references to types.

  Returns:
    size if the size of the type can be determined, otherwise None.
  """
  array_multiplier = 1
  while type_ir.HasField("array_type"):
    if type_ir.array_type.WhichOneof("size") == "automatic":
      return None
    else:
      assert type_ir.array_type.WhichOneof("size") == "element_count", (
          'Expected array size to be "automatic" or "element_count".')
    element_count = type_ir.array_type.element_count
    if not is_constant(element_count):
      return None
    else:
      array_multiplier *= constant_value(element_count)
    assert not type_ir.HasField("size_in_bits"), (
        "TODO(bolms): implement explicitly-sized arrays")
    type_ir = type_ir.array_type.base_type
  assert type_ir.HasField("atomic_type"), "Unexpected type!"
  if type_ir.HasField("size_in_bits"):
    size = constant_value(type_ir.size_in_bits)
  else:
    type_definition = find_object(type_ir.atomic_type.reference, ir)
    size_attr = get_attribute(type_definition.attribute, _FIXED_SIZE_ATTRIBUTE)
    if not size_attr:
      return None
    size = constant_value(size_attr.expression)
  return size * array_multiplier


def field_is_virtual(field_ir):
  """Returns true if the field is virtual."""
  # TODO(bolms): Should there be a more explicit indicator that a field is
  # virtual?
  return not field_ir.HasField("location")


def field_is_read_only(field_ir):
  """Returns true if the field is read-only."""
  # For now, all virtual fields are read-only, and no non-virtual fields are
  # read-only.
  return field_ir.write_method.read_only