From 575a73286598b69c4da9fdd91b63b0ef5480d83d Mon Sep 17 00:00:00 2001
From: Shuhei Kadowaki <40514306+aviatesk@users.noreply.github.com>
Date: Wed, 17 Jan 2024 19:21:15 +0900
Subject: [PATCH] use `IdDict{MethodInstance}` instead of
 `Dict{MethodInstance}` (#534)

---
 src/Cthulhu.jl     | 2 --
 src/interface.jl   | 6 ++++--
 src/interpreter.jl | 9 +++++----
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/src/Cthulhu.jl b/src/Cthulhu.jl
index a61d6ef9..46b98bdc 100644
--- a/src/Cthulhu.jl
+++ b/src/Cthulhu.jl
@@ -266,8 +266,6 @@ end
 get_effects(codeinst::CodeInstance) = CC.decode_effects(codeinst.ipo_purity_bits)
 get_effects(codeinst::CodeInfo) = CC.decode_effects(codeinst.purity)
 get_effects(src::InferredSource) = src.effects
-get_effects(unopt::Dict{Union{MethodInstance, InferenceResult}, InferredSource}, mi::MethodInstance) =
-    haskey(unopt, mi) ? get_effects(unopt[mi]) : Effects()
 get_effects(result::InferenceResult) = result.ipo_effects
 get_effects(result::CC.ConstPropResult) = get_effects(result.result)
 get_effects(result::CC.ConcreteResult) = result.effects
diff --git a/src/interface.jl b/src/interface.jl
index ba37f03a..9c9428b5 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -81,8 +81,10 @@ get_excts(interp::CthulhuInterpreter, key::InferenceKey) = get(interp.exception_
 # a sensible default cursor for a MethodInstance
 AbstractCursor(interp::AbstractInterpreter, mi::MethodInstance) = CthulhuCursor(mi)
 
-get_effects(interp::CthulhuInterpreter, mi::MethodInstance, opt::Bool) =
-    get_effects(opt ? interp.opt : interp.unopt, mi)
+function get_effects(interp::CthulhuInterpreter, mi::MethodInstance, opt::Bool)
+    effects = opt ? interp.opt : interp.unopt
+    return haskey(effects, mi) ? get_effects(effects[mi]) : Effects()
+end
 
 mutable struct CustomToggle
     onoff::Bool
diff --git a/src/interpreter.jl b/src/interpreter.jl
index b03e1d4d..620fb619 100644
--- a/src/interpreter.jl
+++ b/src/interpreter.jl
@@ -23,7 +23,8 @@ struct OptimizedSource
 end
 
 const InferenceKey = Union{MethodInstance,InferenceResult}
-const InferenceDict{T} = Dict{InferenceKey, T}
+const InferenceDict{T} = IdDict{InferenceKey, T}
+const OptimizationDict = IdDict{MethodInstance, CodeInstance}
 const PC2Remarks = Vector{Pair{Int, String}}
 const PC2Effects = Dict{Int, Effects}
 const PC2Excts = Dict{Int, Any}
@@ -32,7 +33,7 @@ struct CthulhuInterpreter <: AbstractInterpreter
     native::AbstractInterpreter
 
     unopt::InferenceDict{InferredSource}
-    opt::Dict{MethodInstance, CodeInstance}
+    opt::OptimizationDict
 
     remarks::InferenceDict{PC2Remarks}
     effects::InferenceDict{PC2Effects}
@@ -43,7 +44,7 @@ function CthulhuInterpreter(interp::AbstractInterpreter=NativeInterpreter())
     return CthulhuInterpreter(
         interp,
         InferenceDict{InferredSource}(),
-        Dict{MethodInstance, CodeInstance}(),
+        OptimizationDict(),
         InferenceDict{PC2Remarks}(),
         InferenceDict{PC2Effects}(),
         InferenceDict{PC2Excts}())
@@ -69,7 +70,7 @@ CC.lock_mi_inference(interp::CthulhuInterpreter, mi::MethodInstance) = nothing
 CC.unlock_mi_inference(interp::CthulhuInterpreter, mi::MethodInstance) = nothing
 CC.method_table(interp::CthulhuInterpreter) = method_table(interp.native)
 struct CthulhuCache
-    cache::Dict{MethodInstance, CodeInstance}
+    cache::OptimizationDict
 end
 
 CC.code_cache(interp::CthulhuInterpreter) = WorldView(CthulhuCache(interp.opt), WorldRange(get_world_counter(interp)))