Merge pull request #15 from MilesCranmer/spawn-macro

Create `@spawn` to validate captures to `Threads.@spawn`

Author: MilesCranmer

README.md

@@ -114,6 +114,7 @@ nested task error: Cannot create mutable reference: `data` is already mutably bo
 ```
 
 This is because in BorrowChecker.jl's ownership model, similar to Rust, an owned object follows strict borrowing rules to prevent data races and ensure safety.
+(Though, in practice, you should use `BorrowChecker.@spawn` instead of `Threads.@spawn`, so that it validates captured variables.)
 
 ## Ownership Rules
 
@@ -155,6 +156,8 @@ In essence: You can have many readers (`Borrowed`) **or** one writer (`BorrowedM
 > ```
 >
 > This will validate that any captured variable is an immutable reference.
+> Similarly, you should generally prefer the `BorrowChecker.@spawn` macro instead of
+> `Threads.@spawn` to validate captured variables.
 
 ## API
 
@@ -181,6 +184,7 @@ In essence: You can have many readers (`Borrowed`) **or** one writer (`BorrowedM
 ### Validation
 
 - `@cc closure_expr`: Verifies that closures only capture immutable references.
+- `BorrowChecker.@spawn [options...] expr`: A safety wrapper around `Threads.@spawn` that applies `@cc` to the expression (which is internally put inside a closure).
 
 ### Loops
 
@@ -452,6 +456,22 @@ let
 end
 ```
 
+For threads, you can use the `BorrowChecker.@spawn` macro instead of the standard `Threads.@spawn`.
+This ensures safe captures by automatically applying `@cc` to the closure (which is generated internally by `@spawn`):
+
+```julia
+@own x = 42
+@lifetime lt begin
+    @ref ~lt safe_ref = x
+
+    tasks = [
+        BorrowChecker.@spawn safe_ref + 1
+        for _ in 1:10
+    ]
+    sum(fetch, tasks)
+end
+```
+
 </details>
 
 ### Automated Borrowing with `@bc`

docs/src/api.md

@@ -27,6 +27,7 @@ CurrentModule = BorrowChecker
 
 ```@docs
 @cc
+BorrowChecker.@spawn
 ```
 
 ## Types

src/BorrowChecker.jl

@@ -26,6 +26,7 @@ export @own, @move, @ref, @take, @take!, @lifetime, @clone, @bc, @mut, @cc
 using .PreferencesModule: disable_by_default!
 using .StaticTraitModule: is_static
 using .TypesModule: AsMutable, Lifetime, LazyAccessorOf, is_moved, get_owner, get_symbol, get_immutable_borrows, get_mutable_borrows
+using .MacrosModule: @spawn
 #! format: on
 
 end

src/macros.jl

@@ -553,4 +553,25 @@ _check_capture_allowed(::Type{<:OrLazy{OwnedMut}}) = false
 _check_capture_allowed(::Type{<:OrLazy{BorrowedMut}}) = false
 # COV_EXCL_STOP
 
+"""
+    BorrowChecker.@spawn [options...] expr
+
+`Threads.@spawn` but with [`@cc`](@ref) applied to the expression
+to ensure safe captures.
+"""
+macro spawn(expr, args...)
+    is_borrow_checker_enabled(__module__) ||
+        return esc(:($(Threads).@spawn($expr, $(args...))))
+    return esc(_spawn(expr, args...))
+end
+
+function _spawn(args...)
+    inner_closure = gensym("inner_closure")
+    return quote
+        let $inner_closure = $(@__MODULE__).@cc () -> $(last(args))
+            $(Threads).@spawn($(args[1:(end - 1)]...), $(inner_closure)())
+        end
+    end
+end
+
 end

test/FakeModule/src/FakeModule.jl

@@ -1,6 +1,7 @@
 module FakeModule
 
 using BorrowChecker
+using BorrowChecker: @spawn
 using Test
 
 function test()
@@ -42,6 +43,12 @@ function test()
         @test @bc(f(z)) === f(z)
         @test @bc(f(@mut(z_mut))) === f(z_mut)
     end
+
+    let
+        # This spawn still works because the borrow checker is disabled
+        @own x = 1
+        @test fetch(@spawn x + 1) == 2
+    end
 end
 
 end

test/bc_macro.jl renamed to test/complex_macros.jl

@@ -728,3 +728,55 @@ end
         )
     end
 end
+
+@testitem "Simple @spawn macro" begin
+    using BorrowChecker
+    using BorrowChecker: @spawn
+
+    @test_throws "The closure function captured a variable `x" let
+        @own x = 1
+        fetch(@spawn x + 1)
+    end
+    @test_throws "The closure function captured" let
+        @own :mut x = 1
+        fetch(@spawn x + 1)
+    end
+    @test_throws "The closure function captured" let
+        @own :mut x = 1
+        @lifetime lt begin
+            @ref ~lt :mut borrowed = x
+            fetch(@spawn borrowed + 1)
+        end
+    end
+    let
+        @own :mut x = 1
+        @lifetime lt begin
+            @ref ~lt borrowed = x
+            @test fetch(@spawn borrowed + 1) == 2
+        end
+    end
+end
+
+@testitem "Valid uses of @spawn macro" begin
+    using BorrowChecker
+    using BorrowChecker: @spawn
+
+    let
+        @own x = [1, 2, 3]
+        ch = Channel(1)
+        t = @lifetime lt begin
+            @ref ~lt borrowed = x
+            @test fetch(@spawn borrowed .+ 1) == [2, 3, 4]
+
+            # We want to verify that the borrow expires in this spawn:
+            @spawn (take!(ch); borrowed .+ 1)
+        end
+        put!(ch, 1)
+        err_msg = try
+            fetch(t)
+        catch e
+            sprint(showerror, e)
+        end
+        @test occursin("Cannot use `borrowed`: value's lifetime has expired", err_msg)
+    end
+end

test/runtests.jl

@@ -6,7 +6,7 @@ include("ownership_tests.jl")
 include("reference_tests.jl")
 include("feature_tests.jl")
 include("integration_tests.jl")
-include("bc_macro.jl")
+include("complex_macros.jl")
 
 @testitem "Aqua" begin
     using Aqua