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

"""Tests for util.traverse_ir."""

import collections

import unittest

from compiler.util import ir_pb2
from compiler.util import traverse_ir

_EXAMPLE_IR = ir_pb2.EmbossIr.from_json("""{
"module": [
  {
    "type": [
      {
        "structure": {
          "field": [
            {
              "location": {
                "start": { "constant": { "value": "0" } },
                "size": { "constant": { "value": "8" } }
              },
              "type": {
                "atomic_type": {
                  "reference": {
                    "canonical_name": {
                      "module_file": "",
                      "object_path": ["UInt"]
                    }
                  }
                }
              },
              "name": { "name": { "text": "field1" } }
            },
            {
              "location": {
                "start": { "constant": { "value": "8" } },
                "size": { "constant": { "value": "16" } }
              },
              "type": {
                "array_type": {
                  "base_type": {
                    "atomic_type": {
                      "reference": {
                        "canonical_name": {
                          "module_file": "",
                          "object_path": ["UInt"]
                        }
                      }
                    }
                  },
                  "element_count": { "constant": { "value": "8" } }
                }
              },
              "name": { "name": { "text": "field2" } }
            }
          ]
        },
        "name": { "name": { "text": "Foo" } },
        "subtype": [
          {
            "structure": {
              "field": [
                {
                  "location": {
                    "start": { "constant": { "value": "24" } },
                    "size": { "constant": { "value": "32" } }
                  },
                  "type": {
                    "atomic_type": {
                      "reference": {
                        "canonical_name": {
                          "module_file": "",
                          "object_path": ["UInt"]
                        }
                      }
                    }
                  },
                  "name": { "name": { "text": "bar_field1" } }
                },
                {
                  "location": {
                    "start": { "constant": { "value": "32" } },
                    "size": { "constant": { "value": "320" } }
                  },
                  "type": {
                    "array_type": {
                      "base_type": {
                        "array_type": {
                          "base_type": {
                            "atomic_type": {
                              "reference": {
                                "canonical_name": {
                                  "module_file": "",
                                  "object_path": ["UInt"]
                                }
                              }
                            }
                          },
                          "element_count": { "constant": { "value": "16" } }
                        }
                      },
                      "automatic": { }
                    }
                  },
                  "name": { "name": { "text": "bar_field2" } }
                }
              ]
            },
            "name": { "name": { "text": "Bar" } }
          }
        ]
      },
      {
        "enumeration": {
          "value": [
            {
              "name": { "name": { "text": "ONE" } },
              "value": { "constant": { "value": "1" } }
            },
            {
              "name": { "name": { "text": "TWO" } },
              "value": {
                "function": {
                  "function": "ADDITION",
                  "args": [
                    { "constant": { "value": "1" } },
                    { "constant": { "value": "1" } }
                  ],
                  "function_name": { "text": "+" }
                }
              }
            }
          ]
        },
        "name": { "name": { "text": "Bar" } }
      }
    ],
    "source_file_name": "t.emb"
  },
  {
    "type": [
      {
        "external": { },
        "name": {
          "name": { "text": "UInt" },
          "canonical_name": { "module_file": "", "object_path": ["UInt"] }
        },
        "attribute": [
          {
            "name": { "text": "statically_sized" },
            "value": { "expression": { "boolean_constant": { "value": true } } }
          },
          {
            "name": { "text": "size_in_bits" },
            "value": { "expression": { "constant": { "value": "64" } } }
          }
        ]
      }
    ],
    "source_file_name": ""
  }
]
}""")


def _count_entries(sequence):
  counts = collections.Counter()
  for entry in sequence:
    counts[entry] += 1
  return counts


def _record_constant(constant, constant_list):
  constant_list.append(int(constant.value))


def _record_field_name_and_constant(constant, constant_list, field):
  constant_list.append((field.name.name.text, int(constant.value)))


def _record_file_name_and_constant(constant, constant_list, source_file_name):
  constant_list.append((source_file_name, int(constant.value)))


def _record_location_parameter_and_constant(constant, constant_list,
                                            location=None):
  constant_list.append((location, int(constant.value)))


def _record_kind_and_constant(constant, constant_list, type_definition):
  if type_definition.HasField("enumeration"):
    constant_list.append(("enumeration", int(constant.value)))
  elif type_definition.HasField("structure"):
    constant_list.append(("structure", int(constant.value)))
  elif type_definition.HasField("external"):
    constant_list.append(("external", int(constant.value)))
  else:
    assert False, "Shouldn't be here."


