[flang] Update ArrayValueCopy to support array_amend and array_access

This patch update the array value copy pass to support fir-array_amend
and fir.array_access.

This patch is part of the upstreaming effort from fir-dev branch.

Reviewed By: PeteSteinfeld, schweitz

Differential Revision: https://reviews.llvm.org/D121300

Co-authored-by: Jean Perier <jperier@nvidia.com>
Co-authored-by: Eric Schweitz <eschweitz@nvidia.com>

flang/include/flang/Optimizer/Builder/Array.h

@@ -0,0 +1,27 @@
+//===-- Array.h -------------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_OPTIMIZER_BUILDER_ARRAY_H
+#define FORTRAN_OPTIMIZER_BUILDER_ARRAY_H
+
+#include "flang/Optimizer/Dialect/FIROps.h"
+
+namespace fir::factory {
+
+/// Return true if and only if the extents are those of an assumed-size array.
+/// An assumed-size array in Fortran is declared with `*` as the upper bound of
+/// the last dimension of the array. Lowering converts the asterisk to an
+/// undefined value.
+inline bool isAssumedSize(const llvm::SmallVectorImpl<mlir::Value> &extents) {
+  return !extents.empty() &&
+         mlir::isa_and_nonnull<UndefOp>(extents.back().getDefiningOp());
+}
+
+} // namespace fir::factory
+
+#endif // FORTRAN_OPTIMIZER_BUILDER_ARRAY_H

flang/include/flang/Optimizer/Builder/FIRBuilder.h

@@ -371,6 +371,12 @@ class FirOpBuilder : public mlir::OpBuilder {
   /// Generate code testing \p addr is a null address.
   mlir::Value genIsNull(mlir::Location loc, mlir::Value addr);
 
+  /// Compute the extent of (lb:ub:step) as max((ub-lb+step)/step, 0). See
+  /// Fortran 2018 9.5.3.3.2 section for more details.
+  mlir::Value genExtentFromTriplet(mlir::Location loc, mlir::Value lb,
+                                   mlir::Value ub, mlir::Value step,
+                                   mlir::Type type);
+
 private:
   const KindMapping &kindMap;
 };

flang/include/flang/Optimizer/Dialect/FIROpsSupport.h

@@ -83,6 +83,10 @@ static constexpr llvm::StringRef getHostAssocAttrName() {
   return "fir.host_assoc";
 }
 
+/// Does the function, \p func, have a host-associations tuple argument?
+/// Some internal procedures may have access to host procedure variables.
+bool hasHostAssociationArgument(mlir::FuncOp func);
+
 /// Tell if \p value is:
 ///   - a function argument that has attribute \p attributeName
 ///   - or, the result of fir.alloca/fir.allocamem op that has attribute \p

flang/lib/Optimizer/Builder/FIRBuilder.cpp

@@ -507,6 +507,23 @@ mlir::Value fir::FirOpBuilder::genIsNull(mlir::Location loc, mlir::Value addr) {
                                   mlir::arith::CmpIPredicate::eq);
 }
 
+mlir::Value fir::FirOpBuilder::genExtentFromTriplet(mlir::Location loc,
+                                                    mlir::Value lb,
+                                                    mlir::Value ub,
+                                                    mlir::Value step,
+                                                    mlir::Type type) {
+  auto zero = createIntegerConstant(loc, type, 0);
+  lb = createConvert(loc, type, lb);
+  ub = createConvert(loc, type, ub);
+  step = createConvert(loc, type, step);
+  auto diff = create<mlir::arith::SubIOp>(loc, ub, lb);
+  auto add = create<mlir::arith::AddIOp>(loc, diff, step);
+  auto div = create<mlir::arith::DivSIOp>(loc, add, step);
+  auto cmp = create<mlir::arith::CmpIOp>(loc, mlir::arith::CmpIPredicate::sgt,
+                                         div, zero);
+  return create<mlir::arith::SelectOp>(loc, cmp, div, zero);
+}
+
 //===--------------------------------------------------------------------===//
 // ExtendedValue inquiry helper implementation
 //===--------------------------------------------------------------------===//

flang/lib/Optimizer/Dialect/FIROps.cpp

@@ -3258,6 +3258,15 @@ fir::GlobalOp fir::createGlobalOp(mlir::Location loc, mlir::ModuleOp module,
   return result;
 }
 
+bool fir::hasHostAssociationArgument(mlir::FuncOp func) {
+  if (auto allArgAttrs = func.getAllArgAttrs())
+    for (auto attr : allArgAttrs)
+      if (auto dict = attr.template dyn_cast_or_null<mlir::DictionaryAttr>())
+        if (dict.get(fir::getHostAssocAttrName()))
+          return true;
+  return false;
+}
+
 bool fir::valueHasFirAttribute(mlir::Value value,
                                llvm::StringRef attributeName) {
   // If this is a fir.box that was loaded, the fir attributes will be on the