Hlfir

class CharExtremumPredicate(Enum):
    """min / max selector for hlfir.char_extremum."""

    MIN = "min"
    MAX = "max"

    @staticmethod
    def try_parse(parser: AttrParser) -> CharExtremumPredicate | None:
        for option in CharExtremumPredicate:
            if parser.parse_optional_characters(option.value) is not None:
                return option
        return None

@dataclass(frozen=True)
class CharExtremumPredicateAttrBase(Data[CharExtremumPredicate]):
    @classmethod
    def parse_parameter(cls, parser: AttrParser) -> CharExtremumPredicate:
        with parser.in_angle_brackets():
            return parser.expect(
                lambda: CharExtremumPredicate.try_parse(parser),
                "char_extremum predicate (min or max) expected",
            )

    def print_parameter(self, printer: Printer):
        with printer.in_angle_brackets():
            printer.print_string(self.data.value)

@irdl_attr_definition
class CharExtremumPredicateAttr(CharExtremumPredicateAttrBase):
    name = "hlfir.char_extremum_predicate"

@irdl_attr_definition
class ExprType(ParametrizedAttribute, TypeAttribute):
    """
    The type of an array, character, or derived type Fortran expression.

    Abstract value type for Fortran arrays, characters and derived types.
    The rank cannot be assumed, and empty shape means that the expression is a scalar.
    When the element type is a derived type, the polymorphic flag may be set to true
    to indicate that the expression dynamic type can differ from its static type.
    """

    name = "hlfir.expr"
    shape: ArrayAttr[IntegerAttr | DeferredAttr | NoneType]
    elementType: Attribute

    def print_parameters(self, printer: Printer) -> None:
        with printer.in_angle_brackets():
            for s in self.shape.data:
                if isinstance(s, DeferredAttr):
                    printer.print_string("?")
                elif isinstance(s, NoneType):
                    raise Exception(
                        "Can not have none type as part of sequence shape with only one type"
                    )
                else:
                    printer.print_int(s.value.data)
                printer.print_string("x")
            printer.print_attribute(self.elementType)

    @classmethod
    def parse_parameters(cls, parser: AttrParser) -> list[Attribute]:
        # We need extra work here as the builtin tuple is not being supported
        # yet, therefore handle this here
        def parse_interval() -> IntegerAttr | DeferredAttr:
            if parser.parse_optional_punctuation("?"):
                return DeferredAttr()
            s = parser.parse_integer(allow_boolean=False)
            return IntegerAttr(s, 32)

        shape: list[IntegerAttr | DeferredAttr] = []
        with parser.in_angle_brackets():
            elementType = parser.parse_optional_type()
            while elementType is None:
                shape.append(parse_interval())
                parser.parse_shape_delimiter()
                elementType = parser.parse_optional_type()
        return [ArrayAttr(shape), elementType]

