reject all uses of @goto and return inside @no_escape (#19)

Project.toml

@@ -1,7 +1,7 @@
 name = "Bumper"
 uuid = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 authors = ["MasonProtter <[email protected]>"]
-version = "0.3.1"
+version = "0.3.2"
 
 [deps]
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"

README.org

@@ -3,10 +3,11 @@
 :END:
 * Bumper.jl
 
-Bumper.jl is an experimental package that aims to make working with bump allocators easy and safer (when used right).
-You can dynamically allocate memory to these bump allocators, and reset them at the end of a code block, just like
-Julia's default stack. Allocating to the a =AllocBuffer= with Bumper.jl can be just as efficient as stack allocation.
-The point of this is to not have to pay the hefty cost of intermediate allocations.
+Bumper.jl is an experimental package that aims to make working with bump allocators (also known as arena allocators)
+easy and safer (though not totally safe!). You can dynamically allocate memory to these bump allocators, and reset
+them at the end of a code block, just like Julia's default stack. Allocating to the a =AllocBuffer= with Bumper.jl
+can be just as efficient as stack allocation. The point of this is to not have to pay the hefty cost of
+intermediate allocations.
 
 Bumper.jl has a task-local default buffer, which can dynamically grow to be one eigth the size of your computer's
 physical memory pool. You can change the default buffer size with =set_default_buffer_size!(nbytes)= where =nbytes=
@@ -124,17 +125,16 @@ end
 
 Some miscellaneous notes:
 + =@no_escape= blocks can be nested as much as you want (so long as the allocator has enough memory to store the objects you're using.
-+ =alloc(T, n...)= is dynamically scoped, meaning that you can have deeply nested =alloc= calls inside a =@no_escape= block, and they'll
-  still use the same default buffer, and be reset once the block ends.
-+ Bumper.jl only supports =isbits= types. You cannot use it for allocating vectors of mutable, abstract, or other pointer-backed objects. 
-+ As mentioned previously, *Do not allow any memory which was initialized inside a* =@no_escape= *block to escape the block.* Doing so can cause memory
-  corruption.
++ =alloc(T, n...)= is dynamically scoped, meaning that you can have deeply nested =alloc= calls inside a =@no_escape= block, and they'll still use the same default buffer, and be reset once the block ends.
++ Bumper.jl only supports =isbits= types. You cannot use it for allocating vectors containing mutable, abstract, or other pointer-backed objects. 
++ As mentioned previously, *Do not allow any array which was initialized inside a* =@no_escape= *block to escape the block.* Doing so will cause incorrect results.
 + You can use =alloc= outside of an =@no_escape= block, but that will leak memory from the buffer and cause it to overflow if you do it too many times.
   If you accidentally do this, and need to reset the buffer, use =Bumper.reset_buffer!=.
 + =alloc(T, n...)= creates a =StrideArraysCore.PtrArray{T, length(n)}=.
-+ In order to be lightweight, Bumper.jl only depends on StrideArraysCore.jl, not the full [[https://github.com/JuliaSIMD/StrideArrays.jl][StrideArrays.jl]], so if you need some of
-  the more advanced functionality from StrideArrays.jl itself, you'll need to do =using StrideArrays= separately.
-+ Bumper.jl is experimental, and may have bugs. Let me know if you find any. Contributing to the test suite would be greatly appreciated.
++ In order to be lightweight, Bumper.jl only depends on [[https://github.com/JuliaSIMD/StrideArraysCore.jl][StrideArraysCore.jl]], not the full [[https://github.com/JuliaSIMD/StrideArrays.jl][StrideArrays.jl]], so if you need some of the more advanced functionality from StrideArrays.jl itself, you'll need to do =using StrideArrays= separately.
++ You are not allowed to use =return= or =@goto= inside a =@no_escape= block, since this could compromise the cleanup it performs after the block finishes.
++ If you use Bumper.jl, please consider submitting a sample of your use-case so I can include it in the test suite.
++ Bumper.jl is experimental, and may have bugs. Let me know if you find any.
 
 ** Concurrency and parallelism
 

src/Bumper.jl

@@ -3,7 +3,7 @@ module Bumper
 export AllocBuffer, alloc, alloc_nothrow, default_buffer, @no_escape, with_buffer
 
 
-## Public
+## Public API
 # ------------------------------------------------------
 mutable struct AllocBuffer{Storage}
     buf::Storage
@@ -18,14 +18,15 @@ function with_buffer end
 function set_default_buffer_size! end
 allow_ptr_array_to_escape() = false
 function alloc_nothrow end
+function reset_buffer! end
 
 ## Private
 # ------------------------------------------------------
 module Internals
 
 using StrideArraysCore, MacroTools
-import Bumper: AllocBuffer,  alloc, default_buffer, allow_ptr_array_to_escape, set_default_buffer_size!, with_buffer, no_escape, @no_escape,
-    alloc_nothrow
+import Bumper: AllocBuffer,  alloc, default_buffer, allow_ptr_array_to_escape, set_default_buffer_size!,
+    with_buffer, no_escape, @no_escape, alloc_nothrow, reset_buffer!
 
 function total_physical_memory()
     @static if isdefined(Sys, :total_physical_memory)
@@ -107,14 +108,19 @@ function _no_escape_macro(b_ex, ex, __module__)
     @gensym b offset
     e_offset = esc(offset)
     e_b = esc(b)
-    cleaned_ex = MacroTools.postwalk(ex) do x
+    # We need to macroexpand the code in case the user has a macro in the body which has return or goto in it
+    MacroTools.postwalk(macroexpand(__module__, ex)) do x
         @gensym rv
-        MacroTools.isexpr(x, :return) ? Expr(:block, :($rv = $(x.args[1])), :($b.offset = $offset), :(return $rv)) : x
+        if MacroTools.isexpr(x, :return) 
+            error("The `return` keyword is not allowed to be used inside the `@no_escape` macro")
+        elseif MacroTools.isexpr(x, :symbolicgoto) # && (x.args[1] ∉ enclosed_labels)
+            error("`@goto` statements are not allowed to be used inside the `@no_escape` macro")
+        end
     end
     quote
         $e_b = $(esc(b_ex))
         $e_offset = getfield($e_b, :offset)
-        res = $(esc(cleaned_ex))
+        res = $(esc(ex))
         $e_b.offset = $e_offset
         if res isa PtrArray && !(allow_ptr_array_to_escape())
            esc_err()

test/runtests.jl

@@ -18,12 +18,6 @@ function f(x, buf::AllocBuffer)
     end
 end
 
-function returns_early()
-    @no_escape begin
-        return sum(alloc(Int, 10) .= 1)
-    end
-end
-
 @testset "basic" begin
     v = [1,2,3]
     b = AllocBuffer(100)
@@ -35,9 +29,51 @@ end
     @test @allocated(f(v, b)) == 0
 
     @test b.offset == 0
+
+    @test_throws Exception alloc(Int, b, 100000)
+    Bumper.reset_buffer!(b)
+    @test_throws Exception @no_escape begin
+        alloc(Int, 10)
+    end
+end
 
-    @test returns_early() == 10
-    @test default_buffer().offset == 0
+macro sneaky_return(ex)
+    esc(:(return $ex))
+end
+
+macro sneaky_goto(label)
+    esc(:(@goto $label))
+end
+
+@testset "trying to break out of no_escape blocks" begin
+    # It is very tricky to properly deal with code which uses @goto or return inside
+    # a @no_escape code block, because could bypass the mechanism for resetting the
+    # buffer's offset after the block completes.
+
+    # I played with some mechanisms for cleaning it up, but they were sometimes incorrect
+    # if one nested multuple @no_escape blocks, so I decided that they should simply be
+    # disabled, and throw an error at macroexpansion time.
+
+    @test_throws Exception Bumper.Internals._no_escape_macro(
+        :(default_buffer()),
+        :(return sum(alloc(Int, 10) .= 1)),
+        @__MODULE__()
+    )
+    @test_throws Exception Bumper.Internals._no_escape_macro(
+        :(default_buffer()),
+        :(@sneaky_return sum(alloc(Int, 10) .= 1)),
+        @__MODULE__()
+    )
+    @test_throws Exception Bumper.Internals._no_escape_macro(
+        :(default_buffer()),
+        :(@goto lab),
+        @__MODULE__()
+    )
+    @test_throws Exception Bumper.Internals._no_escape_macro(
+        :(default_buffer()),
+        :(@sneaky_goto lab),
+        @__MODULE__()
+    )
 end
 
 @testset "tasks and buffer switching" begin