compiler/util/ir_data_utils.py

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

"""Provides a helpers for working with IR data elements.

builder
-------
Instead of:
```
def set_function_name_end(function: Function):
  if not function.function_name:
    function.function_name = Word()
  if not function.function_name.source_location:
    function.function_name.source_location = Location()
  word.source_location.end = Position(line=1,column=2)
```

We can do:
```
def set_function_name_end(function: Function):
  builder(function).function_name.source_location.end = Position(line=1,
  column=2)
```

reader
------
Instead of:
```
def is_leaf_synthetic(data):
  if data:
    if data.attribute:
      if data.attribute.value:
        if data.attribute.value.is_synthetic is not None:
          return data.attribute.value.is_synthetic
  return False
```
We can do:
```
def is_leaf_synthetic(data):
  return reader(data).attribute.value.is_synthetic
```

IrDataSerializer
----------------
Provides methods for serializing and deserializing an IR data object.
"""
import enum
import json
from typing import (
    Any,
    Callable,
    Generic,
    MutableMapping,
    MutableSequence,
    Optional,
    Tuple,
    TypeVar,
    cast,
)

from compiler.util import ir_data
from compiler.util import ir_data_fields


T = TypeVar("T", bound=ir_data.Message)


def field_specs(ir: T | type[T]):
  """Retrieves the field specs for the IR data class"""
  data_type = ir if isinstance(ir, type) else type(ir)
  return ir_data_fields.IrDataclassSpecs.get_specs(data_type).all_field_specs


class IrDataSerializer:
  """Provides methods for serializing IR data objects"""

  def __init__(self, ir: T):
    assert ir is not None
    self.ir = ir

  def _to_dict(
      self,
      ir: T,
      field_func: Callable[[T], list[Tuple[ir_data_fields.FieldSpec, Any]]],
  ) -> MutableMapping[str, Any]:
    assert ir is not None
    values: MutableMapping[str, Any] = {}
    for spec, value in field_func(ir):
      if value is not None and spec.is_dataclass:
        if spec.is_sequence:
          value = [self._to_dict(v, field_func) for v in value]
        else:
          value = self._to_dict(value, field_func)
      values[spec.name] = value
    return values

  def to_dict(self, exclude_none: bool = False):
    """Converts the IR data class to a dictionary."""

    def non_empty(ir):
      return fields_and_values(
          ir, lambda v: v is not None and (not isinstance(v, list) or len(v))
      )

    def all_fields(ir):
      return fields_and_values(ir)

    # It's tempting to use `dataclasses.asdict` here, but that does a deep
    # copy which is overkill for the current usage; mainly as an intermediary
    # for `to_json` and `repr`.
    return self._to_dict(self.ir, non_empty if exclude_none else all_fields)

  def to_json(self, *args, **kwargs):
    """Converts the IR data class to a JSON string"""
    return json.dumps(self.to_dict(exclude_none=True), *args, **kwargs)

  @staticmethod
  def from_json(data_cls, data):
    """Constructs an IR data class from the given JSON string"""
    as_dict = json.loads(data)
    return IrDataSerializer.from_dict(data_cls, as_dict)

  def copy_from_dict(self, data):
    """Deserializes the data and overwrites the IR data class with it"""
    cls = type(self.ir)
    data_copy = IrDataSerializer.from_dict(cls, data)
    for k in field_specs(cls):
      setattr(self.ir, k, getattr(data_copy, k))

  @staticmethod
  def _enum_type_converter(enum_cls: type[enum.Enum], val: Any) -> enum.Enum:
    if isinstance(val, str):
      return getattr(enum_cls, val)
    return enum_cls(val)

  @staticmethod
  def _enum_type_hook(enum_cls: type[enum.Enum]):
    return lambda val: IrDataSerializer._enum_type_converter(enum_cls, val)

  @staticmethod
  def _from_dict(data_cls: type[T], data):
    class_fields: MutableMapping[str, Any] = {}
    for name, spec in ir_data_fields.field_specs(data_cls).items():
      if (value := data.get(name)) is not None:
        if spec.is_dataclass:
          if spec.is_sequence:
            class_fields[name] = [
                IrDataSerializer._from_dict(spec.data_type, v) for v in value
            ]
          else:
            class_fields[name] = IrDataSerializer._from_dict(
                spec.data_type, value
            )
        else:
          if spec.data_type in (
              ir_data.FunctionMapping,
              ir_data.AddressableUnit,
          ):
            class_fields[name] = IrDataSerializer._enum_type_converter(
                spec.data_type, value
            )
          else:
            if spec.is_sequence:
              class_fields[name] = value
            else:
              class_fields[name] = spec.data_type(value)
    return data_cls(**class_fields)

  @staticmethod
  def from_dict(data_cls: type[T], data):
    """Creates a new IR data instance from a serialized dict"""
    return IrDataSerializer._from_dict(data_cls, data)


