Merge pull request #363 from willow-ahrens/wma/extrastatichash

make static hash more static
finch-tensor · Dec 31, 2023 · 2c17510 · 2c17510
2 parents c390952 + 4135c93
commit 2c17510
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 11 deletions.
diff --git a/src/lower.jl b/src/lower.jl
@@ -17,17 +17,11 @@ function contain(f, ctx::LowerJulia; bindings = ctx.bindings, kwargs...)
 end
 
 struct StaticHash
-    counts::Dict{Any, Int}
+    counts::Dict{Tuple{Any, DataType}, UInt}
 end
-StaticHash() = StaticHash(Dict{Any, Int}())
+StaticHash() = StaticHash(Dict{Tuple{Any, DataType}, UInt}())
 
-function (h::StaticHash)(x)
-    if haskey(h.counts, x)
-        return h.counts[x]
-    else
-        return (h.counts[x] = UInt(length(h.counts)))
-    end
-end
+(h::StaticHash)(x) = get!(h.counts, (x, typeof(x)), UInt(length(h.counts)))
 
 (ctx::AbstractCompiler)(root) = ctx(root, Stylize(root, ctx)(root))
 (ctx::AbstractCompiler)(root, style) = lower(root, ctx, style)

diff --git a/src/util/limits.jl b/src/util/limits.jl
@@ -124,14 +124,14 @@ julia> 1.0 < minus_eps(1.0)
 false
 """
 struct Limit{T} <: Number
-    val
+    val::T
     sign::Infinitesimal
     Limit{T}(x::T, y) where {T} = new{T}(x, y)
     Limit{T}(x::T, y) where {T<:Limit} = error()
 end
 
 limit(x::T, s) where {T} = Limit{T}(x, s)
-plus_eps(x)::Limit = limit(0, tiny_positive())
+plus_eps(x)::Limit = limit(x, tiny_positive())
 minus_eps(x)::Limit = limit(x, tiny_negative())
 limit(x) = limit(x, tiny_zero())
 limit(x::Limit) = x