From 42b6be9b19525e7dd4732a4dbe7ba26a87442114 Mon Sep 17 00:00:00 2001
From: Julius Krumbiegel <julius.krumbiegel@gmail.com>
Date: Sun, 19 Nov 2023 19:59:36 +0100
Subject: [PATCH] add `@check_allocs_call`

---
 src/AllocCheck.jl |  2 +-
 src/macro.jl      | 27 +++++++++++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/src/AllocCheck.jl b/src/AllocCheck.jl
index 87a019c..6bfec17 100644
--- a/src/AllocCheck.jl
+++ b/src/AllocCheck.jl
@@ -205,6 +205,6 @@ function check_allocs(@nospecialize(func), @nospecialize(types); ignore_throw=tr
 end
 
 
-export check_allocs, alloc_type, @check_allocs, AllocCheckFailure
+export check_allocs, alloc_type, @check_allocs, AllocCheckFailure, @check_allocs_call
 
 end
diff --git a/src/macro.jl b/src/macro.jl
index bf46a31..424d781 100644
--- a/src/macro.jl
+++ b/src/macro.jl
@@ -2,6 +2,7 @@ using ExprTools: splitdef, combinedef
 using MacroTools: splitarg, combinearg
 
 _is_func_def(ex) = Meta.isexpr(ex, :function) || Base.is_short_function_def(ex) || Meta.isexpr(ex, :->)
+_is_call(ex) = Meta.isexpr(ex, :call)
 
 function extract_keywords(ex0)
     kws = Dict{Symbol, Any}()
@@ -154,3 +155,29 @@ function _check_allocs_macro(ex::Expr, mod::Module, source::LineNumberNode; igno
         $(wrapper_fn)
     end
 end
+
+"""
+    @check_allocs_call ignore_throw=true (function call)
+
+Calls `check_allocs(func, types; ignore_throw)` where `types` is a tuple of the types of the arguments
+passed to `func` in the function call expression that `@check_allocs_call` is applied to.
+For example, `@check_allocs_call 1 + 2.0` is equivalent to `check_allocs(+, (Int, Float64))`.
+"""
+macro check_allocs_call(ex...)
+    kws, body = extract_keywords(ex)
+    if _is_call(body)
+        return _check_allocs_call_macro(body, __module__, __source__; kws...)
+    else
+        error("@check_allocs_call used on something other than a function call")
+    end
+end
+
+function _check_allocs_call_macro(ex::Expr, mod::Module, source::LineNumberNode; ignore_throw=true)
+    func = ex.args[1]
+    args = ex.args[2:end]
+
+    quote
+        types = map(typeof, ($(esc.(args)...),))
+        check_allocs($(esc(func)), types; ignore_throw = $ignore_throw)
+    end
+end