[Feature Request] Assert max allocations in unit tests

[sa-]

Review Mojo's priorities

• [X] I have read the roadmap and priorities and I believe this request falls within the priorities.

What is your request?

Description

As a developer using Mojo, I would like to be able to unit test the number of allocations made by a function, maybe similar to how Julia's @time macro allows for measuring allocations. It would allow for more robust performance testing and ensure that functions are not making unnecessary copies or allocations. I can see this being particularly useful when someone accidentally introduces a performance regression but the unit test is able to catch it before the code is ever merged or published.

Proposed Implementation

One possible implementation of this feature could be the addition of a new @allocations macro, which would allow developers to assert the number of allocations made by a function in a unit test. Below are some possible ways to express this. I'm not sure if either of these are the best, and readers are encouraged to contribute ideas of their own!

Option 1: if Mojo has macros

def test_my_function():
    @assert_max_allocations("64B") my_function()

Option 2: A pythonic take using with

def test_my_function():
    with assert_max_allocations("64B"): 
        my_function()

[edit from feedback] Option 3: Same as option 2 but without strings

def test_my_function():
    with assert_max_allocations(64):
        my_function()

In these examples, assert_max_allocations asserts that my_function should make at most 64B of allocations. If the function makes any more allocations, the unit test will fail.

What is your motivation for this change?

Currently, when writing performance-critical code, developers often have to rely on other tools (e.g. xctrace on macos) to measure allocations and ensure that they are minimized. But these tools are not standard across platforms. Having the ability to assert the number of allocations in unit tests would

1. Make it easier to catch performance issues early on in the development process
2. Create a standard way of doing this across platforms

Any other details?

Julia's @time macro in action

julia> @time h()
  1.151921 seconds (915.60 k allocations: 48.166 MiB, 1.64% gc time)
(0.5, 0.0)
 
julia> @time h()
  0.000013 seconds (21 allocations: 720 bytes)
(0.5, 0.0)