@irdl_op_definition
class DeclareOp(IRDLOperation):
    """
    Declare a variable and produce an SSA value that can be used as a variable in HLFIR operations.

    Tie the properties of a Fortran variable to an address. The properties
    include bounds, length parameters, and Fortran attributes.

    The arguments are the same as for fir.declare.

    The main difference with fir.declare is that hlfir.declare returns two
    values:
      - the first one is an SSA value that allows retrieving the variable
        address, bounds, and type parameters at any point without requiring
        access to the defining operation. This may be:
        - for scalar numerical, logical, or derived type without length
          parameters: a fir.ref<T> (e.g. fir.ref<i32>)
        - for scalar characters: a fir.boxchar<kind> or fir.ref<fir.char<kind,
          cst_len>>
        - for arrays of types without length parameters, without lower bounds,
          that are not polymorphic and with a constant shape:
          fir.ref<fir.array<cst_shapexT>>
        - for all non pointer/non allocatable entities: fir.box<T>, and
          fir.class<T> for polymorphic entities.
        - for all pointers/allocatables:
          fir.ref<fir.box<fir.ptr<T>>>/fir.ref<fir.box<fir.heap<T>>>
      - the second value has the same type as the input memref, and is the
        same. If it is a fir.box or fir.class, it may not contain accurate
        local lower bound values. It is intended to be used when generating FIR
        from HLFIR in order to avoid descriptor creation for simple entities.

    Example:
    CHARACTER(n) :: c(10:n, 20:n)

    Can be represented as:
    func.func @foo(%arg0: !fir.ref<!fir.array<?x?x!fir.char<1,?>>>, %arg1: !fir.ref<i64>) {
      %c10 = arith.constant 10 : index
      %c20 = arith.constant 20 : index
      %1 = fir.load %ag1 : fir.ref<i64>
      %2 = fir.shape_shift %c10, %1, %c20, %1 : (index, index, index, index) -> !fir.shapeshift<2>
      %3 = hfir.declare %arg0(%2) typeparams %1 {uniq_name = "c"} (fir.ref<!fir.array<?x?x!fir.char<1,?>>>, fir.shapeshift<2>, index) -> (fir.box<!fir.array<?x?x!fir.char<1,?>>>, fir.ref<!fir.array<?x?x!fir.char<1,?>>>)
      // ... uses %3#0 as "c"
    }
    """  # noqa: E501

    name = "hlfir.declare"
    memref = operand_def()
    shape = opt_operand_def()
    typeparams = var_operand_def()
    dummy_scope = opt_operand_def()
    storage = opt_operand_def()
    storage_offset = opt_prop_def(IntegerAttr)
    uniq_name = prop_def(StringAttr)
    fortran_attrs = opt_prop_def(FortranVariableFlagsAttr)
    data_attr = opt_prop_def(Attribute)
    skip_rebox = opt_prop_def(UnitAttr)
    dummy_arg_no = opt_prop_def(IntegerAttr)
    result = result_def()
    result2 = result_def()

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

@irdl_op_definition
class DesignateOp(IRDLOperation):
    """
    Designate a Fortran variable.

    This operation represents a Fortran "part-ref", except that it can embed a
    substring or or complex part directly, and that vector subscripts cannot be
    used. It returns a Fortran variable that is a part of the input variable.

    The operands are as follow:
      - memref is the variable being designated.
      - component may be provided if the memref is a derived type to
        represent a reference to a component. It must be the name of a
        component of memref derived type.
      - component_shape represents the shape of the component and must be
        provided if and only if both component and indices appear.
      - indices can be provided to index arrays. The indices may be simple
        indices or triplets.
        If indices are provided and there is a component, the component must be
        an array component and the indices index the array component.
        If memref is an array, and component is provided and is an array
        component, indices must be provided and must not be triplets. This
        ensures hlfir.designate does not create arrays of arrays (which is not
        possible in Fortran).
      - substring may contain two values to represent a substring lower and
        upper bounds.
      - complex_part may be provided to represent a complex part (true
        represents the imaginary part, and false the real part).
      - shape represents the shape of the result and must be provided if the
        result is an array that is not a box address.
      - typeparams represents the length parameters of the result and must be
        provided if the result type has length parameters and is not a box
        address.
    """

    name = "hlfir.designate"
    memref = operand_def()
    component_shape = opt_operand_def()
    indices = var_operand_def()
    substring = var_operand_def()
    shape = opt_operand_def()
    typeparams = var_operand_def()
    component = opt_prop_def(StringAttr)
    complex_part = opt_prop_def(BoolAttr)
    is_triplet = prop_def(DenseArrayBase)
    fortran_attrs = opt_prop_def(FortranVariableFlagsAttr)
    result = result_def()

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

