Add @test_noalloc convenience

Project.toml

@@ -9,6 +9,12 @@ GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 
+[weakdeps]
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[extensions]
+TestAlloc = "Test"
+
 [compat]
 GPUCompiler = "0.24, 0.25"
 LLVM = "6.3"

ext/TestAlloc.jl

@@ -0,0 +1,49 @@
+module TestAlloc
+
+using AllocCheck
+using Test
+
+function AllocCheck._test_noalloc(__module__, __source__, expr, kws...)
+    # Collect the broken/skip/ignore_throw keywords and remove them from the rest of keywords
+    broken = [kw.args[2] for kw in kws if kw.args[1] === :broken]
+    skip = [kw.args[2] for kw in kws if kw.args[1] === :skip]
+    ignore_throw = [kw.args[2] for kw in kws if kw.args[1] === :ignore_throw]
+    kws = filter(kw -> kw.args[1] ∉ (:skip, :broken, :ignore_throw), kws)
+    # Validation of broken/skip keywords
+    for (kw, name) in ((broken, :broken), (skip, :skip), (ignore_throw, :ignore_throw))
+        if length(kw) > 1
+            error("invalid test_noalloc macro call: cannot set $(name) keyword multiple times")
+        end
+    end
+    if length(skip) > 0 && length(broken) > 0
+        error("invalid test_noalloc macro call: cannot set both skip and broken keywords")
+    end
+    if !Meta.isexpr(expr, :call) || isempty(expr.args)
+        error("invalid test_noalloc macro call: must be applied to a function call")
+    end
+    ex = Expr(:inert, Expr(:macrocall, Symbol("@test_noalloc"), nothing, expr))
+    quote
+        if $(length(skip) > 0 && esc(skip[1]))
+            $(Test.record)($(Test.get_testset)(), $(Test.Broken)(:skipped, $ex))
+        else
+            result = try
+                x = $(AllocCheck._check_allocs_call(
+                    expr, __module__, __source__; ignore_throw = length(ignore_throw) == 0 || ignore_throw[1]))
+                Base.donotdelete(x)
+                $(Test.Returned)(true, nothing, $(QuoteNode(__source__)))
+            catch err
+                if err isa InterruptException
+                    rethrow()
+                elseif err isa AllocCheck.AllocCheckFailure
+                    $(Test.Returned)(false, nothing, $(QuoteNode(__source__)))
+                else
+                    $(Test.Threw)(err, Base.current_exceptions(), $(QuoteNode(__source__)))
+                end
+            end
+            test_do = $(length(broken) > 0 && esc(broken[1])) ? $(Test.do_broken_test) : $(Test.do_test)
+            test_do(result, $ex)
+        end
+    end
+end
+
+end

src/AllocCheck.jl

@@ -215,7 +215,28 @@ function check_allocs(@nospecialize(func), @nospecialize(types); ignore_throw=tr
     return allocs
 end
 
+function _test_noalloc end # Implemented in `ext/TestAlloc.jl`
 
-export check_allocs, alloc_type, @check_allocs, AllocCheckFailure
+"""
+    @test_noalloc f(args...; kwargs...) key=val ...
+
+Test that `f(args...; kwargs...)` does not allocate. If executed inside a
+`@testset`, return a `Pass Result` if no allocations occur, a `Fail Result` if
+there are allocations, and a `Error Result` if any other errors occur during
+evaluation. If executed outside a `@testset` throw an exception instead of
+returning `Fail` or `Error`.
+
+The `broken` and `skip` keys can be set to `true`/`false`, and behave the same
+as in `@test`. The `ignore_throw` key can also be set to `true`/`false`, and is
+passed through to `@check_allocs`.
+"""
+macro test_noalloc(expr, kws...)
+    if length(methods(_test_noalloc)) == 0
+        error("@test_noalloc is an extension to Test, but Test is not loaded.")
+    end
+    _test_noalloc(__module__, __source__, expr, kws...)
+end
+
+export check_allocs, alloc_type, @check_allocs, AllocCheckFailure, @test_noalloc
 
 end

test/runtests.jl

@@ -273,3 +273,8 @@ end
     @test !allunique(allocs)
     @test length(unique(allocs)) == 2
 end
+
+@testset "@test_noalloc" begin
+    @test_noalloc 1 + 1
+    @test_noalloc rand(2, 2) * rand(2, 2) broken=true
+end