Package for Rust-like borrow checker in Julia?

I think this is a really subtle design question that needs to look at the tradeoff between convenience and safety.

So, generally speaking, you can always @take an owned variable to unpack it, and ship that off to the external method. This trips the .moved = true so you can’t use the variable after that point. You would have to create a new owned variable, like

@own y = maybe_mutating_function!(@take x)

I think this design might be ok for many things. It’s quite similar to a lot of Rust code, when a variable’s ownership is lost. It also lets you gradually increase the number of functions with borrow checking, without forcing you to update your whole codebase at once.

Now, I’ve also written some of the methods needed for collections and math:

We can probably add some other basic stuff. The main thing is to know when it’s okay to have :write versus :read in accesses as that determines what variable states are allowed when using them.

All other methods not provided here, the user would need to specify explicitly. I thought about having a special BorrowedArray{T} <: AbstractArray{T} type for arrays, and another for Number, but I realised it is problematic to do this sort of thing implicitly.

In some sense, we actually want the MethodError, so that the user is forced to either @take the variable, or explicitly declare functions as being borrow-compatible.

One super tricky method is iterate. I haven’t defined that yet… not sure how. Should the default iterate(x) return references or owned values?

But if we return references, we’ll have to enforce that the user only uses foreach(x) do ... end so that we can constraint the lifetimes of the references.

If we return owned values, then the user can’t use the array after iterating over it.

(Speaking of which I should probably add a @clone that performs a move-free @take followed by deepcopy.)

If I’m understanding Rust correctly (I don’t in general), that’s up to the variable declaration if there is one, not the object x.next() returns.

Yeah in Rust it is either:

let mut x = vec![1, 2, 3];

for xi in x {
    // xi is i32
    // Can't access x[0] here
}
// x no longer available!

/// OR

let mut x = vec![1, 2, 3];
for xi in &x {
    // xi is &i32
    // Can also use println!("{}", x[0]); - another immutable reference
}
// x available again!

/// OR

for xi in &mut x {
    // xi is &mut i32
    // Can change xi if we want
    // However, we CANT access x[0] here
}

So I guess my question is, do we want iterate(x) to generate Borrowed{T} of elements, or Owned{T} of elements? As you can see there are various tradeoffs here.

Maybe Owned{T} is best here by default, since you could get the 2nd format above with a @lifetime block again:

@own x = [MyMutable(1), MyMutable(2)]

@lifetime lt begin
    @own const rx = x in lt
    for xi in rx
        xi::Borrowed{T}
        # CANT edit
        # but we CAN get x[1] here (another Borrowed{T})
    end
end
# x is still available

@lifetime lt begin
    @own rx = x in lt
    for xi in rx
        xi::BorrowedMut{T}
        # CAN edit xi
        # but we CANT access x[1]
    end
end
# x is still available

for xi in x
    xi::OwnedMut{T}
    # Can edit xi
end
# However, now x is no longer available! It was moved!

What do you think?

@Benny Ok I got the above idea working for Borrowed (i.e., @ref const rx = x in lt). iterating over Owned* or BorrowedMut is going to need a redesign though!

This now works:

@own x = [1, 2, 3]
@lifetime lt begin
    @ref const ref = x in lt
    for (i, xi) in enumerate(ref)
        @test xi isa Borrowed{Int}
        @test xi == x[i]
    end
end

It would be nice to something like the following syntax; but sadly it it doesn’t work.

julia> ex = quote
           for const xi in x
ERROR: ParseError:
#= Error @ REPL[1]:2:9
ex = quote
    for const xi in x
#       └───────────┘ ── expected assignment after `const` =#

julia> ex = quote
           for const xi = x
           end
       end
ERROR: ParseError:
#= Error @ REPL[2]:2:21
ex = quote
    for const xi = x
#                   └ ── invalid iteration spec: expected one of `=` `in` or `∈` =#

P.S., after staring at @own const x = [1, 2, 3] for a while I’m not sure I like it. Any ideas for alternative syntaxes?

Random ideas:

# immutable                  # mutable
@own const x = [1, 2, 3];    @own y = [1, 2, 3];

# Separate macros
@own x = [1, 2, 3];          @own_mut y = [1, 2, 3];
@let x = [1, 2, 3];          @let_mut y = [1, 2, 3];  # actual rust syntax; but bad Julia intuition
@const x = [1, 2, 3];        @mut y = [1, 2, 3];
@freeze x = [1, 2, 3];       @melt y = [1, 2, 3];

# Annotation
@own x = [1, 2, 3];          @own y::mut = [1, 2, 3];
@own x = [1, 2, 3];          @own y::(@mut) = [1, 2, 3];

