How to implement custom assignment with `:=` ? (for tinkering/learning purposes)

I stumbled across this post and realized I have no idea how I would actually implement := to mean something like “define/assign the variable if it is not already defined, otherwise throw an error”, so I’m looking for some (any) pointers to get started.

This could also just be pointing me at the right files in the julia source or docs, or commenting on my current understanding of what’s going on under the hood so far (see below). Any hints are appreciated :slight_smile:

What I think is going on so far

My first thoughts were: It should be possible to write this as a macro like so

@newsyntax a := 1

Which could take a := 1 and transform it into code that checks the definition of a and then does a normal assignment or throws the error. But that’s clunky of course!

My next thought was that I probably want to compare what the expressions a = 1 and a := 1 get lowered to and it seems like this is the stage that one would need to customize.

julia> Meta.@lower a = 1
:($(Expr(:thunk, CodeInfo(
    @ none within `top-level scope`
1 ─ %1 = Core.get_binding_type(Main, :a)
│        #s1 = 1
│   %3 = #s1 isa %1
└──      goto #3 if not %3
2 ─      goto #4
3 ─      #s1 = Base.convert(%1, #s1)
4 ┄      a = #s1
└──      return 1
))))
julia> Meta.@lower a := 1
:($(Expr(:error, "unsupported assignment operator \":=\"")))

Aha! So somewhere between the expression and the lowered IR Julia figures out that := is not valid syntax (yet!).

Since Meta.lower seems to just do ccall(:jl_expand, ...) the interesting part happens in a C funtion, which I could customize – is this correct so far?

So the next steps would be

  • find the function jl_expand in the Julia source and edit it accordingly
  • compile Julia from the changed source files

Or is there a more straightforward way that allows this kind of rewriting from within Julia?

Disclaimer:
I don’t care if this becomes Julia syntax or not and the precise meaning of := is also not important for me – I just want to learn how to figure out and change the current behavior of Julia internals.

1 Like

The parser has changed, and something like here to look at (assuming you don’t want to change the deprecated flisp parser too):

1 Like

Oh I see, thanks for the pointer!

Making changes in JuliaSyntax.jl would definitely have less overhead than digging into the previous parser).

It’d be clunky if you had to invoke it for every :=, but if you wrote a function that walked entire expression trees, then it becomes much easier. You could then pass that function to include so you’d only need to import that function to a package top level.

Extremely barebones MWE:

In foo.jl:

x := 1
y = 2
y := 3

println(x)
println(y)

In repl:

julia> function colon_def(expr)
           if expr.head === :(:=)
               expr.head = :(=) # or whatever
               return expr
           else
               return expr
           end
       end
colon_def (generic function with 1 method)

julia> include(colon_def, "./foo.jl")
1
3
2 Likes

include_string can work on inlined text, and that could be done in a string literal, a syntax often used for evaluating code in other languages’ interpreters.

julia> macro tweakedparse_str(code::String)
         :(include_string(colon_def, @__MODULE__, $code)) # iffy on this getting to the right module
       end
@tweakedparse_str (macro with 1 method)

julia> tweakedparse"""
       x := 1
       y = 2
       y := 3

       println(x)
       println(y)
       """
1
3

Note that non-standard string literals are not parsed with $-interpolation and \-unescaping like String literals are. @raw_str is a macro that actually does nothing to the input String, so non-standard string literal inputs can be thought of as raw strings. I’m not aware of a way to opt into the String literal parsing behavior inside a macro, so if you need it, the wordier include_string call can be used directly.

I think a macro should work for inline code too, just needs a manual way to work on each subexpression of a begin block:

@tweakedparse2 begin
#= lines of code =#
end

Sounds like a job for MacroTools.jl’s @capture and postwalk:

julia> using MacroTools: @capture, postwalk

julia> macro tweakedparse2(ex)
           return postwalk(ex) do x
               if @capture(x, a_ := b_)
                  return quote 
                       if @isdefined($a) 
                           error("Variable " * $(sprint(show_unquoted, a)) * " is already defined")
                       end                       
                       $a = $b
                   end 
               else 
                   return x
               end
           end
       end
@tweakedparse2 (macro with 1 method)

julia> @tweakedparse2 begin
           x := 1
           y = 2
           z := 3

           @show x y z x + y + z
       end
x = 1
y = 2
z = 3
x + y + z = 6
6

julia> @tweakedparse2 begin
           x := 1  # Fine
           x = 2   # Fine
           x := 3  # Error
       end
ERROR: Variable x is already defined
Stacktrace:
 [1] error(s::String)
   @ Base .\error.jl:35
 [2] top-level scope
   @ REPL[4]:4

Edit: added variable name to error message (cf. @show macro implementation) and split tweakedparse2 macro into multiple lines for easier readability.

3 Likes

Thanks for these suggestions! That’s really convenient, I wasn’t aware that one can transform every expression in the included file (or string) in this way.