class _IrDataSequenceBuilder(MutableSequence[T]):
  """Wrapper for a list of IR elements

  Simply wraps the returned values during indexed access and iteration with
  IrDataBuilders.
  """

  def __init__(self, target: MutableSequence[T]):
    self._target = target

  def __delitem__(self, key):
    del self._target[key]

  def __getitem__(self, key):
    return _IrDataBuilder(self._target.__getitem__(key))

  def __setitem__(self, key, value):
    self._target[key] = value

  def __iter__(self):
    itr = iter(self._target)
    for i in itr:
      yield _IrDataBuilder(i)

  def __repr__(self):
    return repr(self._target)

  def __len__(self):
    return len(self._target)

  def __eq__(self, other):
    return self._target == other

  def __ne__(self, other):
    return self._target != other

  def insert(self, index, value):
    self._target.insert(index, value)

  def extend(self, values):
    self._target.extend(values)


class _IrDataBuilder(Generic[T]):
  """Wrapper for an IR element"""

  def __init__(self, ir: T) -> None:
    assert ir is not None
    self.ir: T = ir

  def __setattr__(self, __name: str, __value: Any) -> None:
    if __name == "ir":
      # This our proxy object
      object.__setattr__(self, __name, __value)
    else:
      # Passthrough to the proxy object
      ir: T = object.__getattribute__(self, "ir")
      setattr(ir, __name, __value)

  def __getattribute__(self, name: str) -> Any:
    """Hook for `getattr` that handles adding missing fields.

    If the field is missing inserts it, and then returns either the raw value
    for basic types
    or a new IrBuilder wrapping the field to handle the next field access in a
    longer chain.
    """

    # Check if getting one of the builder attributes
    if name in ("CopyFrom", "ir"):
      return object.__getattribute__(self, name)

    # Get our target object by bypassing our getattr hook
    ir: T = object.__getattribute__(self, "ir")
    if ir is None:
      return object.__getattribute__(self, name)

    if name in ("HasField", "WhichOneof"):
      return getattr(ir, name)

    field_spec = field_specs(ir).get(name)
    if field_spec is None:
      raise AttributeError(
          f"No field {name} on {type(ir).__module__}.{type(ir).__name__}."
      )

    obj = getattr(ir, name, None)
    if obj is None:
      # Create a default and store it
      obj = ir_data_fields.build_default(field_spec)
      setattr(ir, name, obj)

    if field_spec.is_dataclass:
      obj = (
          _IrDataSequenceBuilder(obj)
          if field_spec.is_sequence
          else _IrDataBuilder(obj)
      )

    return obj

  def CopyFrom(self, template: T):  # pylint:disable=invalid-name
    """Updates the fields of this class with values set in the template"""
    update(cast(type[T], self), template)


def builder(target: T) -> T:
  """Create a wrapper around the target to help build an IR Data structure"""
  # Check if the target is already a builder.
  if isinstance(target, (_IrDataBuilder, _IrDataSequenceBuilder)):
    return target

  # Builders are only valid for IR data classes.
  if not hasattr(type(target), "IR_DATACLASS"):
    raise TypeError(f"Builder target {type(target)} is not an ir_data.message")

  # Create a builder and cast it to the target type to expose type hinting for
  # the wrapped type.
  return cast(T, _IrDataBuilder(target))


def _field_checker_from_spec(spec: ir_data_fields.FieldSpec):
  """Helper that builds an FieldChecker that pretends to be an IR class"""
  if spec.is_sequence:
    return []
  if spec.is_dataclass:
    return _ReadOnlyFieldChecker(spec)
  return ir_data_fields.build_default(spec)


def _field_type(ir_or_spec: T | ir_data_fields.FieldSpec) -> type:
  if isinstance(ir_or_spec, ir_data_fields.FieldSpec):
    return ir_or_spec.data_type
  return type(ir_or_spec)