# Macro option
@own x = [1, 2, 3];          @own :mut y = [1, 2, 3];

# Chained macros
@own x = [1, 2, 3];          @own @mut y = [1, 2, 3];

Then, what word? @own, @track, @safe, @let? Maybe @bind?

I do think it’s better to have the immutable syntax be easier, as a way to encourage its more frequent use.

I mean, the core challenge is that you’re trying to track references and borrows thereof in a language that simply doesn’t use any notion of first-class references. So you have to introduce the references — those wrappers — for everything, and then you have to change what assignment means, and then you need to think about what function calls mean.

In general, trying to forward methods through wrappers is quite fraught, especially when you start trying to put those wrappers inside structs, containers, and other data structures.

I do appreciate your tenacity though. At the risk of suggesting something terribly cursed, has anyone ever tried doing the Ruby monkey-patch thing of redirecting method_missing to do the unwrap-and-forward step with MethodErrors (or somewhere else in the system) in Julia? An optional checker system like this actually seems like a somewhat defensible use of the technique… if it’s possible at all. I’m thinking it’s probably not possible, though.

Edit: Actually, this is very Cassette.jl ish.

4 Likes

I only had time to skim the rest of this so I’m sorry if this was addressed, but while you can’t override assignment, couldn’t you use @borrow_checker to require all instances of Expr(:(=), ...) to include one of your assignment macros?

If we want the same language-wide logical guarantees as rust, then I agree, it is probably impossible. But if we just want to incrementally improve the design of some library, we can work from inside-out, and do a @take whenever one hits the limits of the methods that have started to use it.

Also, not my first rodeo! DispatchDoctor.jl does the analogous thing but for static return typing :smiley: It’s far from comprehensive (see list of caveats at the end, not to mention completely ignoring internal instabilities) but it has caught so so soooo many unintentional type instabilities for me, that it has become an absolutely essential piece of my toolbox.

DispatchDoctor.jl is imperfect, and goes against some of the original design of the language, and does not cover “everything,” but at the same time, it’s really useful! I think the same thing would be true of a pseudo-borrow checker.

Sounds super cool indeed. Especially if BorrowChecker.jl is only meant to be a tool for test/develop time, I feel like monkey patching is far more acceptable. And it could be yet another thing controlled by Preferences.jl

1 Like

Possibly. Right now what I do is just store the symbol in the struct

mutable struct Bound{T}
    const value::T
    @atomic moved::Bool
    @atomic immutable_borrows::Int
    const symbol::Symbol
    const threadid::Int
end

(used to be named Owned; sorry, I’m still fiddling with the API before it hits 0.1.)

This helps make sure you don’t cheat the system and use regular assignment, by storing the original variable symbol into the object itself:

julia> @bind x = 1
Bound{Int64}(1)

julia> y = x
Bound{Int64}(1)

julia> @move z = y
ERROR: Variable `y` holds an object that was reassigned from `x`.

It does something similar with Threads.threadid(), so that you don’t try to pass a bound variable to a closure-based Threads.@spawn or something. You have to explicitly @take, @clone, or pass an immutable @ref there!

Btw @mbauman do you know if there are any new features I should be aware of regarding aliasing detection? I took a look at What's the aliasing story in Julia.

The reason I’m interested in it is because Rust has this feature of allowing multiple mutable references to an object if they point at non-overlapping fields. However, this is only safe because the compiler can prove whether fields are aliased or not.

Meaning this would work –

struct Point
    x::Vector{Int}
    y::Vector{Int}
end

# Mutable binding of a Point object
@bind @mut p = Point([1], [2])

@lifetime lt begin
    # Get references to all fields we'll need
    @ref @mut p_x = p.x in lt
    @ref @mut p_y = p.y in lt
end

These are mutable references (like C pointers) to p.x and p.y. In Rust, this is allowed, because the compiler can prove whether x and y are aliased or not (for any type of data structure). But in Julia I’m not sure if this is possible?

If not then I think we should just ban multiple mutable references (= current behavior), even if they refer to distinct fields. I think it just violates some of the borrow checker semantics. So we’ll be a bit stricter than Rust in that sense, only because we can’t prove things about the memory layout.

Isn’t that “just” more borrow checking at work? Maybe that’s pedantic, but fundamentally it really shouldn’t matter whether your objects come from fields of the same struct or different ones or even different local “variable” references, should it? They’re all just references, and if they’re distinct then it’s ok. The fact that it’s not aliased just “falls out” of the enforcement of the borrow-checking rules, right?

