use CC.add_invalidation_callback! if available (#573)

aviatesk · Nov 1, 2023 · 6ca37dd · 6ca37dd
1 parent a06c4ad
commit 6ca37dd
Showing 1 changed file with 19 additions and 4 deletions.
diff --git a/src/abstractinterpret/typeinfer.jl b/src/abstractinterpret/typeinfer.jl
@@ -280,6 +280,22 @@ function CC.setindex!(wvc::WorldView{<:AbstractAnalyzerView}, ci::CodeInstance,
     analysis_cache[mi] = ci
 end
 
+struct JETCallback
+    analysis_cache::AnalysisCache
+end
+
+@static if VERSION ≥ v"1.11.0-DEV.798"
+
+function add_jet_callback!(mi::MethodInstance, analysis_cache::AnalysisCache)
+    callback = JETCallback(analysis_cache)
+    CC.add_invalidation_callback!(callback, mi)
+end
+function (callback::JETCallback)(replaced::MethodInstance, max_world::UInt32)
+    delete!(callback.analysis_cache, replaced)
+end
+
+else
+
 function add_jet_callback!(mi::MethodInstance, analysis_cache::AnalysisCache)
     callback = JETCallback(analysis_cache)
     if !isdefined(mi, :callbacks)
@@ -293,10 +309,7 @@ function add_jet_callback!(mi::MethodInstance, analysis_cache::AnalysisCache)
     return nothing
 end
 
-struct JETCallback
-    analysis_cache::AnalysisCache
-end
-function (callback::JETCallback)(replaced::MethodInstance, max_world,
+function (callback::JETCallback)(replaced::MethodInstance, max_world::UInt32,
                                  seen::IdSet{MethodInstance} = IdSet{MethodInstance}())
     push!(seen, replaced)
     delete!(callback.analysis_cache, replaced)
@@ -311,6 +324,8 @@ function (callback::JETCallback)(replaced::MethodInstance, max_world,
     return nothing
 end
 
+end
+
 # local
 # -----