|
| 1 | +""" |
| 2 | + $(TYPEDEF) |
| 3 | +
|
| 4 | +Wrapper over an `AbstractDict{SymbolicT, V} where {V}` which disallows keys that are |
| 5 | +indexed array symbolics. Specifically, if `@variables x[1:4]` exists, then `x` can be |
| 6 | +a key but `x[1]` cannot. |
| 7 | +""" |
| 8 | +struct AtomicArrayDict{V, D <: AbstractDict{SymbolicT, V}} <: AbstractDict{SymbolicT, V} |
| 9 | + dict::D |
| 10 | + |
| 11 | + function AtomicArrayDict(dict::AbstractDict{SymbolicT, V}) where {V} |
| 12 | + for k in keys(dict) |
| 13 | + validate_atomic_array_key(k) |
| 14 | + end |
| 15 | + new{V, typeof(dict)}(dict) |
| 16 | + end |
| 17 | +end |
| 18 | + |
| 19 | +AtomicArrayDict{V, D}(dict::AtomicArrayDict{V, D}) where {V, D} = copy(dict) |
| 20 | +AtomicArrayDict{V, D}() where {V, D} = AtomicArrayDict(D()) |
| 21 | +AtomicArrayDict() = AtomicArrayDict(Dict{SymbolicT, SymbolicT}()) |
| 22 | +AtomicArrayDict(args::Pair...) = AtomicArrayDict(Dict(args...)) |
| 23 | +AtomicArrayDict{V, D}(args::Pair...) where {V, D} = AtomicArrayDict(Dict(args...)) |
| 24 | +AtomicArrayDict{V}(args...) where {V} = AtomicArrayDict(Dict{SymbolicT, V}(args...)) |
| 25 | +AtomicArrayDict{V, D}(args...) where {V, D} = AtomicArrayDict(D(args...)) |
| 26 | + |
| 27 | +struct IndexedArrayKeyError <: Exception |
| 28 | + k::SymbolicT |
| 29 | +end |
| 30 | + |
| 31 | +function Base.showerror(io::IO, err::IndexedArrayKeyError) |
| 32 | + print(io, """ |
| 33 | + `AtomicArrayDict` treats symbolic arrays as atomic. It does not allow keys to be \ |
| 34 | + indexed array symbolics. Got key $(err.k). |
| 35 | + """) |
| 36 | +end |
| 37 | + |
| 38 | +function validate_atomic_array_key(k::SymbolicT) |
| 39 | + split_indexed_var(k)[2] && throw(IndexedArrayKeyError(k)) |
| 40 | +end |
| 41 | + |
| 42 | +Base.copy(dd::AtomicArrayDict) = AtomicArrayDict(copy(dd.dict)) |
| 43 | +function Base.empty(dd::AtomicArrayDict, ::Type{K}, ::Type{V}) where {K, V} |
| 44 | + AtomicArrayDict(empty(dd.dict, K, V)) |
| 45 | +end |
| 46 | + |
| 47 | +Base.get(def::Base.Callable, dd::AtomicArrayDict, k) = def() |
| 48 | +Base.get(def::Base.Callable, dd::AtomicArrayDict, k::SymbolicT) = get(def, dd.dict, k) |
| 49 | +function Base.get(f::Base.Callable, dd::AtomicArrayDict, k::Union{Num, Arr, CallAndWrap}) |
| 50 | + return get(f, dd, unwrap(k)) |
| 51 | +end |
| 52 | +Base.get(dd::AtomicArrayDict, k, default) = get(Returns(default), dd, k) |
| 53 | + |
| 54 | +Base.haskey(dd::AtomicArrayDict, k) = haskey(dd.dict, k) |
| 55 | + |
| 56 | +Base.getindex(dd::AtomicArrayDict, k) = dd.dict[k] |
| 57 | + |
| 58 | +function Base.setindex!(dd::AtomicArrayDict, v, k) |
| 59 | + k = unwrap(k) |
| 60 | + validate_atomic_array_key(unwrap(k)) |
| 61 | + setindex!(dd.dict, v, k) |
| 62 | +end |
| 63 | + |
| 64 | +Base.isempty(dd::AtomicArrayDict) = isempty(dd.dict) |
| 65 | +Base.length(dd::AtomicArrayDict) = length(dd.dict) |
| 66 | +Base.iterate(dd::AtomicArrayDict, args...) = Base.iterate(dd.dict, args...) |
| 67 | +Base.sizehint!(dd::AtomicArrayDict, n; kw...) = sizehint!(dd.dict, n; kw...) |
| 68 | +Base.empty!(dd::AtomicArrayDict) = empty!(dd.dict) |
| 69 | + |
| 70 | +Base.delete!(dd::AtomicArrayDict, k) = delete!(dd.dict, k) |
| 71 | + |
| 72 | +""" |
| 73 | + $TYPEDSIGNATURES |
| 74 | +
|
| 75 | +Convert the symbolic mapping `dict` to an `AtomicArrayDict`. If `dict` contains keys which |
| 76 | +are elements of a symbolic array, the returned mappng will have a key for the array, and |
| 77 | +a value which is a symbolic array where entries specified in `dict` are present and `default` |
| 78 | +otherwise. |
| 79 | +""" |
| 80 | +function as_atomic_dict_with_defaults(dict::AbstractDict{SymbolicT, SymbolicT}, default::SymbolicT) |
| 81 | + dd = AtomicArrayDict(empty(dict)) |
| 82 | + indexed_array_vals = empty(dict, SymbolicT, Array{SymbolicT}) |
| 83 | + for (k, v) in dict |
| 84 | + arr, isarr = split_indexed_var(k) |
| 85 | + if isarr |
| 86 | + buffer = get!(() -> fill(default, size(arr)), indexed_array_vals, arr) |
| 87 | + si = get_stable_index(k) |
| 88 | + buffer[si] = v |
| 89 | + else |
| 90 | + dd[k] = v |
| 91 | + end |
| 92 | + end |
| 93 | + for (k, v) in indexed_array_vals |
| 94 | + if all(SU.isconst, v) |
| 95 | + dd[k] = BSImpl.Const{VartypeT}(unwrap_const.(v)) |
| 96 | + else |
| 97 | + dd[k] = BSImpl.Const{VartypeT}(v) |
| 98 | + end |
| 99 | + end |
| 100 | + return dd |
| 101 | +end |
| 102 | + |
| 103 | +""" |
| 104 | + $TYPEDSIGNATURES |
| 105 | +
|
| 106 | +Modify an atomic array mapping `dd` to map `k` to `v`. If `k` is an indexed array symbolic, |
| 107 | +update the array to have value `v` at the corresponding index. If the array is not a key, |
| 108 | +create the key and set all other entries to `default`. |
| 109 | +""" |
| 110 | +function write_possibly_indexed_array!(dd::AtomicArrayDict{SymbolicT}, k::SymbolicT, v::SymbolicT, default::SymbolicT) |
| 111 | + arr, isarr = split_indexed_var(k) |
| 112 | + if isarr |
| 113 | + buffer::Array{SymbolicT} = if haskey(dd, arr) |
| 114 | + collect(dd[arr]) |
| 115 | + else |
| 116 | + fill(default, size(arr)) |
| 117 | + end |
| 118 | + idx = get_stable_index(k) |
| 119 | + buffer[idx] = v |
| 120 | + if all(SU.isconst, buffer) |
| 121 | + dd[arr] = BSImpl.Const{VartypeT}(unwrap_const.(buffer)) |
| 122 | + else |
| 123 | + dd[arr] = BSImpl.Const{VartypeT}(buffer) |
| 124 | + end |
| 125 | + else |
| 126 | + dd[k] = v |
| 127 | + end |
| 128 | + return dd |
| 129 | +end |
| 130 | + |
| 131 | +""" |
| 132 | + $TYPEDSIGNATURES |
| 133 | +
|
| 134 | +Check if `dd` has the key `k`. If `k` is indexed, check if `dd` has the array as a key. |
| 135 | +""" |
| 136 | +function has_possibly_indexed_key(dd::AtomicArrayDict, k::SymbolicT) |
| 137 | + arr, _ = split_indexed_var(k) |
| 138 | + return haskey(dd, arr) |
| 139 | +end |
| 140 | + |
| 141 | +""" |
| 142 | + $TYPEDSIGNATURES |
| 143 | +
|
| 144 | +Equivalent to `get(dd, k, default)`. If `k` is an indexed array, then return |
| 145 | +`dd[arr][idxs...]` for the corresponding array `arr` and indices, or `default` |
| 146 | +if `arr` does not exist. |
| 147 | +""" |
| 148 | +function get_possibly_indexed(dd::AtomicArrayDict, k::SymbolicT, default) |
| 149 | + arr, isarr = split_indexed_var(k) |
| 150 | + res = get(dd, arr, default) |
| 151 | + isarr || return res |
| 152 | + res === default && return default |
| 153 | + idx = get_stable_index(k) |
| 154 | + return res[idx] |
| 155 | +end |
| 156 | + |
| 157 | +struct AtomicArraySet{D <: AbstractDict{SymbolicT, Nothing}} <: AbstractSet{SymbolicT} |
| 158 | + dd::AtomicArrayDict{Nothing, D} |
| 159 | + |
| 160 | + function AtomicArraySet{D}(dd::AtomicArrayDict{Nothing, D}) where {D} |
| 161 | + new{D}(dd) |
| 162 | + end |
| 163 | +end |
| 164 | + |
| 165 | +AtomicArraySet() = AtomicArraySet{Dict{SymbolicT, Nothing}}() |
| 166 | +AtomicArraySet{D}() where {D} = AtomicArraySet{D}(D()) |
| 167 | +AtomicArraySet{D}(x::D) where {D} = AtomicArraySet{D}(AtomicArrayDict(x)) |
| 168 | + |
| 169 | +Base.isempty(x::AtomicArraySet) = isempty(x.dd) |
| 170 | +Base.length(x::AtomicArraySet) = length(x.dd) |
| 171 | +Base.sizehint!(x::AtomicArraySet, n::Integer) = (sizehint!(x.dd, n); x) |
| 172 | +Base.in(item, x::AtomicArraySet) = haskey(x.dd, item) |
| 173 | +Base.push!(x::AtomicArraySet, item) = (x.dd[item] = nothing; x) |
| 174 | +Base.delete!(x::AtomicArraySet, item) = (delete!(x.dd, item); x) |
| 175 | +Base.empty(::AtomicArraySet{D}) where {D} = AtomicArraySet{D}() |
| 176 | +Base.copy(x::AtomicArraySet{D}) where {D} = AtomicArraySet{D}(copy(x.dd)) |
| 177 | +Base.iterate(x::AtomicArraySet, args...) = iterate(keys(x.dd), args...) |
| 178 | + |
| 179 | +function Base.filter!(f::F, x::AtomicArraySet) where {F} |
| 180 | + filter!(f ∘ first, x.dd) |
| 181 | + return x |
| 182 | +end |
| 183 | + |
| 184 | +""" |
| 185 | + $TYPEDSIGNATURES |
| 186 | +
|
| 187 | +Add `item` to `x`. If `item` is an indexed array, add the array instead. |
| 188 | +""" |
| 189 | +function push_as_atomic_array!(x::AtomicArraySet, item::SymbolicT) |
| 190 | + push!(x, split_indexed_var(item)[1]) |
| 191 | +end |
| 192 | + |
| 193 | +""" |
| 194 | + $METHODLIST |
| 195 | +
|
| 196 | +Convert an array of possibly scalarized variables into an `AtomicArraySet`. |
| 197 | +""" |
| 198 | +as_atomic_array_set(vars::Vector{SymbolicT}) = as_atomic_array_set(Dict{SymbolicT, Nothing}, vars) |
| 199 | +function as_atomic_array_set(::Type{D}, vars::Vector{SymbolicT}) where {D} |
| 200 | + set = AtomicArraySet{D}() |
| 201 | + for v in vars |
| 202 | + push_as_atomic_array!(set, v) |
| 203 | + end |
| 204 | + return set |
| 205 | +end |
| 206 | + |
| 207 | +function contains_possibly_indexed_element(x::AtomicArraySet, k::SymbolicT) |
| 208 | + has_possibly_indexed_key(x.dd, k) |
| 209 | +end |
0 commit comments