diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp --- a/flang/lib/Lower/Bridge.cpp +++ b/flang/lib/Lower/Bridge.cpp @@ -2188,9 +2188,15 @@ /// pointer assignment.) void genArrayAssignment( const Fortran::evaluate::Assignment &assign, - Fortran::lower::StatementContext &stmtCtx, + Fortran::lower::StatementContext &localStmtCtx, llvm::Optional> lbounds = llvm::None, llvm::Optional> 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 @@ -2230,9 +2236,7 @@ // 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) diff --git a/flang/test/Lower/where-allocatable-assignments.f90 b/flang/test/Lower/where-allocatable-assignments.f90 new file mode 100644 --- /dev/null +++ b/flang/test/Lower/where-allocatable-assignments.f90 @@ -0,0 +1,94 @@ +! 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.bindc_name = "a"}, +! CHECK-SAME: %[[VAL_1:.*]]: !fir.ref>>> {fir.bindc_name = "b"}) { +subroutine foo(a, b) + integer :: a(:) + integer, allocatable :: b(:) +! CHECK: %[[VAL_2:.*]] = fir.load %[[VAL_1]] : !fir.ref>>> +! CHECK: %[[VAL_3:.*]] = arith.constant 0 : index +! CHECK: %[[VAL_4:.*]]:3 = fir.box_dims %[[VAL_2]], %[[VAL_3]] : (!fir.box>>, index) -> (index, index, index) + ! WHERE mask temp allocation +! CHECK: %[[VAL_9:.*]] = fir.allocmem !fir.array>, %[[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>, !fir.array>, !fir.heap>> + + ! 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.shape<1>, index) -> !fir.ref> +! CHECK: fir.if {{.*}} { +! CHECK: } else { + ! elsewhere +! CHECK: fir.array_update +! CHECK: } +! CHECK: } + ! WHERE temp clean-up +! CHECK: fir.freemem %[[VAL_9]] : !fir.heap>> +! 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