JuliaLowering/src/binding_analysis.jl - master - githubmirror/julia

JuliaLowering/src/binding_analysis.jl
374 строки13 KB
Em Chu
[JuliaLowering] Make `SyntaxTree` a standard tree (#62474)
29 июл 2026, 15:34
29 июл 2026, 15:3492cca2d
#-------------------------------------------------------------------------------
# Syntactic-block-local dominance analysis to optimize Box allocations.
# Nearly identical (by design) to flisp's `lambda-optimize-vars!` in julia-syntax.scm.
#
# This pass attempts to prove
# (for the special case of assigned-once variables):
#   1. A variable is always defined at every use / capture
#   2. A variable will not be modified anywhere after any capture
#
# Within inner syntactic blocks, an outer assignment effectively "guards" the
# variable so that uses / captures do not taint the "always-defined" status.
#
# This "guard" behavior is disabled when a variable was introduced outside of
# an enclosing loop, since that would allow condition (2) to be violated. In
# contrast, straight-line captures (i.e. an assignment followed by a capture
# in the same block) is allowed even in loops, since it's unconditional.
#
# In the implementation, any variables in `unused` / `live` are considered
# "always-defined- when-used-or-captured-and-only-modified-once-dynamically".
# These variables may temporarily lose their status when considering uses /
# captures in inner blocks, but this is restored later if dominated by an
# outer assignment.

"""
    analyze_def_and_use!(ctx, ex)

Perform tree-based def-use analysis to find captured variables that are assigned
before any closure captures them and not modified afterward. For such variables,
as an abuse of binding flags we can mark them as `unboxed=true` to avoid
unnecessary `Core.Box` allocations during closure conversion.

This is called on the outermost lambda, and recursively processes nested lambdas.
"""
function analyze_def_and_use!(ctx, ex)
    @stm ex begin
        [K"lambda" _ _ _ body _...] -> begin
            _analyze_nested_lambdas!(ctx, body)
            _analyze_lambda_vars!(ctx, ex)
        end
        [K"toplevel_lambda" _ _ _ body _...] -> begin
            _analyze_nested_lambdas!(ctx, body)
            _analyze_lambda_vars!(ctx, ex)
        end
        [K"generated_lambda" _ _ _ body _...] -> begin
            _analyze_nested_lambdas!(ctx, body)
            _analyze_lambda_vars!(ctx, ex)
        end
    end
end

function _analyze_nested_lambdas!(ctx, ex)
    k = kind(ex)
    if k in KSet"lambda toplevel_lambda generated_lambda"
        analyze_def_and_use!(ctx, ex)
    elseif !is_leaf(ex) && !is_quoted(ex)
        for child in children(ex)
            _analyze_nested_lambdas!(ctx, child)
        end
    end
end

"""
    DefUseState

State for def-use analysis (flisp-compatible tables for tracking variable def and use).

Fields:
- `unused`: candidate variables not yet used (read) in current block
- `live`: variables that have been assigned in current block
- `seen`: all variables we've seen assigned
- `decl`: variables scoped in current scope (via `local` or an argument)
- `decl_outside_loop`: variables scoped in scope outside loop (via `local` or an argument)
- `args`: argument variables (never undefined, special handling in mark_used!)
"""
mutable struct DefUseState
    const lambda_id::ScopeId
    const unused::Set{IdTag}
    const live::Set{IdTag}
    const seen::Set{IdTag}
    decl::Set{IdTag}
    decl_outside_loop::Set{IdTag}
    const args::Set{IdTag}

    function DefUseState(lambda_id, ctx, candidates)
        unused = copy(candidates)
        live = Set{IdTag}()
        seen = Set{IdTag}()
        decl = Set{IdTag}()
        decl_outside_loop = Set{IdTag}()
        args = Set{IdTag}()
        # Initialize decl and args with arguments since they're implicitly declared outside any loop
        for id in candidates
            binfo = get_binding(ctx, id)
            if binfo.kind == :argument
                push!(decl, id)
                push!(args, id)
            end
        end
        return new(lambda_id, unused, live, seen, decl, decl_outside_loop, args)
    end
end

# At CFG merge points, we lose certainty about which path was taken,
# so variables assigned in one branch may not have been assigned.
# Move live variables back to unused to require re-assignment.
# NOTE: This is NOT needed at branch points (return/break/goto) because
# code after them is unreachable - only at merge points (if/while/label).
function du_kill!(state::DefUseState)
    union!(state.unused, state.live)
    empty!(state.live)
end

# Restore live to a previous state, moving new additions back to unused
function du_restore!(state::DefUseState, prev)
    for id in state.decl_outside_loop
        if (id in prev) && !(id in state.unused)
            # This variable was 'used' inside this branch, but it's declared
            # outside of a loop so it may see the dominating assignment execute
            # multiple times. Invalidate it here for soundness.
            delete!(prev, id)
        end
    end
    for id in state.live
        if !(id in prev)
            push!(state.unused, id)
        end
    end
    empty!(state.live)
    union!(state.live, prev)
end

# At the beginning of a loop, move all active decls into the "decl_outside_loop" set.
function du_enter_loop!(state::DefUseState)
    prev_decl_outside_loop = state.decl_outside_loop
    state.decl_outside_loop = state.decl
    state.decl = copy(state.decl)
    return prev_decl_outside_loop
end

# At the end of a loop, restore the previous set of "declared" variables.
function du_leave_loop!(state::DefUseState, prev_decl_outside_loop)
    state.decl = state.decl_outside_loop
    state.decl_outside_loop = prev_decl_outside_loop
end

# When a variable is used (read), remove from unused.
# Note: arguments are only "used" for purposes of this analysis when
# they are captured, since they are never undefined.
function du_mark_used!(state::DefUseState, var_id)
    if var_id in state.unused && !(var_id in state.args)
        delete!(state.unused, var_id)
    end
end

# When a variable is captured by a nested lambda before being assigned
function du_mark_captured!(state::DefUseState, var_id)
    if var_id in state.unused
        delete!(state.unused, var_id)
    end
end

# When a variable is assigned, move from unused to live
function du_assign!(state::DefUseState, var_id)
    if var_id in state.unused
        push!(state.live, var_id)
        push!(state.seen, var_id)
        delete!(state.unused, var_id)
    end
end

# Track local declarations for loop handling
function du_declare!(state::DefUseState, var_id)
    if var_id in state.unused
        push!(state.decl, var_id)
    end
end

# Returns whether e contained a symboliclabel
function du_visit!(ctx, state::DefUseState, e)
    k = kind(e)

    if k == K"BindingId"
        du_mark_used!(state, syntax_id(e))
        return false

    elseif k == K"symboliclabel"
        # Must check BEFORE is_leaf since symboliclabel is a leaf node
        du_kill!(state)
        return true

    elseif k == K"label"
        du_kill!(state)
        return false

    elseif k in KSet"break symbolicgoto oldsymbolicgoto"
        # this kill!() is not required for soundness since these are branch points
        # not merge points, but it's here for parity with flisp
        du_kill!(state)
        return false

    elseif k == K"="
        # Visit RHS first, then record assignment
        has_label = du_visit!(ctx, state, e[2])
        lhs = e[1]
        if kind(lhs) == K"BindingId"
            du_assign!(state, syntax_id(lhs))
        end
        return has_label

    elseif k == K"lambda"
        # Check captures from nested lambda
        for (id, is_capt) in lambda_bindings(e[1]).locals_capt
            if is_capt
                du_mark_captured!(state, id)
            end
        end
        # Don't recurse into nested lambdas - they have their own analysis
        return false

    elseif k == K"local"
        # Track local declarations for loop handling
        # Note: For typed locals like `local x::T`, the K"local" node only
        # contains the BindingId after desugaring. The type info is in
        # a separate K"decl" node. So we only need to handle K"BindingId" here.
        for child in children(e)
            if kind(child) == K"BindingId"
                du_declare!(state, syntax_id(child))
            end
        end
        return false

    elseif k == K"decl"
        # Don't recurse into decl nodes - the BindingId is just a declaration,
        # not a use. We only need to visit the type expression.
        if numchildren(e) >= 2
            return du_visit!(ctx, state, e[2])
        end
        return false

    elseif k == K"function_decl"
        # [function_decl] defines and instantiates the closure type and assigns
        # it to its first argument (but only once per unique closure key).
        @assert kind(e[1]) == K"BindingId"
        func_id = syntax_id(e[1])
        func_id in state.seen && return false
        ck = ClosureKey(func_id, state.lambda_id)
        if haskey(ctx.closure_bindings, ck)
            for lam in ctx.closure_bindings[ck].lambdas
                for (id, capt) in lam.locals_capt
                    capt && du_mark_captured!(state, id)
                end
            end
        end
        du_assign!(state, func_id)
        return false

    elseif k == K"method_defs"
        # Process nested lambdas within
        has_label = false
        for child in children(e)
            has_label |= du_visit!(ctx, state, child)
        end
        return has_label

    elseif k == K"return"
        has_label = numchildren(e) >= 1 ? du_visit!(ctx, state, e[1]) : false
        du_kill!(state) # not necessary, but included for flisp parity
        return has_label

    elseif k in KSet"if elseif trycatchelse tryfinally"
        prev = copy(state.live)
        has_label = false
        for child in children(e)
            has_label |= du_visit!(ctx, state, child)
            du_kill!(state)
        end
        if has_label
            # If there's a label inside, we could have skipped a prior
            # variable initialization
            return true
        else
            du_restore!(state, prev)
            return false
        end

    elseif k in KSet"_while _do_while"
        prev = copy(state.live)
        old_decl = du_enter_loop!(state)
        has_label = false
        for child in children(e)
            has_label |= du_visit!(ctx, state, child)
        end
        du_leave_loop!(state, old_decl)
        if has_label
            du_kill!(state)
            return true
        else
            du_restore!(state, prev)
            return false
        end

    elseif k == K"symbolicblock"
        # Skip the first child (break target label) - it's not a @goto target
        # No save/restore needed: the body always executes (break just exits early)
        has_label = false
        for child in children(e)[2:end]
            has_label |= du_visit!(ctx, state, child)
        end
        return has_label

    elseif is_leaf(e) || is_quoted(e) ||
        k in KSet"local always_defined meta inbounds boundscheck noinline
            loopinfo decl with_static_parameters toplevel_butfirst global
            globalref constdecl atomic isdefined toplevel module error
            gc_preserve_begin gc_preserve_end export public inline"

        # Forms that don't interact with locals or affect control flow (likely more than is necessary).
        # flisp: `lambda-opt-ignored-exprs`
        return false

    else
        has_label = false
        for child in children(e)
            has_label |= du_visit!(ctx, state, child)
        end
        return has_label
    end
end

function _analyze_lambda_vars!(ctx::VariableAnalysisContext, ex)
    # Collect candidate variables: captured and single-assigned
    candidates = Set{IdTag}()
    for (id, from_outer_lambda) in lambda_bindings(ex[1]).locals_capt
        b = get_binding(ctx, id)
        !b.is_captured && continue
        from_outer_lambda && continue
        if b.is_assigned_once && b.kind in (:local, :argument)
            push!(candidates, id)
        end
    end
    isempty(candidates) && return

    state = DefUseState(lambda_bindings(ex[1]).scope_id, ctx, candidates)
    @stm ex begin
        [K"lambda" _ _ _ body] -> du_visit!(ctx, state, body)
        [K"lambda" _ _ _ body rett] -> (du_visit!(ctx, state, body);
                                      du_visit!(ctx, state, rett))
        [K"toplevel_lambda" _ _ _ body] -> du_visit!(ctx, state, body)
        [K"generated_lambda" _ _ _ body] -> du_visit!(ctx, state, body)
    end

    for id in union(state.live, state.unused)
        if id in state.seen
            b = get_binding(ctx, id)
            b.unboxed = true
            b.is_always_defined = true
        end
    end

    # A single (scope-dominating) assignment implies unboxed even if we gave up above
    for id in candidates
        b = get_binding(ctx, id)
        # XXX: This uses is-always-defined to imply that the assignment is defined
        #      everywhere in its scope, which then implies that the one definition
        #      executes only once dynamically.
        #      (i.e. it forbids single-assignment to `x` in an inner loop)
        #
        #      If this flag becomes broader and only considers definedness-at-use
        #      then this check (taken from `julia-syntax.scm`) becomes unsound.
        if b.kind === :local && b.is_always_defined && b.is_assigned_once
            b.unboxed = true
        end
    end
end