[flang] Deallocate WHERE masks after all assignments

[jeanPerier]

Summary

Allocatable assignments were triggering lowering to clean-up
any WHERE mask temporaries, even if some assignments where left
in the WHERE construct.

This is because allocatable assignments lowering was being passed the
wrong statement context. Fix this by selecting the where/forall statement
context instead of a local one when there is one.

[clementval]

LGTM

[PeteSteinfeld]

All builds and tests correctly and looks good.

flang/lib/Lower/Bridge.cpp

   /// This is an assignment expression with rank > 0. The assignment may or may
   /// not be in a WHERE and/or FORALL context.
   /// In a FORALL context, the assignment may be a pointer assignment and the \p
   /// lbounds and \p ubounds parameters should only be used in such a pointer
   /// assignment case. (If both are None then the array assignment cannot be a
   /// pointer assignment.)
   void genArrayAssignment(
       const Fortran::evaluate::Assignment &assign,
-      Fortran::lower::StatementContext &stmtCtx,
+      Fortran::lower::StatementContext &localStmtCtx,
       llvm::Optional<llvm::SmallVector<mlir::Value>> lbounds = llvm::None,
       llvm::Optional<llvm::SmallVector<mlir::Value>> ubounds = llvm::None) {
+
+    Fortran::lower::StatementContext &stmtCtx =
+        explicitIterationSpace()
+            ? explicitIterSpace.stmtContext()
+            : (implicitIterationSpace() ? implicitIterSpace.stmtContext()
+                                        : localStmtCtx);
     if (Fortran::lower::isWholeAllocatable(assign.lhs)) {
       // Assignment to allocatables may require the lhs to be
       // deallocated/reallocated. See Fortran 2018 10.2.1.3 p3
       Fortran::lower::createAllocatableArrayAssignment(
           *this, assign.lhs, assign.rhs, explicitIterSpace, implicitIterSpace,
           localSymbols, stmtCtx);
       return;
     }
void genArrayAssignment(
     // expression is evaluated at all points in the user-defined space much like
     // an ordinary array assignment. More specifically, the semantics inside the
     // FORALL much more closely resembles that of WHERE than a scalar
     // assignment.
     // Otherwise, generate a masked array assignment. The iteration space is
     // implied by the lhs array expression.
     Fortran::lower::createAnyMaskedArrayAssignment(
         *this, assign.lhs, assign.rhs, explicitIterSpace, implicitIterSpace,
-        localSymbols,
-        explicitIterationSpace() ? explicitIterSpace.stmtContext()
-                                 : implicitIterSpace.stmtContext());
+        localSymbols, stmtCtx);
   }
 
 #if !defined(NDEBUG)
   static bool isFuncResultDesignator(const Fortran::lower::SomeExpr &expr) {
     const Fortran::semantics::Symbol *sym =
         Fortran::evaluate::GetFirstSymbol(expr);
     return sym && sym->IsFuncResult();
   }

flang/test/Lower/where-allocatable-assignments.f90

+! Test that WHERE mask clean-up occurs at the right time when the
+! WHERE contains whole allocatable assignments.
+! RUN: bbc -emit-fir %s -o - | FileCheck %s
+
+module mtest
+contains
+
+! CHECK-LABEL: func.func @_QMmtestPfoo(
+! CHECK-SAME:       %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>> {fir.bindc_name = "a"},
+! CHECK-SAME:       %[[VAL_1:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>> {fir.bindc_name = "b"}) {
+subroutine foo(a, b)
+  integer :: a(:)
+  integer, allocatable :: b(:)
+! CHECK:  %[[VAL_2:.*]] = fir.load %[[VAL_1]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+! CHECK:  %[[VAL_3:.*]] = arith.constant 0 : index
+! CHECK:  %[[VAL_4:.*]]:3 = fir.box_dims %[[VAL_2]], %[[VAL_3]] : (!fir.box<!fir.heap<!fir.array<?xi32>>>, index) -> (index, index, index)
+          ! WHERE mask temp allocation
+! CHECK:  %[[VAL_9:.*]] = fir.allocmem !fir.array<?x!fir.logical<4>>, %[[VAL_4]]#1 {uniq_name = ".array.expr"}
+! CHECK:  %[[VAL_15:.*]] = fir.do_loop {{.*}} {
+!           ! WHERE mask element computation
+! CHECK:  }
+! CHECK:  fir.array_merge_store %{{.*}}, %[[VAL_15]] to %[[VAL_9]] : !fir.array<?x!fir.logical<4>>, !fir.array<?x!fir.logical<4>>, !fir.heap<!fir.array<?x!fir.logical<4>>>
+
+          ! First assignment to a whole allocatable (in WHERE)
+! CHECK:  fir.if {{.*}} {
+! CHECK:    fir.if {{.*}} {
+            ! assignment into new storage (`b` allocated with bad shape)
+! CHECK:      fir.allocmem
+! CHECK:      fir.do_loop {{.*}} {
+! CHECK:        fir.array_coor %[[VAL_9]]
+! CHECK:        fir.if %{{.*}} {
+                  ! WHERE
+! CHECK:          fir.array_update {{.*}}
+! CHECK:        } else {
+! CHECK:        }
+! CHECK:      }
+! CHECK:    } else {
+              ! assignment into old storage (`b` allocated with the same shape)
+! CHECK:      fir.do_loop {{.*}} {
+! CHECK:        fir.array_coor %[[VAL_9]]
+! CHECK:        fir.if %{{.*}} {
+                  ! WHERE
+! CHECK:          fir.array_update {{.*}}
+! CHECK:        } else {
+! CHECK:        }
+! CHECK:      }
+! CHECK:    }
+! CHECK:  } else {
+            ! assignment into new storage (`b` unallocated)
+! CHECK:    fir.allocmem
+! CHECK:    fir.do_loop %{{.*}} {
+! CHECK:      fir.array_coor %[[VAL_9]]
+! CHECK:      fir.if %{{.*}} {
+                ! WHERE
+! CHECK:        fir.array_update {{.*}}
+! CHECK:      } else {
+! CHECK:      }
+! CHECK:    }
+! CHECK:  }
+! CHECK:  fir.if {{.*}} {
+! CHECK:    fir.if {{.*}} {
+              ! deallocation of `b` old allocatable data store
+! CHECK:    }
+            ! update of `b` descriptor
+! CHECK:  }
+          ! Second assignment (in ELSEWHERE)
+! CHECK:  fir.do_loop {{.*}} {
+! CHECK:    fir.array_coor %[[VAL_9]]{{.*}} : (!fir.heap<!fir.array<?x!fir.logical<4>>>, !fir.shape<1>, index) -> !fir.ref<!fir.logical<4>>
+! CHECK:    fir.if {{.*}} {
+! CHECK:    } else {
+              ! elsewhere
+! CHECK:      fir.array_update
+! CHECK:    }
+! CHECK:  }
+          ! WHERE temp clean-up
+! CHECK:  fir.freemem %[[VAL_9]] : !fir.heap<!fir.array<?x!fir.logical<4>>>
+! CHECK-NEXT:  return
+  where (b > 0)
+    b = a
+  elsewhere
+    b(:) = 0
+  end where
+end
+end module
+
+  use mtest
+  integer, allocatable :: a(:), b(:)
+  allocate(a(10),b(10))
+  a = 5
+  b = 1
+  call foo(a, b)
+  print*, b
+  deallocate(a,b)
+end