|
| 1 | +# This file is a part of Julia. License is MIT: https://julialang.org/license |
| 2 | + |
| 3 | +module ScopedVariables |
| 4 | + |
| 5 | +export ScopedVariable, scoped |
| 6 | + |
| 7 | +mutable struct Scope |
| 8 | + const parent::Union{Nothing, Scope} |
| 9 | +end |
| 10 | + |
| 11 | +current_scope() = current_task().scope::Union{Nothing, Scope} |
| 12 | + |
| 13 | +""" |
| 14 | + ScopedVariable(x) |
| 15 | +
|
| 16 | +Create a container that propagates values across scopes. |
| 17 | +Use [`scoped`](@ref) to create and enter a new scope. |
| 18 | +
|
| 19 | +Values can only be set when entering a new scope, |
| 20 | +and the value referred to will be constant during the |
| 21 | +execution of a scope. |
| 22 | +
|
| 23 | +Dynamic scopes are propagated across tasks. |
| 24 | +
|
| 25 | +# Examples |
| 26 | +```jldoctest |
| 27 | +julia> const svar = ScopedVariable(1); |
| 28 | +
|
| 29 | +julia> svar[] |
| 30 | +1 |
| 31 | +
|
| 32 | +julia> scoped(svar => 2) do |
| 33 | + svar[] |
| 34 | + end |
| 35 | +2 |
| 36 | +``` |
| 37 | +
|
| 38 | +!!! compat "Julia 1.11" |
| 39 | + This method requires at least Julia 1.11. In Julia 1.7+ this |
| 40 | + is available from the package ScopedVariables.jl. |
| 41 | +""" |
| 42 | +mutable struct ScopedVariable{T} |
| 43 | + const values::WeakKeyDict{Scope, T} |
| 44 | + const initial_value::T |
| 45 | + ScopedVariable{T}(initial_value) where {T} = new{T}(WeakKeyDict{Scope, T}(), initial_value) |
| 46 | +end |
| 47 | +ScopedVariable(initial_value::T) where {T} = ScopedVariable{T}(initial_value) |
| 48 | + |
| 49 | +Base.eltype(::Type{ScopedVariable{T}}) where {T} = T |
| 50 | + |
| 51 | +function Base.getindex(var::ScopedVariable{T})::T where T |
| 52 | + scope = current_scope() |
| 53 | + if scope === nothing |
| 54 | + return var.initial_value |
| 55 | + end |
| 56 | + @lock var.values begin |
| 57 | + while scope !== nothing |
| 58 | + if haskey(var.values.ht, scope) |
| 59 | + return var.values.ht[scope] |
| 60 | + end |
| 61 | + scope = scope.parent |
| 62 | + end |
| 63 | + end |
| 64 | + return var.initial_value |
| 65 | +end |
| 66 | + |
| 67 | +function Base.show(io::IO, var::ScopedVariable) |
| 68 | + print(io, ScopedVariable) |
| 69 | + print(io, '{', eltype(var), '}') |
| 70 | + print(io, '(') |
| 71 | + show(io, var[]) |
| 72 | + print(io, ')') |
| 73 | +end |
| 74 | + |
| 75 | +function __set_var!(scope::Scope, var::ScopedVariable{T}, val::T) where T |
| 76 | + # PRIVATE API! Wrong usage will break invariants of ScopedVariable. |
| 77 | + if scope === nothing |
| 78 | + error("ScopedVariable: Currently not in scope.") |
| 79 | + end |
| 80 | + @lock var.values begin |
| 81 | + if haskey(var.values.ht, scope) |
| 82 | + error("ScopedVariable: Variable is already set for this scope.") |
| 83 | + end |
| 84 | + var.values[scope] = val |
| 85 | + end |
| 86 | +end |
| 87 | + |
| 88 | +""" |
| 89 | + scoped(f, var::ScopedVariable{T} => val::T) |
| 90 | +
|
| 91 | +Execute `f` in a new scope with `var` set to `val`. |
| 92 | +""" |
| 93 | +function scoped(f, pair::Pair{<:ScopedVariable{T}, T}) where T |
| 94 | + @nospecialize |
| 95 | + ct = Base.current_task() |
| 96 | + current_scope = ct.scope::Union{Nothing, Scope} |
| 97 | + try |
| 98 | + scope = Scope(current_scope) |
| 99 | + __set_var!(scope, pair...) |
| 100 | + ct.scope = scope |
| 101 | + return f() |
| 102 | + finally |
| 103 | + ct.scope = current_scope |
| 104 | + end |
| 105 | +end |
| 106 | + |
| 107 | +""" |
| 108 | + scoped(f, vars...::ScopedVariable{T} => val::T) |
| 109 | +
|
| 110 | +Execute `f` in a new scope with each scoped variable set to the provided `val`. |
| 111 | +""" |
| 112 | +function scoped(f, pairs::Pair{<:ScopedVariable}...) |
| 113 | + @nospecialize |
| 114 | + ct = Base.current_task() |
| 115 | + current_scope = ct.scope::Union{Nothing, Scope} |
| 116 | + try |
| 117 | + scope = Scope(current_scope) |
| 118 | + for (var, val) in pairs |
| 119 | + __set_var!(scope, var, val) |
| 120 | + end |
| 121 | + ct.scope = scope |
| 122 | + return f() |
| 123 | + finally |
| 124 | + ct.scope = current_scope |
| 125 | + end |
| 126 | +end |
| 127 | + |
| 128 | +end # module ScopedVariables |
0 commit comments