@irdl_op_definition
class AssignOp(IRDLOperation):
    """
    Assign an expression or variable value to a Fortran variable.

    Assign rhs to lhs following Fortran intrinsic assignments rules.
    The operation deals with inserting a temporary if the lhs and rhs
    may overlap.

    The optional "realloc" flag allows indicating that this assignment
    has the Fortran 95 semantics for assignments to a whole allocatable.
    In such case, the left hand side must be an allocatable that may be
    unallocated or allocated with a different type and shape than the right
    hand side. It will be allocated or re-allocated as needed during the
    assignment.

    When "realloc" is set and this is a character assignment, the optional
    flag "keep_lhs_length_if_realloc" indicates that the character
    left hand side should retain its length after the assignment. If the
    right hand side has a different length, truncation and padding will
    occur. This covers the case of explicit and assumed length character
    allocatables.

    Otherwise, the left hand side will be allocated or reallocated to match the
    right hand side length if they differ. This covers the case of deferred
    length character allocatables.

    The optional "temporary_lhs" flag indicates that the LHS is a compiler
    generated temporary. In this case the temporary is initialized if needed
    (e.g. the LHS is of derived type with allocatable/pointer components),
    and the assignment is done without LHS (or its subobjects) finalization
    and with automatic allocation.

    If "temporary_lhs" and "keep_lhs_length_if_realloc" are both set,
    this assign operation denotes special case of character allocatable
    LHS with explicit length. The LHS that must preserve its length
    during the assignment regardless of the the RHS's length or/and
    allocation status. This assign operation will be lowered into a call
    to AssignExplicitLengthCharacter().
    """

    name = "hlfir.assign"
    rhs = operand_def()
    lhs = operand_def()
    realloc = opt_prop_def(UnitAttr)
    keep_lhs_length_if_realloc = opt_prop_def(UnitAttr)
    temporary_lhs = opt_prop_def(UnitAttr)

@irdl_op_definition
class ParentComponentOp(IRDLOperation):
    """
    Designate the parent component of a variable.

    This operation represents a Fortran component reference where the
    component name is a parent type of the variable's derived type.
    These component references cannot be represented with an hlfir.designate
    because the parent type names are not embedded in fir.type<> types
    as opposed to the actual component names.

    The operands are as follow:
      - memref is a derived type variable whose parent component is being
        designated.
      - shape is the shape of memref and the result and must be provided if
        memref is an array. Parent component reference lower bounds are ones,
        so the provided shape must be a fir.shape.
      - typeparams are the type parameters of the parent component type if any.
        It is a subset of memref type parameters.
    The parent component type and name is reflected in the result type.
    """

    name = "hlfir.parent_comp"
    memref = operand_def()
    shape = opt_operand_def()
    typeparams = var_operand_def()
    result = result_def()

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

@irdl_op_definition
class ConcatOp(IRDLOperation):
    """
    Concatenate characters.

    Concatenate two or more character strings of a same character kind.
    """

    name = "hlfir.concat"
    strings = var_operand_def()
    length = operand_def()
    result = result_def()

@irdl_op_definition
class AllOp(IRDLOperation):
    """
    ALL transformational intrinsic.

    Takes a logical array MASK as argument, optionally along a particular dimension,
    and returns true if all elements of MASK are true.
    """

    name = "hlfir.all"
    mask = operand_def()
    dim = opt_operand_def()
    result = result_def()

@irdl_op_definition
class AnyOp(IRDLOperation):
    """
    ANY transformational intrinsic.

    Takes a logical array MASK as argument, optionally along a particular dimension,
    and returns true if any element of MASK is true.
    """

    name = "hlfir.any"
    mask = operand_def()
    dim = opt_operand_def()
    result = result_def()

@irdl_op_definition
class CountOp(IRDLOperation):
    """
    COUNT transformational intrinsic.

    Takes a logical and counts the number of true values.
    """

    name = "hlfir.count"
    mask = operand_def()
    dim = opt_operand_def()
    result = result_def()

@irdl_op_definition
class MaxvalOp(IRDLOperation):
    """
    MAXVAL transformational intrinsic.

    Maximum value(s) of an array.
    If DIM is absent, the result is a scalar.
    If DIM is present, the result is an array of rank n-1, where n is the rank of ARRAY.
    """

    name = "hlfir.maxval"
    array = operand_def()
    dim = opt_operand_def()
    mask = opt_operand_def()
    fastmath = opt_prop_def(FastMathFlagsAttr)
    result = result_def()

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

@irdl_op_definition
class MinvalOp(IRDLOperation):
    """
    MINVAL transformational intrinsic.

    Minimum value(s) of an array.
    If DIM is absent, the result is a scalar.
    If DIM is present, the result is an array of rank n-1, where n is the rank of ARRAY.
    """

    name = "hlfir.minval"
    array = operand_def()
    dim = opt_operand_def()
    mask = opt_operand_def()
    fastmath = opt_prop_def(FastMathFlagsAttr)
    result = result_def()

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

