[flang][hlfir] address post-commit comments from D151247 and D151251

Addresses comments not addressed in https://reviews.llvm.org/D151251
and https://reviews.llvm.org/D151247

- Fix typo in comments.
- Update an expected test output to include the fir.allocmem argument.
- Make a more generic type comparisons and cast when fetching value
  back from the AnyValueStack temporary storage.

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

flang/lib/Optimizer/Builder/TemporaryStorage.cpp

@@ -208,12 +208,22 @@ mlir::Value fir::factory::AnyValueStack::fetch(mlir::Location loc,
                                                fir::FirOpBuilder &builder) {
   mlir::Value indexValue = counter.getAndIncrementIndex(loc, builder);
   fir::runtime::genValueAt(loc, builder, opaquePtr, indexValue, retValueBox);
-  /// Dereference the allocatable "retValueBox", and load if trivial scalar
-  /// value.
+  // Dereference the allocatable "retValueBox", and load if trivial scalar
+  // value.
   mlir::Value result =
       hlfir::loadTrivialScalar(loc, builder, hlfir::Entity{retValueBox});
-  if (valueStaticType == builder.getI1Type())
-    return builder.createConvert(loc, valueStaticType, result);
+  if (valueStaticType != result.getType()) {
+    // Cast back saved simple scalars stored with another type to their original
+    // type (like i1).
+    if (fir::isa_trivial(valueStaticType))
+      return builder.createConvert(loc, valueStaticType, result);
+    // Memory type mismatches (e.g. fir.ref vs fir.heap) or hlfir.expr vs
+    // variable type mismatches are OK, but the base Fortran type must be the
+    // same.
+    assert(hlfir::getFortranElementOrSequenceType(valueStaticType) ==
+               hlfir::getFortranElementOrSequenceType(result.getType()) &&
+           "non trivial values must be saved with their original type");
+  }
   return result;
 }
 

flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp

@@ -146,7 +146,8 @@ class OrderedAssignmentRewriter {
 
   /// Is this an assignment to a vector subscripted entity?
   static bool hasVectorSubscriptedLhs(hlfir::RegionAssignOp regionAssignOp);
-  /// Are they any leaf region in node that must be saved in the current run?
+  /// Are there any leaf region in the node that must be saved in the current
+  /// run?
   bool mustSaveRegionIn(
       hlfir::OrderedAssignmentTreeOpInterface node,
       llvm::SmallVectorImpl<hlfir::SaveEntity> &saveEntities) const;
@@ -216,7 +217,7 @@ class OrderedAssignmentRewriter {
   }
 
   /// Can the current loop nest iteration number be computed? For simplicity,
-  /// this is true if an only if all the bounds and steps of the fir.do_loop
+  /// this is true if and only if all the bounds and steps of the fir.do_loop
   /// nest dominates the outer loop. The argument is filled with the current
   /// loop nest on success.
   bool currentLoopNestIterationNumberCanBeComputed(
@@ -250,10 +251,10 @@ class OrderedAssignmentRewriter {
   /// Map of temporary storage to keep track of saved entity once the run
   /// that saves them has been lowered. It is kept in-between runs.
   llvm::DenseMap<mlir::Region *, fir::factory::TemporaryStorage> savedEntities;
-  /// Map holding the value that were saved in the current run and that also
+  /// Map holding the values that were saved in the current run and that also
   /// need to be used (because their construct will be visited). It is reset
   /// after each run. It avoids having to store and fetch in the temporary
-  /// during the same run, which would required the temporary to have different
+  /// during the same run, which would require the temporary to have different
   /// fetching and storing counters.
   llvm::DenseMap<mlir::Region *, ValueAndCleanUp> savedInCurrentRunBeforeUse;
 

flang/test/HLFIR/order_assignments/impure-where.fir

@@ -34,7 +34,8 @@ func.func @test_elsewhere_impure_mask(%x: !fir.ref<!fir.array<10xi32>>, %y: !fir
 }
 // CHECK-LABEL:   func.func @test_elsewhere_impure_mask(
 // CHECK:           %[[VAL_12:.*]] = fir.call @impure() : () -> !fir.heap<!fir.array<10x!fir.logical<4>>>
-// CHECK:           %[[VAL_21:.*]] = fir.allocmem !fir.array<?x!fir.logical<4>>
+// CHECK:           %[[VAL_21:.*]] = fir.allocmem !fir.array<?x!fir.logical<4>>, %[[extent:[^ ]*]]
+// CHECK:           %[[VAL_22:.*]] = fir.shape %[[extent]] : (index) -> !fir.shape<1>
 // CHECK:           %[[VAL_23:.*]]:2 = hlfir.declare %[[VAL_21]](%{{.*}}) {uniq_name = ".tmp.forall"}
 // CHECK:           fir.do_loop
 // CHECK:             fir.if {{.*}} {