Split allocatable assignment to avoid deallocation before RHS evaluation

jeanPerier · jeanPerier · commit c329c4f8de1a · 2021-11-04T15:02:09.000+01:00
The previous allocatable assignment lowering was both allocating the new
storage and deallocating the old one/updating the allocatable before
generating a normal assignment in the new storage.

This is wrong in general because the previous allocatable value may be
used in the assignment RHS, so the deallocation/allocatable update must
be done after the assignment.

This patch split genReallocIfNeeded in two parts: the first part that
create a new storage if needed, to be called before lowering the
assignment, and a second part that creates the deallocation/allocatable
descriptor update after the assignment if a new storage was allocated.

This fixes runtime error for allocatable assignments like `a = a(1:100)`.
diff --git a/flang/include/flang/Optimizer/Builder/BoxValue.h b/flang/include/flang/Optimizer/Builder/BoxValue.h
@@ -467,6 +467,14 @@ inline bool isUnboxedValue(const ExtendedValue &exv) {
       [](const fir::UnboxedValue &box) { return box ? true : false; },
       [](const auto &) { return false; });
 }
+
+/// Is the extended value `exv` a derived type with length parameters ?
+inline bool isDerivedWithLengthParameters(const ExtendedValue &exv) {
+  auto type = fir::unwrapPassByRefType(fir::getBase(exv).getType());
+  auto record = fir::unwrapSequenceType(type).dyn_cast<fir::RecordType>();
+  return record && record.getNumLenParams() != 0;
+}
+
 } // namespace fir
 
 #endif // FORTRAN_OPTIMIZER_BUILDER_BOXVALUE_H
diff --git a/flang/include/flang/Optimizer/Builder/MutableBox.h b/flang/include/flang/Optimizer/Builder/MutableBox.h
@@ -13,6 +13,7 @@
 #ifndef FORTRAN_OPTIMIZER_BUILDER_MUTABLEBOX_H
 #define FORTRAN_OPTIMIZER_BUILDER_MUTABLEBOX_H
 
+#include "flang/Optimizer/Builder/BoxValue.h"
 #include "llvm/ADT/StringRef.h"
 
 namespace mlir {
@@ -24,8 +25,6 @@ class Location;
 
 namespace fir {
 class FirOpBuilder;
-class MutableBoxValue;
-class ExtendedValue;
 } // namespace fir
 
 namespace fir::factory {
@@ -86,10 +85,23 @@ void disassociateMutableBox(fir::FirOpBuilder &builder, mlir::Location loc,
 /// parameter mismatch can trigger a reallocation. See Fortran 10.2.1.3 point 3
 /// that this function is implementing for more details. The polymorphic
 /// requirements are not yet covered by this function.
-void genReallocIfNeeded(fir::FirOpBuilder &builder, mlir::Location loc,
-                        const fir::MutableBoxValue &box,
-                        mlir::ValueRange lbounds, mlir::ValueRange shape,
-                        mlir::ValueRange lengthParams);
+struct MutableBoxReallocation {
+  fir::ExtendedValue newValue;
+  mlir::Value oldAddress;
+  mlir::Value wasReallocated;
+  mlir::Value oldAddressWasAllocated;
+};
+
+MutableBoxReallocation genReallocIfNeeded(fir::FirOpBuilder &builder,
+                                          mlir::Location loc,
+                                          const fir::MutableBoxValue &box,
+                                          mlir::ValueRange shape,
+                                          mlir::ValueRange lengthParams);
+
+void finalizeRealloc(fir::FirOpBuilder &builder, mlir::Location loc,
+                     const fir::MutableBoxValue &box, mlir::ValueRange lbounds,
+                     bool takeLboundsIfRealloc,
+                     const MutableBoxReallocation &realloc);
 
 /// Finalize a mutable box if it is allocated or associated. This includes both
 /// calling the finalizer, if any, and deallocating the storage.
diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
@@ -1942,26 +1942,25 @@ class FirConverter : public Fortran::lower::AbstractConverter {
                   isNumericScalarCategory(lhsType->category());
               auto rhs = isNumericScalar ? genExprValue(assign.rhs, stmtCtx)
                                          : genExprAddr(assign.rhs, stmtCtx);
-              auto lowerAllocatableLHS = [&]() -> fir::ExtendedValue {
-                auto lhs = genExprMutableBox(loc, assign.lhs);
-                llvm::SmallVector<mlir::Value> lengthParams;
-                if (auto *charBox = rhs.getCharBox())
-                  lengthParams.push_back(charBox->getLen());
-                else if (lhs.isDerivedWithLengthParameters())
-                  TODO(loc, "assignment to derived type allocatable with "
-                            "length parameters");
-                fir::factory::genReallocIfNeeded(
-                    *builder, loc, lhs, /*lbounds=*/llvm::None,
-                    /*shape=*/llvm::None, lengthParams);
-                // Assume lhs is not polymorphic for now given TODO above,
-                // otherwise, the read would is conservative and returns
-                // BoxValue for derived types.
-                return fir::factory::genMutableBoxRead(
-                    *builder, loc, lhs, /*mayBePolymorphic=*/false);
-              };
-              auto lhs = isWholeAllocatable(assign.lhs)
-                             ? lowerAllocatableLHS()
-                             : genExprAddr(assign.lhs, stmtCtx);
+              bool lhsIsWholeAllocatable = isWholeAllocatable(assign.lhs);
+              llvm::Optional<fir::factory::MutableBoxReallocation> lhsRealloc;
+              llvm::Optional<fir::MutableBoxValue> lhsMutableBox;
+              auto lhs = [&]() -> fir::ExtendedValue {
+                if (lhsIsWholeAllocatable) {
+                  lhsMutableBox = genExprMutableBox(loc, assign.lhs);
+                  llvm::SmallVector<mlir::Value> lengthParams;
+                  if (auto *charBox = rhs.getCharBox())
+                    lengthParams.push_back(charBox->getLen());
+                  else if (fir::isDerivedWithLengthParameters(rhs))
+                    TODO(loc, "assignment to derived type allocatable with "
+                              "length parameters");
+                  lhsRealloc = fir::factory::genReallocIfNeeded(
+                      *builder, loc, *lhsMutableBox,
+                      /*shape=*/llvm::None, lengthParams);
+                  return lhsRealloc->newValue;
+                }
+                return genExprAddr(assign.lhs, stmtCtx);
+              }();
 
               if (isNumericScalar) {
                 // Fortran 2018 10.2.1.3 p8 and p9
@@ -1983,21 +1982,22 @@ class FirConverter : public Fortran::lower::AbstractConverter {
                       toLocation(), builder->getRefType(toTy), addr);
                 }
                 builder->create<fir::StoreOp>(loc, cast, addr);
-                return;
-              }
-              if (isCharacterCategory(lhsType->category())) {
+              } else if (isCharacterCategory(lhsType->category())) {
                 // Fortran 2018 10.2.1.3 p10 and p11
                 fir::factory::CharacterExprHelper{*builder, loc}.createAssign(
                     lhs, rhs);
-                return;
-              }
-              if (isDerivedCategory(lhsType->category())) {
+              } else if (isDerivedCategory(lhsType->category())) {
                 // Fortran 2018 10.2.1.3 p13 and p14
                 // Recursively gen an assignment on each element pair.
                 fir::factory::genRecordAssignment(*builder, loc, lhs, rhs);
-                return;
+              } else {
+                llvm_unreachable("unknown category");
               }
-              llvm_unreachable("unknown category");
+              if (lhsIsWholeAllocatable)
+                fir::factory::finalizeRealloc(
+                    *builder, loc, lhsMutableBox.getValue(),
+                    /*lbounds=*/llvm::None, /*takeLboundsIfRealloc=*/false,
+                    lhsRealloc.getValue());
             },
 
             // [2] User defined assignment. If the context is a scalar
diff --git a/flang/lib/Lower/ConvertExpr.cpp b/flang/lib/Lower/ConvertExpr.cpp
@@ -2741,12 +2741,11 @@ class ArrayExprLowering {
       auto lbs = fir::factory::getOrigins(arrayOperands[0].shape);
       lbounds.append(lbs.begin(), lbs.end());
     }
-    fir::factory::genReallocIfNeeded(builder, loc, mutableBox, lbounds,
-                                     destShape, lengthParams);
+    auto realloc = fir::factory::genReallocIfNeeded(builder, loc, mutableBox,
+                                                    destShape, lengthParams);
     // Create ArrayLoad for the mutable box and save it into `destination`.
     PushSemantics(ConstituentSemantics::ProjectedCopyInCopyOut);
-    ccStoreToDest =
-        genarr(fir::factory::genMutableBoxRead(builder, loc, mutableBox));
+    ccStoreToDest = genarr(realloc.newValue);
     // If the rhs is scalar, get shape from the allocatable ArrayLoad.
     if (destShape.empty())
       destShape = getShape(destination);
@@ -2761,6 +2760,8 @@ class ArrayExprLowering {
           loc, destination, fir::getBase(exv), destination.memref(),
           destination.slice(), destination.typeparams());
     }
+    fir::factory::finalizeRealloc(builder, loc, mutableBox, lbounds,
+                                  takeLboundsIfRealloc, realloc);
   }
 
   /// Entry point for when an array expression appears in a context where the
diff --git a/flang/lib/Optimizer/Builder/MutableBox.cpp b/flang/lib/Optimizer/Builder/MutableBox.cpp
diff --git a/flang/test/Lower/allocatable-assignment.f90 b/flang/test/Lower/allocatable-assignment.f90