Julia’s alias detection is far less definitive and much more rudimentary.

1 Like

Yeah I think you’re right. I guess in my head I was assuming that objects being passed into @bind might be unsafe due to being created outside of the borrow checker. But if you did create it within the framework, then I think (?) you are correct?

e.g., if I pass in data from the “real world” to the following function:

function foo(x::Array, y::Array)
    @bind @mut z = (x, y)

    @lifetime lt begin
        @ref @mut rx = z[1] in lt
        @ref @mut ry = z[2] in lt

        push!(ry, 2)
    end
end

with

julia> c = [1, 2, 3];

julia> foo(c, c)  # aliased input => breaks multiple mut ref assumption

But if we had tried to start this using BorrowChecker.jl, we’d get an error:

julia> @bind c = [1, 2, 3];

julia> foo(@take(c), @take(c))
ERROR: Cannot use c: value has been moved

So it’s only as safe as the inputs. Maybe that’s the best we can do and I’m being overkill here? Ugh. Maybe the rules need to be configurable too…

(P.S., for monkey patching, I think method_missing could try calling @take on any Bound objects. That’s kinda getting the right behavior - then a user also can’t cheat and do c1 = @take(c); foo(c1, c1))

If it’s statically knowable like literal/constant fields then it should be possible. Base.mightalias works on a lot of AbstractArrays where the information for aliasing is only known at runtime. For example, view(x, 1:2:9) and view(x, 2:2:10) heavily imply they don’t alias, which mightalias detects, but they share the same SubArray type, so the compiler can’t even begin to tell. Even if 1:2:9 and 2:2:10 were constants, the compiler can’t immediately tell that they are indices that don’t alias; maybe x’s type is weird and only has 1 actual element.

I think this could actually be the proper way to pull this off? It looks perfect.

2 Likes

Would the newer compiler plugin interface be preferred these days over using Cassette.jl?

Is there an equivalent of Cassette.overdub in the new interface?

(I don’t know. That question was for someone more knowledgeable than me. :slight_smile:)

@mbauman nice recommendation – it works!! Check this out:

julia> using BorrowChecker

julia> @bind x = 1
Bound{Int64}(1)

julia> f(x::Int) = x^2
f (generic function with 1 method)

julia> result = BorrowChecker.managed() do
           # Enters Cassette.@overdub context
           f(x)
       end
1

julia> x  # Successfully moved
[moved]

It succesfully called @take(x) on our bound variable x despite us not doing anything.

The code was surprisingly concise too:

maybe_take!(x) = x
function maybe_take!(arg::AllBound)
    is_moved(arg) && throw(MovedError(arg.symbol))
    value = unsafe_get_value(arg)
    mark_moved!(arg)
    return value
end
const SKIP_METHODS = (
    Base.getindex, Base.setindex!, Base.getproperty,
    Base.setproperty!, Base.getfield, Base.setfield!
)
function skip_method(f)
    parentmodule(parentmodule(f)) == parentmodule(@__MODULE__) || f in SKIP_METHODS
end

Cassette.@context ManagedCtx

function Cassette.overdub(ctx::ManagedCtx, f, args...)
    if skip_method(f)
        Cassette.fallback(ctx, f, args...)
    else
        Cassette.recurse(ctx, f, map(maybe_take!, args)...)
    end
end
function managed(f)
    Cassette.@overdub(Cassette.disablehooks(ManagedCtx()), f())
end

That’s it!!

All it does is: goes through the arguments of EVERY SINGLE external function call, run maybe_take! on the arguments, and then call the function. This results in all Bound and BoundMut objects being taken, which voids their ownership in their original scope, thus

julia> x == 2
ERROR: Cannot use x: value has been moved

i.e., it is no longer owned, and can’t be used!

The only functions where x won’t be taken are ones we have defined explicitly. Otherwise we assume all external functions take ownership unless specified otherwise.

5 Likes

Of course now the trouble is that Cassette doesn’t work on 1.12, and I’m similarly out of touch on what’s happening in this space. But yes, I suspect this sort of approach will work much better.

Weird, it still works for me on nightly.

Make `import Cassette` work on Julia 1.12 (although the package proba… · JuliaLabs/Cassette.jl@7e798c2 · GitHub – I guess I just haven’t hit a code path that used this removed bit yet??


Edit: nevermind, while the unittest work on 1.12, the REPL attempt fails:

julia> BorrowChecker.managed() do
           foo(x)
       end
ERROR: The function body AST defined by this @generated function is not pure. This likely means it contains a closure, a comprehension or a generator.