search

c42f / JuliaLowering.jl

Julia code lowering with precise provenance
MIT License
52 stars 2 forks source link

readme

JuliaLowering

Build Status

JuliaLowering.jl is an experimental port of Julia's code lowering compiler passes, written in Julia itself. "Code lowering" is the set of compiler passes which symbolically transform and simplify Julia's syntax prior to type inference.

Goals

This work is intended to

Trying it out

Note this is a very early work in progress; most things probably don't work!

  1. Use a recent dev version of Julia (need at least version 1.12.0-DEV.512)
  2. Check out the main branch of JuliaSyntax
  3. Get the latest version of JuliaSyntaxFormatter
  4. Run the demo include("test/demo.jl")

Design Notes

Lowering has five symbolic simplification passes:

  1. Macro expansion - expanding user-defined syntactic constructs by running the user's macros. This pass also includes a small amount of other symbolic simplification.
  2. Syntax desugaring - simplifying Julia's rich surface syntax down to a small number of syntactic forms.
  3. Scope analysis - analyzing identifier names used in the code to discover local variables, closure captures, and associate global variables to the appropriate module.
  4. Closure conversion - convert closures to types and deal with captured variables efficiently where possible.
  5. Flattening to linear IR - convert code in hierarchical tree form to a flat array of statements; convert control flow into gotos.

Syntax trees

Want something something better than JuliaSyntax.SyntaxNode! SyntaxTree and SyntaxGraph provide this. (These will probably end up in JuliaSyntax.)

We want to allow arbitrary attributes to be attached to tree nodes by analysis passes. This separates the analysis pass implementation from the data structure, allowing passes which don't know about each other to act on a shared data structure.

Design and implementation inspiration comes in several analogies:

Analogy 1: the ECS (Entity-Component-System) pattern for computer game design. This pattern is highly successful because it separates game logic (systems) from game objects (entities) by providing flexible storage

Analogy 2: The AoS to SoA transformation. But here we've got a kind of tree-of-structs-with-optional-attributes to struct-of-Dicts transformation. The data alignment / packing efficiency and concrete type safe storage benefits are similar.

Analogy 3: Graph algorithms which represent graphs as a compact array of node ids and edges with integer indices, rather than using a linked data structure.

Provenance tracking

Expression provenance is tracked through lowering by attaching provenance information in the source attribute to every expression as it is generated. For example when parsing a source file we have

julia> ex = parsestmt(SyntaxTree, "a + b", filename="foo.jl")
SyntaxTree with attributes kind,value,name_val,syntax_flags,source
[call-i]                                │ 
  a                                     │ 
  +                                     │ 
  b                                     │ 

julia> ex[3].source
a + b
#   ╙ ── these are the bytes you're looking for 😊

The provenance function should be used to look up the source attribute and the showprov function used to inspect the content (this is preferred because the encoding of source is an implementation detail). For example:

julia> showprov(ex[3])
a + b
#   ╙ ── in source
# @ foo.jl:1

During macro expansion and lowering provenance gets more complicated because an expression can arise from multiple sources. For example, we want to keep track of the entire stack of macro expansions an expression was generated by, while also recording where it occurred in the original source file.

For this, we use a tree data structure. Let's look at the following pair of macros

julia> JuliaLowering.include_string(Main, raw"""
       module M
           macro inner()
               :(2)
           end

           macro outer()
               :((1, @inner))
           end
       end
       """, "some_macros.jl")

The tree which arises from macro expanding this is pretty simple:

julia> expanded = JuliaLowering.macroexpand(Main, parsestmt(SyntaxTree, "M.@outer()"))
SyntaxTree with attributes scope_layer,kind,value,var_id,name_val,syntax_flags,source
[tuple-p]                               │ 
  1                                     │ 
  2                                     │

but the provenance information recorded for the second element 2 of this tuple is not trivial; it includes the macro call expressions for @inner and @outer. We can show this in tree form:

julia> showprov(expanded[2], tree=true)
2
├─ 2
│  └─ @ some_macros.jl:3
└─ (macrocall @inner)
   ├─ (macrocall @inner)
   │  └─ @ some_macros.jl:7
   └─ (macrocall-p (. M @outer))
      └─ @ foo.jl:1

or as a more human readable flattened list highlighting of source ranges:

module M
    macro inner()
        :(2)
#         ╙ ── in source
    end

# @ some_macros.jl:3

    macro outer()
        :((1, @inner))
#             └────┘ ── in macro expansion
    end
end
# @ some_macros.jl:7

M.@outer()
└────────┘ ── in macro expansion
# @ foo.jl:1

Hygiene

Problems with Hygiene in Julia's exiting macro system

To write correct hygienic macros in Julia (as of 2024), macro authors must use esc() on any any syntax passed to the macro so that passed identifiers escape to the macro caller scope. However

