Add @test_noalloc convenience

For use in unit-testing, @test_noalloc is implemented as a way to check
that a given call is not allocating. To avoid extra dependencies, this
is implemented as a package extension (loaded with Test). Since package
extensions can not themselves export bindings, the macro is implemented
in AllocCheck.jl, but emits an informative error when called without the
extension loaded.

Author: tecosaur

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