[fir] add fir.array_modify op to allow custom assignment semantics

fir.array_update is only handling intrinsic assignments.
They are two big differences with user defined assignments:
1. The LHS and RHS types may not match, this does not play well
   with fir.array_update that relies on both the merge and the
   updated element to have the same type.
2. user defined assignment has a call semantics, with potential
   side effects. So if a fir.array_update can hide a call, it traits
   would need to be updated.

Instead of hiding more semantic in the fir.array_update, introduce
a new fir.array_modify op that allows de-correlating indicating that
an array value element is modified, and how it is modified.
This allow the ArrayValueCopy pass to still perform copy elision
while not having to implement the call itself, and could in general
be used for all kind of assignments (e.g. character assignment).

Update the alias analysis to not rely on the merge arguments (since
fir.array_modify has none).
Instead, analyze what is done with the element address.
This implies adding the ability to follow the users of fir.array_modify,
as well as being able to go through fir.store that may be generated to
store the RHS value in order to pass it to a user define routine.
This is done by adding a ReachCollector class to gather all array
accesses.

Author: jeanPerier

flang/include/flang/Optimizer/Dialect/FIROps.td

@@ -1479,6 +1479,64 @@ def fir_ArrayUpdateOp : fir_Op<"array_update", [AttrSizedOperandSegments,
   let verifier = "return ::verify(*this);";
 }
 
+def fir_ArrayModifyOp : fir_Op<"array_modify", [AttrSizedOperandSegments,
+    NoSideEffect]> {
+  let summary = "Get an address for an array value to modify it.";
+
+  let description = [{
+    Modify the value of an element in an array value through actions done
+    on the returned address. A new array value is also
+    returned where all element values of the input array are identical except
+    for the selected element which is the value after the modification done
+    on the element address.
+
+    ```fortran
+      real :: a(n)
+      ...
+      ! Elemental user defined assignment from type(SomeType) to real.
+      a = value_of_some_type
+    ```
+
+    One can use `fir.array_modify` to update the (implied) value of `a(i)`
+    in an array expression as shown above.
+
+    ```mlir
+      %s = fir.shape %n : (index) -> !fir.shape<1>
+      // Load the entire array 'a'.
+      %v = fir.array_load %a(%s) : (!fir.ref<!fir.array<?xf32>>, !fir.shape<1>) -> !fir.array<?xf32>
+      // Update the value of one of the array value's elements with a user
+      // defined assignment from %rhs.
+      %new = fir.do_loop %i = ... (%inner = %v) {
+        %rhs = ...
+        %addr, %r = fir.array_modify %inner, %i, %j : (!fir.array<?xf32>,  index) -> fir.ref<f32>, !fir.array<?xf32>
+        fir.call @user_def_assign(%addr, %rhs) (fir.ref<f32>, fir.ref<!fir.type<SomeType>>) -> ()
+        fir.result %r : !fir.ref<!fir.array<?xf32>>
+      }
+      fir.array_merge_store %v, %new to %a : !fir.ref<!fir.array<?xf32>>
+    ```
+
+    An array value modification behaves as if a mapping function from the indices
+    to the new value has been added, replacing the previous mapping. These
+    mappings can be added to the ssa-value, but will not be materialized in
+    memory until the `fir.array_merge_store` is performed.
+  }];
+
+  let arguments = (ins
+    fir_SequenceType:$sequence,
+    Variadic<AnyCoordinateType>:$indices,
+    Variadic<AnyIntegerType>:$typeparams
+  );
+
+  let results = (outs fir_ReferenceType, fir_SequenceType);
+
+  let assemblyFormat = [{
+    $sequence `,` $indices (`typeparams` $typeparams^)? attr-dict
+      `:` functional-type(operands, results)
+  }];
+
+  let verifier = [{ return ::verify(*this); }];
+}
+
 def fir_ArrayMergeStoreOp : fir_Op<"array_merge_store",
     [AttrSizedOperandSegments]> {
   let summary = "Store merged array value to memory.";

flang/lib/Optimizer/Dialect/FIROps.cpp

@@ -499,6 +499,18 @@ static mlir::LogicalResult verify(fir::ArrayUpdateOp op) {
   return mlir::success();
 }
 
+//===----------------------------------------------------------------------===//
+// ArrayModifyOp
+//===----------------------------------------------------------------------===//
+
+static mlir::LogicalResult verify(fir::ArrayModifyOp op) {
+  auto arrTy = op.sequence().getType().cast<fir::SequenceType>();
+  auto indSize = op.indices().size();
+  if (indSize < arrTy.getDimension())
+    return op.emitOpError("number of indices must match array dimension");
+  return mlir::success();
+}
+
 //===----------------------------------------------------------------------===//
 // BoxAddrOp
 //===----------------------------------------------------------------------===//