The requirement to use esc() stems from Julia's pervasive use of the simple Expr data structure which represents a unadorned AST in which names are plain symbols. For example, a macro call @foo x gets passed the symbol :x which is just a name without any information attached to indicate that it came from the scope where @foo was called.

Hygiene in JuliaLowering

In JuliaLowering we make hygiene automatic and remove esc() by combining names with scope information. In the language of the paper Towards the Essence of Hygiene by Michael Adams, this combination is called a "syntax object". In JuliaLowering our representation is the tuple (name,scope_layer), also called VarId in the scope resolution pass.

JuliaLowering's macro expander attaches a unique scope layer to each identifier in a piece of syntax. A "scope layer" is an integer identifer combined with the module in which the syntax was created.

When expanding macros,

Subsequently, the (name,scope_layer) pairs are used when resolving bindings. This ensures that, by default, we satisfy the basic rules for hygenic macros discussed in Adams' paper:

  1. A macro can't insert a binding that can capture references other than those inserted by the macro.
  2. A macro can't insert a reference that can be captured by bindings other than those inserted by the macro.

TODO: Write more here...

References

Julia's existing lowering implementation

How does macro expansion work?

macroexpand(m::Module, x) calls jl_macroexpand in ast.c:

jl_value_t *jl_macroexpand(jl_value_t *expr, jl_module_t *inmodule)
{
    expr = jl_copy_ast(expr);
    expr = jl_expand_macros(expr, inmodule, NULL, 0, jl_world_counter, 0);
    expr = jl_call_scm_on_ast("jl-expand-macroscope", expr, inmodule);
    return expr;
}

First we copy the AST here. This is mostly a trivial deep copy of Exprs and shallow copy of their non-Expr children, except for when they contain embedded CodeInfo/phi/phic nodes which are also deep copied.

Second we expand macros recursively by calling

jl_expand_macros(expr, inmodule, macroctx, onelevel, world, throw_load_error)

This relies on state indexed by inmodule and world, which gives it some funny properties:

Expansion proceeds from the outermost to innermost macros. So macros see any macro calls or quasiquote (quote/$) in their children as unexpanded forms.

Things which are expanded:

Scope resolution

Scopes are documented in the Juila documentation on Scope of Variables

This pass disambiguates variables which have the same name in different scopes and fills in the list of local variables within each lambda.

Which data is needed to define a scope?

As scope is a collection of variable names by category:

How does scope resolution work?

We traverse the AST starting at the root paying attention to certian nodes:

scope-block is the complicated bit. It's processed by

Intermediate forms used in lowering

There's also this comment in https://github.com/JuliaLang/julia/issues/22314:

mark the [...] variable as local-def, which would prevent it from getting Core.Boxed during the closure conversion it'll be detected as known-SSA

But maybe that's confusing. It seems like local-def is a local which lowering asserts is "always defined" / "definitely initialized before use". But it's not necessarily single-assign, so not SSA.

Lowered IR

See https://docs.julialang.org/en/v1/devdocs/ast/#Lowered-form

CodeInfo

mutable struct CodeInfo
    code::Vector{Any}             # IR statements
    codelocs::Vector{Int32}       # `length(code)` Vector of indices into `linetable`
    ssavaluetypes::Any            # `length(code)` or Vector of inferred types after opt
    ssaflags::Vector{UInt32}      # flag for every statement in `code`
                                  #   0 if meta statement
                                  #   inbounds_flag - 1 bit (LSB)
                                  #   inline_flag   - 1 bit
                                  #   noinline_flag - 1 bit
                                  #   ... other 8 flags which are defined in compiler/optimize.jl
                                  #   effects_flags - 9 bits
    method_for_inference_limit_heuristics::Any
    linetable::Any
    slotnames::Vector{Symbol}     # names of parameters and local vars used in the code
    slotflags::Vector{UInt8}      # vinfo flags from flisp
    slottypes::Any                # nothing (used by typeinf)
    rettype::Any                  # Any (used by typeinf)
    parent::Any                   # nothing (used by typeinf)
    edges::Any
    min_world::UInt64
    max_world::UInt64
    inferred::Bool
    propagate_inbounds::Bool
    has_fcall::Bool
    nospecializeinfer::Bool
    inlining::UInt8
    constprop::UInt8
    purity::UInt16
    inlining_cost::UInt16
end

Notes on toplevel-only forms and eval-related functions

In the current Julia runtime,

Base.eval()

jl_toplevel_eval_flex(mod, ex)

Should we lower the above blessed top level forms to julia runtime calls? Pros:

In general, we'd be replacing current declarative lowering targets like Expr(:using) with an imperative call to a Core API instead. The call and the setup of its arguments would need to go in a thunk. We've currently got an odd mixture of imperative and declarative lowered code.

Notes on Racket's hygiene

People look at Racket as an example of a very complete system of hygienic macros. We should learn from them, but keeping in mind that Racket's macro system is inherently more complicated. Racket's current approach to hygiene is described in an accessible talk and in more depth in a paper.

Some differences which makes Racket's macro expander different from Julia: