|
| 1 | +# This file is a part of Julia. License is MIT: https://julialang.org/license |
| 2 | + |
| 3 | +module ScopedValues |
| 4 | + |
| 5 | +export ScopedValue, with, @with |
| 6 | + |
| 7 | +""" |
| 8 | + ScopedValue(x) |
| 9 | +
|
| 10 | +Create a container that propagates values across dynamic scopes. |
| 11 | +Use [`with`](@ref) to create and enter a new dynamic scope. |
| 12 | +
|
| 13 | +Values can only be set when entering a new dynamic scope, |
| 14 | +and the value referred to will be constant during the |
| 15 | +execution of a dynamic scope. |
| 16 | +
|
| 17 | +Dynamic scopes are propagated across tasks. |
| 18 | +
|
| 19 | +# Examples |
| 20 | +```jldoctest |
| 21 | +julia> const sval = ScopedValue(1); |
| 22 | +
|
| 23 | +julia> sval[] |
| 24 | +1 |
| 25 | +
|
| 26 | +julia> with(sval => 2) do |
| 27 | + sval[] |
| 28 | + end |
| 29 | +2 |
| 30 | +
|
| 31 | +julia> sval[] |
| 32 | +1 |
| 33 | +``` |
| 34 | +
|
| 35 | +!!! compat "Julia 1.11" |
| 36 | + Scoped values were introduced in Julia 1.11. In Julia 1.8+ a compatible |
| 37 | + implementation is available from the package ScopedValues.jl. |
| 38 | +""" |
| 39 | +mutable struct ScopedValue{T} |
| 40 | + const initial_value::T |
| 41 | +end |
| 42 | + |
| 43 | +Base.eltype(::Type{ScopedValue{T}}) where {T} = T |
| 44 | + |
| 45 | +## |
| 46 | +# Notes on the implementation. |
| 47 | +# We want lookup to be unreasonably fast. |
| 48 | +# - IDDict/Dict are ~10ns |
| 49 | +# - ImmutableDict is faster up to about ~15 entries |
| 50 | +# - ScopedValue are meant to be constant, Immutabilty |
| 51 | +# is thus a boon |
| 52 | +# - If we were to use IDDict/Dict we would need to split |
| 53 | +# the cache portion and the value portion of the hash-table, |
| 54 | +# the value portion is read-only/write-once, but the cache version |
| 55 | +# would need a lock which makes ImmutableDict incredibly attractive. |
| 56 | +# We could also use task-local-storage, but that added about 12ns. |
| 57 | +# - Values are GC'd when scopes become unreachable, one could use |
| 58 | +# a WeakKeyDict to also ensure that values get GC'd when ScopedValues |
| 59 | +# become unreachable. |
| 60 | +# - Scopes are an inline implementation of an ImmutableDict, if we wanted |
| 61 | +# be really fancy we could use a CTrie or HAMT. |
| 62 | + |
| 63 | +mutable struct Scope |
| 64 | + const parent::Union{Nothing, Scope} |
| 65 | + const key::ScopedValue |
| 66 | + const value::Any |
| 67 | + Scope(parent, key::ScopedValue{T}, value::T) where T = new(parent, key, value) |
| 68 | +end |
| 69 | +Scope(parent, key::ScopedValue{T}, value) where T = |
| 70 | + Scope(parent, key, convert(T, value)) |
| 71 | + |
| 72 | +function Scope(scope, pairs::Pair{<:ScopedValue}...) |
| 73 | + for pair in pairs |
| 74 | + scope = Scope(scope, pair...) |
| 75 | + end |
| 76 | + return scope |
| 77 | +end |
| 78 | + |
| 79 | +""" |
| 80 | + current_scope()::Union{Nothing, Scope} |
| 81 | +
|
| 82 | +Return the current dynamic scope. |
| 83 | +""" |
| 84 | +current_scope() = current_task().scope::Union{Nothing, Scope} |
| 85 | + |
| 86 | +function Base.show(io::IO, scope::Scope) |
| 87 | + print(io, Scope, "(") |
| 88 | + seen = Set{ScopedValue}() |
| 89 | + while scope !== nothing |
| 90 | + if scope.key ∉ seen |
| 91 | + if !isempty(seen) |
| 92 | + print(io, ", ") |
| 93 | + end |
| 94 | + print(io, typeof(scope.key), "@") |
| 95 | + show(io, Base.objectid(scope.key)) |
| 96 | + print(io, " => ") |
| 97 | + show(IOContext(io, :typeinfo => eltype(scope.key)), scope.value) |
| 98 | + push!(seen, scope.key) |
| 99 | + end |
| 100 | + scope = scope.parent |
| 101 | + end |
| 102 | + print(io, ")") |
| 103 | +end |
| 104 | + |
| 105 | +function Base.getindex(var::ScopedValue{T})::T where T |
| 106 | + scope = current_scope() |
| 107 | + while scope !== nothing |
| 108 | + if scope.key === var |
| 109 | + return scope.value::T |
| 110 | + end |
| 111 | + scope = scope.parent |
| 112 | + end |
| 113 | + return var.initial_value |
| 114 | +end |
| 115 | + |
| 116 | +function Base.show(io::IO, var::ScopedValue) |
| 117 | + print(io, ScopedValue) |
| 118 | + print(io, '{', eltype(var), '}') |
| 119 | + print(io, '(') |
| 120 | + show(IOContext(io, :typeinfo => eltype(var)), var[]) |
| 121 | + print(io, ')') |
| 122 | +end |
| 123 | + |
| 124 | +""" |
| 125 | + with(f, (var::ScopedValue{T} => val::T)...) |
| 126 | +
|
| 127 | +Execute `f` in a new scope with `var` set to `val`. |
| 128 | +""" |
| 129 | +function with(f, pair::Pair{<:ScopedValue}, rest::Pair{<:ScopedValue}...) |
| 130 | + @nospecialize |
| 131 | + ct = Base.current_task() |
| 132 | + current_scope = ct.scope::Union{Nothing, Scope} |
| 133 | + scope = Scope(current_scope, pair, rest...) |
| 134 | + ct.scope = scope |
| 135 | + try |
| 136 | + return f() |
| 137 | + finally |
| 138 | + ct.scope = current_scope |
| 139 | + end |
| 140 | +end |
| 141 | + |
| 142 | +with(@nospecialize(f)) = f() |
| 143 | + |
| 144 | +""" |
| 145 | + @with vars... expr |
| 146 | +
|
| 147 | +Macro version of `with(f, vars...)` but with `expr` instead of `f` function. |
| 148 | +This is similar to using [`with`](@ref) with a `do` block, but avoids creating |
| 149 | +a closure. |
| 150 | +""" |
| 151 | +macro with(exprs...) |
| 152 | + if length(exprs) > 1 |
| 153 | + ex = last(exprs) |
| 154 | + exprs = exprs[1:end-1] |
| 155 | + elseif length(exprs) == 1 |
| 156 | + ex = only(exprs) |
| 157 | + exprs = () |
| 158 | + else |
| 159 | + error("@with expects at least one argument") |
| 160 | + end |
| 161 | + for expr in exprs |
| 162 | + if expr.head !== :call || first(expr.args) !== :(=>) |
| 163 | + error("@with expects arguments of the form `A => 2` got $expr") |
| 164 | + end |
| 165 | + end |
| 166 | + exprs = map(esc, exprs) |
| 167 | + ct = gensym(:ct) |
| 168 | + current_scope = gensym(:current_scope) |
| 169 | + body = Expr(:tryfinally, esc(ex), :($(ct).scope = $(current_scope))) |
| 170 | + quote |
| 171 | + $(ct) = $(Base.current_task)() |
| 172 | + $(current_scope) = $(ct).scope::$(Union{Nothing, Scope}) |
| 173 | + $(ct).scope = $(Scope)($(current_scope), $(exprs...)) |
| 174 | + $body |
| 175 | + end |
| 176 | +end |
| 177 | + |
| 178 | +end # module ScopedValues |
0 commit comments