"improve" allocations macro with more functions (#58057)

Change the implementation of the allocations to always add a function
call (closure) wrapper so that the inner code can be optimized, even if
the outer scope is `@latestworld`. The implementation creates a closure
if necessary (the code is complex) or just makes a call if the
expression is simple (just a call on symbols). Many packages in the
wild are observed to already be doing this themselves.

Author: vtjnash

base/timing.jl

@@ -472,6 +472,35 @@ function gc_bytes()
     b[]
 end
 
+function allocated(f, args::Vararg{Any,N}) where {N}
+    b0 = Ref{Int64}(0)
+    b1 = Ref{Int64}(0)
+    Base.gc_bytes(b0)
+    f(args...)
+    Base.gc_bytes(b1)
+    return b1[] - b0[]
+end
+only(methods(allocated)).called = 0xff
+
+function allocations(f, args::Vararg{Any,N}) where {N}
+    stats = Base.gc_num()
+    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)
+    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
+        return false
+    end
+    return true
+end
+
 """
     @allocated
 
@@ -487,15 +516,11 @@ julia> @allocated rand(10^6)
 ```
 """
 macro allocated(ex)
-    quote
-        Experimental.@force_compile
-        local b0 = Ref{Int64}(0)
-        local b1 = Ref{Int64}(0)
-        gc_bytes(b0)
-        $(esc(ex))
-        gc_bytes(b1)
-        b1[] - b0[]
+    if !is_simply_call(ex)
+        ex = :((() -> $ex)())
     end
+    pushfirst!(ex.args, GlobalRef(Base, :allocated))
+    return esc(ex)
 end
 
 """
@@ -516,15 +541,14 @@ julia> @allocations rand(10^6)
     This macro was added in Julia 1.9.
 """
 macro allocations(ex)
-    quote
-        Experimental.@force_compile
-        local stats = Base.gc_num()
-        $(esc(ex))
-        local diff = Base.GC_Diff(Base.gc_num(), stats)
-        Base.gc_alloc_count(diff)
+    if !is_simply_call(ex)
+        ex = :((() -> $ex)())
     end
+    pushfirst!(ex.args, GlobalRef(Base, :allocations))
+    return esc(ex)
 end
 
+
 """
     @lock_conflicts
 

test/boundscheck_exec.jl

@@ -350,7 +350,7 @@ if bc_opt == bc_default
         m1 === m2
     end
     no_alias_prove(1)
-    @test_broken (@allocated no_alias_prove(5)) == 0
+    @test (@allocated no_alias_prove(5)) == 0
 end
 
 end

test/math.jl

@@ -46,8 +46,7 @@ has_fma = Dict(
         @test clamp(100, Int8) === Int8(100)
         @test clamp(200, Int8) === typemax(Int8)
 
-        begin
-            x = [0.0, 1.0, 2.0, 3.0, 4.0]
+        let x = [0.0, 1.0, 2.0, 3.0, 4.0]
             clamp!(x, 1, 3)
             @test x == [1.0, 1.0, 2.0, 3.0, 3.0]
         end
@@ -59,12 +58,14 @@ has_fma = Dict(
         @test clamp(typemax(UInt16), Int16) === Int16(32767)
 
         # clamp should not allocate a BigInt for typemax(Int16)
-        x = big(2) ^ 100
-        @test (@allocated clamp(x, Int16)) == 0
+        let x = big(2) ^ 100
+            @test (@allocated clamp(x, Int16)) == 0
+        end
 
-        x = clamp(2.0, BigInt)
-        @test x isa BigInt
-        @test x == big(2)
+        let x = clamp(2.0, BigInt)
+            @test x isa BigInt
+            @test x == big(2)
+        end
     end
 end