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

[jeanPerier]

Summary

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.

[tblah]

Thanks for fixing these nitpics!

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

149–150

flang/lib/Optimizer/Builder/TemporaryStorage.cpp

void fir::factory::AnyValueStack::resetFetchPosition(
     mlir::Location loc, fir::FirOpBuilder &builder) {
   counter.reset(loc, builder);
 }
 
 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())
+  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;
 }
 
 void fir::factory::AnyValueStack::destroy(mlir::Location loc,
                                           fir::FirOpBuilder &builder) {
   fir::runtime::genDestroyValueStack(loc, builder, opaquePtr);
 }

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

private:
   void post(hlfir::WhereOp);
   void post(hlfir::ElseWhereOp);
   /// Enter (and maybe create) the fir.if else block of an ElseWhereOp,
   /// but do not generate the elswhere mask or the new fir.if.
   void enterElsewhere(hlfir::ElseWhereOp);
 
   /// 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;
   /// Should this node be evaluated in the current run? Saving a region in a
   /// node does not imply the node needs to be evaluated.
   bool
   isRequiredInCurrentRun(hlfir::OrderedAssignmentTreeOpInterface node) const;
 
void doBeforeLoopNest(const std::function<void()> &callback) {
     }
     auto insertionPoint = builder.saveInsertionPoint();
     builder.setInsertionPoint(constructStack[0]);
     callback();
     builder.restoreInsertionPoint(insertionPoint);
   }
 
   /// 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(
       llvm::SmallVectorImpl<fir::DoLoopOp> &loopNest);
 
   template <typename T>
   fir::factory::TemporaryStorage *insertSavedEntity(mlir::Region &region,
                                                     T &&temp) {
private:
   /// operation.
   llvm::SmallVector<mlir::Operation *> constructStack;
   /// Current where loop nest, if any.
   std::optional<hlfir::LoopNest> whereLoopNest;
 
   /// 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;
 
   /// Root of the order assignment tree being lowered.
   hlfir::OrderedAssignmentTreeOpInterface root;
   /// Pointer to the current run of the schedule being lowered.
   hlfir::Run *currentRun = nullptr;
 

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

hlfir.elsewhere mask {
         hlfir.yield %z : !fir.ref<!fir.array<10xi32>>
       }
     }
   }
   return
 }
 // 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 {{.*}} {
 // CHECK:             } else {
 // CHECK:               %[[VAL_28:.*]] = hlfir.designate %[[VAL_12]] (%{{.*}})
 // CHECK:               %[[VAL_29:.*]] = fir.load %[[VAL_28]] : !fir.ref<!fir.logical<4>>
 // CHECK:               %[[VAL_32:.*]] = hlfir.designate %[[VAL_23]]#0 (%{{.*}})  : (!fir.box<!fir.array<?x!fir.logical<4>>>, index) -> !fir.ref<!fir.logical<4>>
 // CHECK:               hlfir.assign %[[VAL_29]] to %[[VAL_32]] : !fir.logical<4>, !fir.ref<!fir.logical<4>>