class _ReadOnlyFieldChecker:
  """Class used the chain calls to fields that aren't set"""

  def __init__(self, ir_or_spec: T | ir_data_fields.FieldSpec) -> None:
    self.ir_or_spec = ir_or_spec

  def __setattr__(self, name: str, value: Any) -> None:
    if name == "ir_or_spec":
      return object.__setattr__(self, name, value)

    raise AttributeError(f"Cannot set {name} on read-only wrapper")

  def __getattribute__(self, name: str) -> Any:  # pylint:disable=too-many-return-statements
    ir_or_spec = object.__getattribute__(self, "ir_or_spec")
    if name == "ir_or_spec":
      return ir_or_spec

    field_type = _field_type(ir_or_spec)
    spec = field_specs(field_type).get(name)
    if not spec:
      if isinstance(ir_or_spec, ir_data_fields.FieldSpec):
        if name == "HasField":
          return lambda x: False
        if name == "WhichOneof":
          return lambda x: None
      return object.__getattribute__(ir_or_spec, name)

    if isinstance(ir_or_spec, ir_data_fields.FieldSpec):
      # Just pretending
      return _field_checker_from_spec(spec)

    value = getattr(ir_or_spec, name)
    if value is None:
      return _field_checker_from_spec(spec)

    if spec.is_dataclass:
      if spec.is_sequence:
        return [_ReadOnlyFieldChecker(i) for i in value]
      return _ReadOnlyFieldChecker(value)

    return value

  def __eq__(self, other):
    if isinstance(other, _ReadOnlyFieldChecker):
      other = other.ir_or_spec
    return self.ir_or_spec == other

  def __ne__(self, other):
    return not self == other


def reader(obj: T | _ReadOnlyFieldChecker) -> T:
  """Builds a read-only wrapper that can be used to check chains of possibly

  unset fields.

  This wrapper explicitly does not alter the wrapped object and is only
  intended for reading contents.

  For example, a `reader` lets you do:
  ```
  def get_function_name_end_column(function: ir_data.Function):
    return reader(function).function_name.source_location.end.column
  ```

  Instead of:
  ```
  def get_function_name_end_column(function: ir_data.Function):
    if function.function_name:
      if function.function_name.source_location:
        if function.function_name.source_location.end:
          return function.function_name.source_location.end.column
    return 0
  ```
  """
  # Create a read-only wrapper if it's not already one.
  if not isinstance(obj, _ReadOnlyFieldChecker):
    obj = _ReadOnlyFieldChecker(obj)

  # Cast it back to the original type.
  return cast(T, obj)


def _extract_ir(
    ir_or_wrapper: T | _ReadOnlyFieldChecker | _IrDataBuilder | None,
) -> ir_data_fields.IrDataclassInstance | None:
  if isinstance(ir_or_wrapper, _ReadOnlyFieldChecker):
    ir_or_spec = ir_or_wrapper.ir_or_spec
    if isinstance(ir_or_spec, ir_data_fields.FieldSpec):
      # This is a placeholder entry, no fields are set.
      return None
    ir_or_wrapper = ir_or_spec
  elif isinstance(ir_or_wrapper, _IrDataBuilder):
    ir_or_wrapper = ir_or_wrapper.ir
  return cast(ir_data_fields.IrDataclassInstance, ir_or_wrapper)


def fields_and_values(
    ir_wrapper: T | _ReadOnlyFieldChecker,
    value_filt: Optional[Callable[[Any], bool]] = None,
) -> list[Tuple[ir_data_fields.FieldSpec, Any]]:
  """Retrieves the fields and their values for a given IR data class.

  Args:
    ir: The IR data class or a read-only wrapper of an IR data class.
    value_filt: Optional filter used to exclude values.
  """
  if (ir := _extract_ir(ir_wrapper)) is None:
    return []

  return ir_data_fields.fields_and_values(ir, value_filt)


def get_set_fields(ir: T):
  """Retrieves the field spec and value of fields that are set in the given IR data class.

  A value is considered "set" if it is not None.
  """
  return fields_and_values(ir, lambda v: v is not None)


def copy(ir_wrapper: T | None) -> T | None:
  """Creates a copy of the given IR data class"""
  if (ir := _extract_ir(ir_wrapper)) is None:
    return None
  ir_copy = ir_data_fields.copy(ir)
  return cast(T, ir_copy)


def update(ir: T, template: T):
  """Updates `ir`s fields with all set fields in the template."""
  if not (template_ir := _extract_ir(template)):
    return

  ir_data_fields.update(
      cast(ir_data_fields.IrDataclassInstance, ir), template_ir
  )