From 6ca37dd99e66c42648ae0befe9089aa17240c161 Mon Sep 17 00:00:00 2001
From: Shuhei Kadowaki <40514306+aviatesk@users.noreply.github.com>
Date: Wed, 1 Nov 2023 17:49:25 +0900
Subject: [PATCH] use `CC.add_invalidation_callback!` if available (#573)

---
 src/abstractinterpret/typeinfer.jl | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/src/abstractinterpret/typeinfer.jl b/src/abstractinterpret/typeinfer.jl
index 5e2494c8e..2abe160ff 100644
--- 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
 # -----