Skip to content

Commit a515f5a

Browse files
committed
Add SmallTag type
This is an alternative to `Tag` that provides largely the same functionality, but carries around only the hash of the function / array types instead of the full types themselves. This can make these types much less bulky to print and easier to visually scan for.
1 parent f00147e commit a515f5a

File tree

1 file changed

+46
-2
lines changed

1 file changed

+46
-2
lines changed

src/config.jl

Lines changed: 46 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,46 @@ end
2020

2121
Tag(::Nothing, ::Type{V}) where {V} = nothing
2222

23-
2423
@inline function (::Type{Tag{F1,V1}}, ::Type{Tag{F2,V2}}) where {F1,V1,F2,V2}
2524
tagcount(Tag{F1,V1}) < tagcount(Tag{F2,V2})
2625
end
2726

27+
# SmallTag is similar to a Tag, but carries just a small UInt64 hash, instead
28+
# of the full type, which makes stacktraces / types easier to read while still
29+
# providing good resilience to perturbation confusion.
30+
struct SmallTag{H}
31+
end
32+
33+
@generated function tagcount(::Type{SmallTag{H}}) where {H}
34+
:($(Threads.atomic_add!(TAGCOUNT, UInt(1))))
35+
end
36+
37+
function SmallTag(f::F, ::Type{V}) where {F,V}
38+
H = if F <: Tuple
39+
# no easy way to check Jacobian tag used with Hessians as multiple functions may be used
40+
# see checktag(::Type{Tag{FT,VT}}, f::F, x::AbstractArray{V}) where {FT<:Tuple,VT,F,V}
41+
nothing
42+
else
43+
hash(F) hash(V)
44+
end
45+
tagcount(SmallTag{H}) # trigger generated function
46+
SmallTag{H}()
47+
end
48+
49+
SmallTag(::Nothing, ::Type{V}) where {V} = nothing
50+
51+
@inline function (::Type{SmallTag{H1}}, ::Type{Tag{F2,V2}}) where {H1,F2,V2}
52+
tagcount(SmallTag{H1}) < tagcount(Tag{F2,V2})
53+
end
54+
55+
@inline function (::Type{Tag{F1,V1}}, ::Type{SmallTag{H2}}) where {F1,V1,H2}
56+
tagcount(Tag{F1,V1}) < tagcount(SmallTag{H2})
57+
end
58+
59+
@inline function (::Type{SmallTag{H1}}, ::Type{SmallTag{H2}}) where {H1,H2}
60+
tagcount(SmallTag{H1}) < tagcount(SmallTag{H2})
61+
end
62+
2863
struct InvalidTagException{E,O} <: Exception
2964
end
3065

@@ -36,13 +71,22 @@ checktag(::Type{Tag{FT,VT}}, f::F, x::AbstractArray{V}) where {FT,VT,F,V} =
3671

3772
checktag(::Type{Tag{F,V}}, f::F, x::AbstractArray{V}) where {F,V} = true
3873

74+
# SmallTag is a smaller tag, that only confirms the hash
75+
function checktag(::Type{SmallTag{HT}}, f::F, x::AbstractArray{V}) where {HT,F,V}
76+
H = hash(F) hash(V)
77+
if HT == H || HT === nothing
78+
true
79+
else
80+
throw(InvalidTagException{SmallTag{H},SmallTag{HT}}())
81+
end
82+
end
83+
3984
# no easy way to check Jacobian tag used with Hessians as multiple functions may be used
4085
checktag(::Type{Tag{FT,VT}}, f::F, x::AbstractArray{V}) where {FT<:Tuple,VT,F,V} = true
4186

4287
# custom tag: you're on your own.
4388
checktag(z, f, x) = true
4489

45-
4690
##################
4791
# AbstractConfig #
4892
##################

0 commit comments

Comments
 (0)