@irdl_op_definition
class ProductOp(IRDLOperation):
    """
    PRODUCT transformational intrinsic.

    Multiplies the elements of an array, optionally along a particular dimension,
    optionally if a mask is true.
    """

    name = "hlfir.product"
    array = operand_def()
    dim = opt_operand_def()
    mask = opt_operand_def()
    fastmath = opt_prop_def(FastMathFlagsAttr)
    result = result_def()

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

@irdl_op_definition
class SetLengthOp(IRDLOperation):
    """
    Change the length of a character entity.

    This trims or pads the
    character argument according to the new length.
    """

    name = "hlfir.set_length"
    string = operand_def()
    length = operand_def()
    result = result_def()

@irdl_op_definition
class GetLengthOp(IRDLOperation):
    """
    Get the length of a character entity.

    Get the length of character entity represented as hlfir.expr.
    """

    name = "hlfir.get_length"
    expr = operand_def()
    result = result_def()

    traits = traits_def(Pure())

@irdl_op_definition
class SumOp(IRDLOperation):
    """
    SUM transformational intrinsic.

    Sums the elements of an array, optionally along a particular dimension,
    optionally if a mask is true.
    """

    name = "hlfir.sum"
    array = operand_def()
    dim = opt_operand_def()
    mask = opt_operand_def()
    fastmath = opt_prop_def(FastMathFlagsAttr)
    result = result_def()

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

@irdl_op_definition
class DotProductOp(IRDLOperation):
    """
    DOT_PRODUCT transformational intrinsic,

    Dot product of two vectors,
    """

    name = "hlfir.dot_product"
    lhs = operand_def()
    rhs = operand_def()
    fastmath = opt_prop_def(FastMathFlagsAttr)
    result = result_def()

@irdl_op_definition
class MatmulOp(IRDLOperation):
    """
    MATMUL transformational intrinsic.

    Matrix multiplication.
    """

    name = "hlfir.matmul"
    lhs = operand_def()
    rhs = operand_def()
    fastmath = opt_prop_def(FastMathFlagsAttr)
    result = result_def()

@irdl_op_definition
class TransposeOp(IRDLOperation):
    """
    TRANSPOSE transformational intrinsic.

    Transpose a rank 2 array.
    """

    name = "hlfir.transpose"
    array = operand_def()
    result = result_def()

@irdl_op_definition
class MatmulTransposeOp(IRDLOperation):
    """
    Optimized MATMUL(TRANSPOSE(...), ...).

    Matrix multiplication where the left hand side is transposed.
    """

    name = "hlfir.matmul_transpose"
    lhs = operand_def()
    rhs = operand_def()
    fastmath = opt_prop_def(FastMathFlagsAttr)
    result = result_def()

@irdl_op_definition
class AssociateOp(IRDLOperation):
    """
    Create a variable from an expression value.

    For expressions, this operation is an incentive to re-use the expression
    storage, if any, after the bufferization pass when possible (if the
    expression is not used afterwards).
    """

    name = "hlfir.associate"
    source = operand_def()
    shape = opt_operand_def()
    typeparams = var_operand_def()
    uniq_name = opt_prop_def(StringAttr)
    fortran_attrs = opt_prop_def(FortranVariableFlagsAttr)
    result = result_def()
    fir_base = result_def()
    must_free = result_def()

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

@irdl_op_definition
class EndAssociateOp(IRDLOperation):
    """
    Mark the end of life of a variable associated to an expression.

    If the expression has a derived type that may contain allocatable
    components, the variable operand must be a Fortran entity.
    """

    name = "hlfir.end_associate"
    var = operand_def()
    must_free = operand_def()

@irdl_op_definition
class AsExprOp(IRDLOperation):
    """
    Take the value of an array, character or derived variable.

    In general, this operation will lead to a copy of the variable
    in the bufferization pass if it was not transformed.

    However, if it is known that the variable storage will not be used anymore
    afterwards, the variable storage ownership can be passed to the hlfir.expr
    by providing the $must_free argument that is a boolean that indicates if
    the storage must be freed (when it was allocated on the heap).
    This allows Fortran lowering to build some expression value in memory when
    there is no adequate hlfir operation, and to promote the result to an
    hlfir.expr value without paying the price of introducing a copy.
    """

    name = "hlfir.as_expr"
    var = operand_def()
    must_free = opt_operand_def()
    result = result_def()

