Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add @check_allocs_call #46

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Add @check_allocs_call #46

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/AllocCheck.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
27 changes: 27 additions & 0 deletions src/macro.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}()
Expand Down Expand Up @@ -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...)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this can be a method/variant of @check_allocs. There is precedent for that with e.g. @timeit from TimerOutputs.jl: when applied to a function definition expression it enables timings for all calls, when applied to a call expression it times just that call.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good idea, I wasn't aware this was being done elsewhere but seems relatively straightforward to explain, with no specific edge cases I could think of

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