make @allocated x = f(...) work as before

Also add `@constprop :none` and `@noinline` to prevent the measuring code from
interfering with the subject code.

Fixes #58780

(cherry picked from commit 23633c2)

Author: JeffBezanson

base/reduce.jl

@@ -212,7 +212,7 @@ Like [`mapreduce`](@ref), but with guaranteed right associativity, as in [`foldr
 provided, the keyword argument `init` will be used exactly once. In general, it will be
 necessary to provide `init` to work with empty collections.
 """
-mapfoldr(f, op, itr; init=_InitialValue()) = mapfoldr_impl(f, op, init, itr)
+mapfoldr(f::F, op::F2, itr; init=_InitialValue()) where {F,F2} = mapfoldr_impl(f, op, init, itr)
 
 
 """
@@ -231,7 +231,7 @@ julia> foldr(=>, 1:4; init=0)
 1 => (2 => (3 => (4 => 0)))
 ```
 """
-foldr(op, itr; kw...) = mapfoldr(identity, op, itr; kw...)
+foldr(op::F, itr; kw...) where {F} = mapfoldr(identity, op, itr; kw...)
 
 ## reduce & mapreduce
 

base/timing.jl

@@ -472,41 +472,93 @@ function gc_bytes()
     b[]
 end
 
-function allocated(f, args::Vararg{Any,N}) where {N}
+@constprop :none function allocated(f, args::Vararg{Any,N}) where {N}
     b0 = Ref{Int64}(0)
     b1 = Ref{Int64}(0)
     Base.gc_bytes(b0)
-    f(args...)
+    @noinline f(args...)
     Base.gc_bytes(b1)
     return b1[] - b0[]
 end
 only(methods(allocated)).called = 0xff
 
-function allocations(f, args::Vararg{Any,N}) where {N}
+@constprop :none function allocations(f, args::Vararg{Any,N}) where {N}
     stats = Base.gc_num()
-    f(args...)
+    @noinline f(args...)
     diff = Base.GC_Diff(Base.gc_num(), stats)
     return Base.gc_alloc_count(diff)
 end
 only(methods(allocations)).called = 0xff
 
 function is_simply_call(@nospecialize ex)
+    is_simple_atom(a) = a isa QuoteNode || a isa Symbol || is_self_quoting(a)
     Meta.isexpr(ex, :call) || return false
     for a in ex.args
-        a isa QuoteNode && continue
-        a isa Symbol && continue
-        Base.is_self_quoting(a) && continue
+        is_simple_atom(a) && continue
+        Meta.isexpr(a, :..., 1) && is_simple_atom(a.args[1]) && continue
         return false
     end
     return true
 end
 
+function _gen_allocation_measurer(ex, fname::Symbol)
+    if isexpr(ex, :call)
+        if !is_simply_call(ex)
+            ex = :((() -> $ex)())
+        end
+        pushfirst!(ex.args, GlobalRef(Base, fname))
+        return quote
+            Experimental.@force_compile
+            $(esc(ex))
+        end
+    elseif fname === :allocated
+        # v1.11-compatible implementation
+        return quote
+            Experimental.@force_compile
+            local b0 = Ref{Int64}(0)
+            local b1 = Ref{Int64}(0)
+            gc_bytes(b0)
+            $(esc(ex))
+            gc_bytes(b1)
+            b1[] - b0[]
+        end
+    else
+        @assert fname === :allocations
+        return quote
+            Experimental.@force_compile
+            # Note this value is unused, but without it `allocated` and `allocations`
+            # are sufficiently different that the compiler can remove allocations here
+            # that it cannot remove there, giving inconsistent numbers.
+            local b1 = Ref{Int64}(0)
+            local stats = Base.gc_num()
+            $(esc(ex))
+            local diff = Base.GC_Diff(Base.gc_num(), stats)
+            gc_bytes(b1)
+            Base.gc_alloc_count(diff)
+        end
+    end
+end
+
 """
     @allocated
 
 A macro to evaluate an expression, discarding the resulting value, instead returning the
 total number of bytes allocated during evaluation of the expression.
 
+If the expression is a function call, an effort is made to measure only allocations from
+the argument expressions and during the function, excluding any overhead from calling it
+and not performing constant propagation with the provided argument values. If you want to
+include those effects, i.e. measuring the call site as well, use the syntax
+`@allocated (()->f(1))()`.
+
+It is recommended to measure function calls with only simple argument expressions, e.g.
+`x = []; @allocated f(x)` instead of `@allocated f([])` to clarify that only `f` is
+being measured.
+
+For more complex expressions, the code is simply run in place and therefore may see
+allocations due to the surrounding context. For example it is possible for
+`@allocated f(1)` and `@allocated x = f(1)` to give different results.
+
 See also [`@allocations`](@ref), [`@time`](@ref), [`@timev`](@ref), [`@timed`](@ref),
 and [`@elapsed`](@ref).
 
@@ -516,11 +568,7 @@ julia> @allocated rand(10^6)
 ```
 """
 macro allocated(ex)
-    if !is_simply_call(ex)
-        ex = :((() -> $ex)())
-    end
-    pushfirst!(ex.args, GlobalRef(Base, :allocated))
-    return esc(ex)
+    _gen_allocation_measurer(ex, :allocated)
 end
 
 """
@@ -541,11 +589,7 @@ julia> @allocations rand(10^6)
     This macro was added in Julia 1.9.
 """
 macro allocations(ex)
-    if !is_simply_call(ex)
-        ex = :((() -> $ex)())
-    end
-    pushfirst!(ex.args, GlobalRef(Base, :allocations))
-    return esc(ex)
+    _gen_allocation_measurer(ex, :allocations)
 end
 
 

test/boundscheck_exec.jl

@@ -349,8 +349,9 @@ if bc_opt == bc_default
         m2 = Memory{Int}(undef,n)
         m1 === m2
     end
-    no_alias_prove(1)
-    @test (@allocated no_alias_prove(5)) == 0
+    no_alias_prove5() = no_alias_prove(5)
+    no_alias_prove5()
+    @test (@allocated no_alias_prove5()) == 0
 end
 
 end

test/misc.jl

@@ -1522,15 +1522,40 @@ end
 # issue #41656
 run(`$(Base.julia_cmd()) -e 'isempty(x) = true'`)
 
+function treshape59278(X::AbstractArray, n, m)
+    Y = reshape(X, n, m)
+    Y .= 1.0
+    return X
+end
+
+# a function that allocates iff no constprop
+@inline maybealloc59278(n, _) = ntuple(i->rand(), n)
+
 @testset "Base/timing.jl" begin
     @test Base.jit_total_bytes() >= 0
 
     # sanity check `@allocations` returns what we expect in some very simple cases.
-    # These are inside functions because `@allocations` uses `Experimental.@force_compile`
-    # so can be affected by other code in the same scope.
     @test (() -> @allocations "a")() == 0
-    @test (() -> @allocations "a" * "b")() == 0 # constant propagation
+    "a" * Base.inferencebarrier("b")
     @test (() -> @allocations "a" * Base.inferencebarrier("b"))() == 1
+    # test that you can grab the value from @allocated
+    @allocated _x = 1+2
+    @test _x === 3
+
+    n, m = 10, 20
+    X = rand(n, m)
+    treshape59278(X, n, m)
+    # test that @allocated and @allocations are consistent about whether anything was
+    # allocated in a case where the compiler can sometimes remove an allocation
+    # https://github.com/JuliaLang/julia/issues/58634#issuecomment-2940840651
+    @test ((@allocated treshape59278(X, n, m))==0) == ((@allocations treshape59278(X, n, m))==0)
+    # TODO: would be nice to have but not yet reliable
+    #@test ((@allocated begin treshape59278(X, n, m) end)==0) == ((@allocations begin treshape59278(X, n, m) end)==0)
+
+    # test that all wrapped allocations are counted and constprop is not done
+    @test (@allocated @noinline maybealloc59278(10, [])) > (@allocated maybealloc59278(10, 0)) > 0
+    # but if you wrap it in another function it can be constprop'd
+    @test (@allocated (()->maybealloc59278(10, []))()) == 0
 
     _lock_conflicts, _nthreads = eval(Meta.parse(read(`$(Base.julia_cmd()) -tauto -E '
         _lock_conflicts = @lock_conflicts begin