@irdl_op_definition
class NoReassocOp(IRDLOperation):
    """
    Synthetic op to prevent reassociation.

    Same as fir.reassoc, except it accepts hlfir.expr arguments.
    """

    name = "hlfir.no_reassoc"
    val = operand_def()
    result = result_def()

@irdl_op_definition
class ElementalOp(IRDLOperation):
    """
    Elemental expression.

    Represent an elemental expression as a function of the indices.
    This operation contain a region whose block arguments are one
    based indices iterating over the elemental expression shape.
    Given these indices, the element value for the given iteration
    can be computed in the region and yielded with the hlfir.yield_element
    operation.

    The shape and typeparams operands represent the extents and type
    parameters of the resulting array value.

    The optional mold is an entity carrying the information about
    the dynamic type of the polymorphic result. Note that the shape
    of the mold does not necessarily match the shape of the result,
    for example, the result of `merge(poly_scalar1, poly_scalar2, mask_array)`
    will have the shape of `mask_array` and the dynamic type of `poly_scalar*`.

    The unordered attribute can be set to allow out of order processing
    of the indices. This is safe only if the operations in the body
    of the elemental do not have side effects.


    Example: Y + X,  with Integer :: X(10, 20), Y(10,20)

    %0 = fir.shape %c10, %c20 : (index, index) -> !fir.shape<2>
    %5 = hlfir.elemental %0 : (!fir.shape<2>) -> !hlfir.expr<10x20xi32> {
    ^bb0(%i: index, %j: index):
      %6 = hlfir.designate %x (%i, %j)  : (!fir.ref<!fir.array<10x20xi32>>, index, index) -> !fir.ref<i32>
      %7 = hlfir.designate %y (%i, %j)  : (!fir.ref<!fir.array<10x20xi32>>, index, index) -> !fir.ref<i32>
      %8 = fir.load %6 : !fir.ref<i32>
      %9 = fir.load %7 : !fir.ref<i32>
      %10 = arith.addi %8, %9 : i32
      hlfir.yield_element %10 : i32
    }
    """

    name = "hlfir.elemental"
    shape = operand_def()
    mold = opt_operand_def()
    typeparams = var_operand_def()
    unordered = opt_prop_def(UnitAttr)
    result = result_def()
    region = region_def("single_block")

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

@irdl_op_definition
class YieldElementOp(IRDLOperation):
    """
    Yield the elemental value in an ElementalOp.

    Yield the element value of the current elemental expression iteration
    in an hlfir.elemental region. See hlfir.elemental description for an
    example.
    """

    name = "hlfir.yield_element"
    element_value = operand_def()

    traits = traits_def(IsTerminator(), HasParent(ElementalOp), Pure())

@irdl_op_definition
class ApplyOp(IRDLOperation):
    """
    Get the element value of an expression.

    Given an hlfir.expr array value, hlfir.apply allow retrieving
    the value for an element given one based indices.

    When hlfir.apply is used on an hlfir.elemental, and if the hlfir.elemental
    operation evaluation can be moved to the location of the hlfir.apply, it is
    as if the hlfir.elemental body was evaluated given the hlfir.apply indices.
    Therefore, apply operations on hlfir.elemental expressions should be located
    such that evaluating the hlfir.elemental at the position of the hlfir.apply
    operation produces the same result as evaluating the hlfir.elemental at its
    location in the instruction stream. Attention should be paid to
    hlfir.elemental memory side effects (in practice these are unlikely).
    "10.1.4 Evaluation of operations" says that expression evaluation shall not
    impact/be impacted by other expression evaluation in the statement.
    """

    name = "hlfir.apply"
    expr = operand_def()
    indices = var_operand_def()
    typeparams = var_operand_def()
    result = result_def()

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

@irdl_op_definition
class NullOp(IRDLOperation):
    """
    Create a NULL() address.

    So far is not intended to represent NULL(MOLD).
    """

    name = "hlfir.null"
    result = result_def()