class TraverseIrTest(unittest.TestCase):

  def test_filter_on_type(self):
    constants = []
    traverse_ir.fast_traverse_ir_top_down(
        _EXAMPLE_IR, [ir_pb2.NumericConstant], _record_constant,
        parameters={"constant_list": constants})
    self.assertEqual(
        _count_entries([0, 8, 8, 8, 16, 24, 32, 16, 32, 320, 1, 1, 1, 64]),
        _count_entries(constants))

  def test_filter_on_type_in_type(self):
    constants = []
    traverse_ir.fast_traverse_ir_top_down(
        _EXAMPLE_IR,
        [ir_pb2.Function, ir_pb2.Expression, ir_pb2.NumericConstant],
        _record_constant,
        parameters={"constant_list": constants})
    self.assertEqual([1, 1], constants)

  def test_filter_on_type_star_type(self):
    struct_constants = []
    traverse_ir.fast_traverse_ir_top_down(
        _EXAMPLE_IR, [ir_pb2.Structure, ir_pb2.NumericConstant],
        _record_constant,
        parameters={"constant_list": struct_constants})
    self.assertEqual(_count_entries([0, 8, 8, 8, 16, 24, 32, 16, 32, 320]),
                     _count_entries(struct_constants))
    enum_constants = []
    traverse_ir.fast_traverse_ir_top_down(
        _EXAMPLE_IR, [ir_pb2.Enum, ir_pb2.NumericConstant], _record_constant,
        parameters={"constant_list": enum_constants})
    self.assertEqual(_count_entries([1, 1, 1]), _count_entries(enum_constants))

  def test_filter_on_not_type(self):
    notstruct_constants = []
    traverse_ir.fast_traverse_ir_top_down(
        _EXAMPLE_IR, [ir_pb2.NumericConstant], _record_constant,
        skip_descendants_of=(ir_pb2.Structure,),
        parameters={"constant_list": notstruct_constants})
    self.assertEqual(_count_entries([1, 1, 1, 64]),
                     _count_entries(notstruct_constants))

  def test_field_is_populated(self):
    constants = []
    traverse_ir.fast_traverse_ir_top_down(
        _EXAMPLE_IR, [ir_pb2.Field, ir_pb2.NumericConstant],
        _record_field_name_and_constant,
        parameters={"constant_list": constants})
    self.assertEqual(_count_entries([
        ("field1", 0), ("field1", 8), ("field2", 8), ("field2", 8),
        ("field2", 16), ("bar_field1", 24), ("bar_field1", 32),
        ("bar_field2", 16), ("bar_field2", 32), ("bar_field2", 320)
    ]), _count_entries(constants))

  def test_file_name_is_populated(self):
    constants = []
    traverse_ir.fast_traverse_ir_top_down(
        _EXAMPLE_IR, [ir_pb2.NumericConstant], _record_file_name_and_constant,
        parameters={"constant_list": constants})
    self.assertEqual(_count_entries([
        ("t.emb", 0), ("t.emb", 8), ("t.emb", 8), ("t.emb", 8), ("t.emb", 16),
        ("t.emb", 24), ("t.emb", 32), ("t.emb", 16), ("t.emb", 32),
        ("t.emb", 320), ("t.emb", 1), ("t.emb", 1), ("t.emb", 1), ("", 64)
    ]), _count_entries(constants))

  def test_type_definition_is_populated(self):
    constants = []
    traverse_ir.fast_traverse_ir_top_down(
        _EXAMPLE_IR, [ir_pb2.NumericConstant], _record_kind_and_constant,
        parameters={"constant_list": constants})
    self.assertEqual(_count_entries([
        ("structure", 0), ("structure", 8), ("structure", 8), ("structure", 8),
        ("structure", 16), ("structure", 24), ("structure", 32),
        ("structure", 16), ("structure", 32), ("structure", 320),
        ("enumeration", 1), ("enumeration", 1), ("enumeration", 1),
        ("external", 64)
    ]), _count_entries(constants))

  def test_keyword_args_dict_in_action(self):
    call_counts = {"populated": 0, "not": 0}

    def check_field_is_populated(node, **kwargs):
      del node  # Unused.
      self.assertTrue(kwargs["field"])
      call_counts["populated"] += 1

    def check_field_is_not_populated(node, **kwargs):
      del node  # Unused.
      self.assertFalse("field" in kwargs)
      call_counts["not"] += 1

    traverse_ir.fast_traverse_ir_top_down(
        _EXAMPLE_IR, [ir_pb2.Field, ir_pb2.Type], check_field_is_populated)
    self.assertEqual(7, call_counts["populated"])

    traverse_ir.fast_traverse_ir_top_down(
        _EXAMPLE_IR, [ir_pb2.Enum, ir_pb2.EnumValue],
        check_field_is_not_populated)
    self.assertEqual(2, call_counts["not"])

  def test_pass_only_to_sub_nodes(self):
    constants = []

    def pass_location_down(field):
      return {
          "location": (int(field.location.start.constant.value),
                       int(field.location.size.constant.value))
      }

    traverse_ir.fast_traverse_ir_top_down(
        _EXAMPLE_IR, [ir_pb2.NumericConstant],
        _record_location_parameter_and_constant,
        incidental_actions={ir_pb2.Field: pass_location_down},
        parameters={"constant_list": constants, "location": None})
    self.assertEqual(_count_entries([
        ((0, 8), 0), ((0, 8), 8), ((8, 16), 8), ((8, 16), 8), ((8, 16), 16),
        ((24, 32), 24), ((24, 32), 32), ((32, 320), 16), ((32, 320), 32),
        ((32, 320), 320), (None, 1), (None, 1), (None, 1), (None, 64)
    ]), _count_entries(constants))


if __name__ == "__main__":
  unittest.main()