[flang] convert stack arrays allocation to match old type

[tblah]

Summary

The old fir.allocmem operation returned a !fir.heap<.> type. The new
fir.alloca operation returns a !fir.ref<.> type. This patch inserts a
fir.convert so that the old type is preserved. This prevents verifier
failures when types returned from fir.if statements don't match the
expected type.

[vzakhari]

LGTM, except for the casts. Thank you!

flang/lib/Optimizer/Transforms/StackArrays.cpp

478

Please use new style for the casts (see https://reviews.llvm.org/D151556).

[tblah]

Thanks for review!

Changes: use the free function variants of casts

[vzakhari]

Thank you, Tom!

flang/lib/Optimizer/Transforms/StackArrays.cpp

 const StackArraysAnalysisWrapper::AllocMemMap *
 StackArraysAnalysisWrapper::getCandidateOps(mlir::Operation *func) {
   if (!funcMaps.contains(func))
     if (mlir::failed(analyseFunction(func)))
       return nullptr;
   return &funcMaps[func];
 }
 
+/// Restore the old allocation type exected by existing code
+static mlir::Value convertAllocationType(mlir::PatternRewriter &rewriter,
+                                         const mlir::Location &loc,
+                                         mlir::Value heap, mlir::Value stack) {
+  mlir::Type heapTy = heap.getType();
+  mlir::Type stackTy = stack.getType();
+
+  if (heapTy == stackTy)
+    return stack;
+
+  fir::HeapType firHeapTy = mlir::cast<fir::HeapType>(heapTy);
+  fir::ReferenceType firRefTy = mlir::cast<fir::ReferenceType>(stackTy);
+  assert(firHeapTy.getElementType() == firRefTy.getElementType() &&
+         "Allocations must have the same type");
+
+  auto insertionPoint = rewriter.saveInsertionPoint();
+  rewriter.setInsertionPointAfter(stack.getDefiningOp());
+  mlir::Value conv =
+      rewriter.create<fir::ConvertOp>(loc, firHeapTy, stack).getResult();
+  rewriter.restoreInsertionPoint(insertionPoint);
+  return conv;
+}
+
 mlir::LogicalResult
 AllocMemConversion::matchAndRewrite(fir::AllocMemOp allocmem,
                                     mlir::PatternRewriter &rewriter) const {
   auto oldInsertionPt = rewriter.saveInsertionPoint();
   // add alloca operation
   std::optional<fir::AllocaOp> alloca = insertAlloca(allocmem, rewriter);
   rewriter.restoreInsertionPoint(oldInsertionPt);
   if (!alloca)
     return mlir::failure();
 
   // remove freemem operations
   llvm::SmallVector<mlir::Operation *> erases;
   for (mlir::Operation *user : allocmem.getOperation()->getUsers())
     if (mlir::isa<fir::FreeMemOp>(user))
       erases.push_back(user);
   // now we are done iterating the users, it is safe to mutate them
   for (mlir::Operation *erase : erases)
     rewriter.eraseOp(erase);
 
   // replace references to heap allocation with references to stack allocation
-  rewriter.replaceAllUsesWith(allocmem.getResult(), alloca->getResult());
+  mlir::Value newValue = convertAllocationType(
+      rewriter, allocmem.getLoc(), allocmem.getResult(), alloca->getResult());
+  rewriter.replaceAllUsesWith(allocmem.getResult(), newValue);
 
   // remove allocmem operation
   rewriter.eraseOp(allocmem.getOperation());
 
   return mlir::success();
 }
 
 static bool isInLoop(mlir::Block *block) {

flang/test/Transforms/stack-arrays.fir

fir.if %arg0 {
     %ref = fir.convert %a : (!fir.heap<!fir.array<1xi8>>) -> !fir.ref<!fir.array<1xi8>>
     func.call @dfa3a_bar(%ref) : (!fir.ref<!fir.array<1xi8>>) -> ()
     fir.freemem %a : !fir.heap<!fir.array<1xi8>>
   }
   return
 }
 // CHECK:     func.func @dfa3a(%arg0: i1) {
 // CHECK-NEXT:  %[[MEM:.*]] = fir.alloca !fir.array<1xi8>
+// CHECK-NEXT:  %[[HEAP:.*]] = fir.convert %[[MEM]] : (!fir.ref<!fir.array<1xi8>>) -> !fir.heap<!fir.array<1xi8>>
 // CHECK-NEXT:  fir.if %arg0 {
-// CHECK-NEXT:    func.call @dfa3a_foo(%[[MEM]])
+// CHECK-NEXT:    %[[REF:.*]] = fir.convert %[[HEAP]] : (!fir.heap<!fir.array<1xi8>>) -> !fir.ref<!fir.array<1xi8>>
+// CHECK-NEXT:    func.call @dfa3a_foo(%[[REF]])
 // CHECK-NEXT:  } else {
-// CHECK-NEXT:    func.call @dfa3a_bar(%[[MEM]])
+// CHECK-NEXT:    %[[REF:.*]] = fir.convert %[[HEAP]] : (!fir.heap<!fir.array<1xi8>>) -> !fir.ref<!fir.array<1xi8>>
+// CHECK-NEXT:    func.call @dfa3a_bar(%[[REF]])
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return
 // CHECK-NEXT:  }
 
 // check the alloca is placed after all operands become available
 func.func @placement1() {
   // do some stuff with other ssa values
   %1 = arith.constant 1 : index