From 5bd972a3a121176d1c1f912d82fbb8ae9667a596 Mon Sep 17 00:00:00 2001
From: Shuhei Kadowaki <aviatesk@gmail.com>
Date: Thu, 15 Feb 2024 20:30:06 +0900
Subject: [PATCH] add more external absint precompile test

---
 test/precompile.jl | 76 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 76 insertions(+)

diff --git a/test/precompile.jl b/test/precompile.jl
index c4cceced8235b1..015cf49f228ac8 100644
--- a/test/precompile.jl
+++ b/test/precompile.jl
@@ -1776,6 +1776,82 @@ let newinterp_path = abspath("compiler/newinterp.jl")
                 @test found
             end
         end
+
+        write(joinpath(load_path, "CustomAbstractInterpreterCaching2.jl"), :(module CustomAbstractInterpreterCaching2
+            import SimpleModule: basic_caller, basic_callee
+
+            module Custom
+                const CC = Core.Compiler
+                include("$($newinterp_path)")
+                @newinterp PrecompileInterpreter
+                struct CustomData
+                    inferred
+                    CustomData(@nospecialize inferred) = new(inferred)
+                end
+                function CC.transform_result_for_cache(interp::PrecompileInterpreter,
+                        mi::Core.MethodInstance, valid_worlds::CC.WorldRange, result::CC.InferenceResult)
+                    inferred_result = @invoke CC.transform_result_for_cache(interp::CC.AbstractInterpreter,
+                        mi::Core.MethodInstance, valid_worlds::CC.WorldRange, result::CC.InferenceResult)
+                    return CustomData(inferred_result)
+                end
+                function CC.inlining_policy(interp::PrecompileInterpreter, @nospecialize(src),
+                                            @nospecialize(info::CC.CallInfo), stmt_flag::UInt32)
+                    if src isa CustomData
+                        src = src.inferred
+                    end
+                    return @invoke CC.inlining_policy(interp::CC.AbstractInterpreter, src::Any,
+                                                    info::CC.CallInfo, stmt_flag::UInt32)
+                end
+            end
+
+            Base.return_types((Float64,)) do x
+                basic_caller(x)
+            end
+            Base.return_types((Float64,); interp=Custom.PrecompileInterpreter()) do x
+                basic_caller(x)
+            end
+            Base.return_types((Vector{Float64},)) do x
+                sum(x)
+            end
+            Base.return_types((Vector{Float64},); interp=Custom.PrecompileInterpreter()) do x
+                sum(x)
+            end
+        end) |> string)
+        Base.compilecache(Base.PkgId("CustomAbstractInterpreterCaching2"))
+        @eval let
+            using CustomAbstractInterpreterCaching2
+            cache_owner = Core.Compiler.cache_owner(
+                CustomAbstractInterpreterCaching2.Custom.PrecompileInterpreter())
+            let m = only(methods(CustomAbstractInterpreterCaching.basic_callee))
+                mi = only(Base.specializations(m))
+                ci = mi.cache
+                @test isdefined(ci, :next)
+                @test ci.owner === nothing
+                @test ci.max_world == typemax(UInt)
+                ci = ci.next
+                @test !isdefined(ci, :next)
+                @test ci.owner === cache_owner
+                @test ci.max_world == typemax(UInt)
+            end
+            let m = only(methods(sum, (Vector{Float64},)))
+                found = false
+                for mi = Base.specializations(m)
+                    if mi isa Core.MethodInstance && mi.specTypes == Tuple{typeof(sum),Vector{Float64}}
+                        ci = mi.cache
+                        @test isdefined(ci, :next)
+                        @test ci.owner === cache_owner
+                        @test ci.max_world == typemax(UInt)
+                        ci = ci.next
+                        @test !isdefined(ci, :next)
+                        @test ci.owner === nothing
+                        @test ci.max_world == typemax(UInt)
+                        found = true
+                        break
+                    end
+                end
+                @test found
+            end
+        end
     end
 end