# This file is a part of Julia. License is MIT: https://julialang.org/license

# setup
# -----

include("setup_Compiler.jl")
include("irutils.jl")

using Test

struct InvalidationTesterToken end

struct InvalidationTester <: Compiler.AbstractInterpreter
    world::UInt
    inf_params::Compiler.InferenceParams
    opt_params::Compiler.OptimizationParams
    inf_cache::Vector{Compiler.InferenceResult}
    function InvalidationTester(;
                                world::UInt = Base.get_world_counter(),
                                inf_params::Compiler.InferenceParams = Compiler.InferenceParams(),
                                opt_params::Compiler.OptimizationParams = Compiler.OptimizationParams(),
                                inf_cache::Vector{Compiler.InferenceResult} = Compiler.InferenceResult[])
        return new(world, inf_params, opt_params, inf_cache)
    end
end

Compiler.InferenceParams(interp::InvalidationTester) = interp.inf_params
Compiler.OptimizationParams(interp::InvalidationTester) = interp.opt_params
Compiler.get_inference_world(interp::InvalidationTester) = interp.world
Compiler.get_inference_cache(interp::InvalidationTester) = interp.inf_cache
Compiler.cache_owner(::InvalidationTester) = InvalidationTesterToken()

# basic functionality test
# ------------------------

basic_callee(x) = x
basic_caller(x) = basic_callee(x)

# run inference and check that cache exist
@test Base.return_types((Float64,); interp=InvalidationTester()) do x
    basic_caller(x)
end |> only === Float64

let mi = Base.method_instance(basic_callee, (Float64,))
    ci = mi.cache
    @test !isdefined(ci, :next)
    @test ci.owner === InvalidationTesterToken()
    @test ci.max_world == typemax(UInt)
end

let mi = Base.method_instance(basic_caller, (Float64,))
    ci = mi.cache
    @test !isdefined(ci, :next)
    @test ci.owner === InvalidationTesterToken()
    @test ci.max_world == typemax(UInt)
end

# this redefinition below should invalidate the cache
const BASIC_CALLER_WORLD = Base.get_world_counter()+1
basic_callee(x) = x, x
@test !isdefined(Base.method_instance(basic_callee, (Float64,)), :cache)
let mi = Base.method_instance(basic_caller, (Float64,))
    ci = mi.cache
    @test !isdefined(ci, :next)
    @test ci.owner === InvalidationTesterToken()
    @test ci.max_world == BASIC_CALLER_WORLD
end

# re-run inference and check the result is updated (and new cache exists)
@test Base.return_types((Float64,); interp=InvalidationTester()) do x
    basic_caller(x)
end |> only === Tuple{Float64,Float64}
let mi = Base.method_instance(basic_callee, (Float64,))
    ci = mi.cache
    @test !isdefined(ci, :next)
    @test ci.owner === InvalidationTesterToken()
    @test ci.max_world == typemax(UInt)
end

let mi = Base.method_instance(basic_caller, (Float64,))
    ci = mi.cache
    @test isdefined(ci, :next)
    @test ci.owner === InvalidationTesterToken()
    @test ci.max_world == typemax(UInt)
    ci = ci.next
    @test !isdefined(ci, :next)
    @test ci.owner === InvalidationTesterToken()
    @test ci.max_world != typemax(UInt)
end


# backedge optimization
# ---------------------

const GLOBAL_BUFFER = IOBuffer()

# test backedge optimization when the callee's type and effects information are maximized
begin
    take!(GLOBAL_BUFFER)

    pr48932_callee(x) = (print(GLOBAL_BUFFER, x); Base.inferencebarrier(x))
    pr48932_caller(x) = pr48932_callee(Base.inferencebarrier(x))

    # assert that type and effects information inferred from `pr48932_callee(::Any)` are the top
    let rt = only(Base.return_types(pr48932_callee, (Any,)))
        @test rt === Any
        effects = Base.infer_effects(pr48932_callee, (Any,))
        @test effects == Compiler.Effects()
    end

    # run inference on both `pr48932_caller` and `pr48932_callee`
    let (src, rt) = code_typed((Int,); interp=InvalidationTester()) do x
            @inline pr48932_caller(x)
        end |> only
        @test rt === Any
        @test any(iscall((src, pr48932_callee)), src.code)
    end

    let mi = only(Base.specializations(Base.only(Base.methods(pr48932_callee))))
        # Base.method_instance(pr48932_callee, (Any,))
        ci = mi.cache
        @test isdefined(ci, :next)
        @test ci.owner === nothing
        @test ci.max_world == typemax(UInt)

        # In cache due to Base.return_types(pr48932_callee, (Any,))
        ci = ci.next
        @test !isdefined(ci, :next)
        @test ci.owner === InvalidationTesterToken()
        @test ci.max_world == typemax(UInt)
    end
    let mi = Base.method_instance(pr48932_caller, (Int,))
        ci = mi.cache
        @test !isdefined(ci, :next)
        @test ci.owner === InvalidationTesterToken()
        @test ci.max_world == typemax(UInt)
    end

    @test 42 == pr48932_caller(42)
    @test "42" == String(take!(GLOBAL_BUFFER))

    # test that we didn't add the backedge from `pr48932_callee` to `pr48932_caller`:
    # this redefinition below should invalidate the cache of `pr48932_callee` but not that of `pr48932_caller`
    pr48932_callee(x) = (print(GLOBAL_BUFFER, x); nothing)

    @test length(Base.methods(pr48932_callee)) == 1
    @test Base.only(Base.methods(pr48932_callee, Tuple{Any})) === only(Base.methods(pr48932_callee))
    @test isempty(Base.specializations(Base.only(Base.methods(pr48932_callee, Tuple{Any}))))
    let mi = only(Base.specializations(Base.only(Base.methods(pr48932_caller))))
        # Base.method_instance(pr48932_callee, (Any,))
        ci = mi.cache
        @test isdefined(ci, :next)
        @test ci.owner === InvalidationTesterToken()
        @test_broken ci.max_world == typemax(UInt)
        ci = ci.next
        @test !isdefined(ci, :next)
        @test ci.owner === nothing
        @test_broken ci.max_world == typemax(UInt)
    end

    @test isnothing(pr48932_caller(42))
    @test "42" == String(take!(GLOBAL_BUFFER))
end

begin
    deduped_callee(x::Int) = @noinline rand(Int)
    deduped_caller1(x::Int) = @noinline deduped_callee(x)
    deduped_caller2(x::Int) = @noinline deduped_callee(x)

    # run inference on both `deduped_callerx` and `deduped_callee`
    let (src, rt) = code_typed((Int,); interp=InvalidationTester()) do x
            @inline deduped_caller1(x)
            @inline deduped_caller2(x)
        end |> only
        @test rt === Int
        @test any(isinvoke(:deduped_callee), src.code)
    end

    # Verify that adding the backedge again does not actually add a new backedge
    let mi = Base.method_instance(deduped_caller1, (Int,)),
        ci = mi.cache

        callee_mi = Base.method_instance(deduped_callee, (Int,))

        # Inference should have added the callers to the callee's backedges
        @test ci in callee_mi.backedges

        # In practice, inference will never end up calling `store_backedges`
        # twice on the same CodeInstance like this - we only need to check
        # that de-duplication works for a single invocation
        N = length(callee_mi.backedges)
        Core.Compiler.store_backedges(ci, Core.svec(callee_mi, callee_mi))
        N′ = length(callee_mi.backedges)

        # A single `store_backedges` invocation should de-duplicate any of the
        # edges it is adding.
        @test N′ - N == 1
    end
end

# we can avoid adding backedge even if the callee's return type is not the top
# when the return value is not used within the caller
begin take!(GLOBAL_BUFFER)
    pr48932_callee_inferable(x) = (print(GLOBAL_BUFFER, x); Base.inferencebarrier(1)::Int)
    pr48932_caller_unuse(x) = (pr48932_callee_inferable(Base.inferencebarrier(x)); nothing)

    # assert that type and effects information inferred from `pr48932_callee(::Any)` are the top
    let rt = only(Base.return_types(pr48932_callee_inferable, (Any,)))
        @test rt === Int
        effects = Base.infer_effects(pr48932_callee_inferable, (Any,))
        @test effects == Compiler.Effects()
    end

    # run inference on both `pr48932_caller` and `pr48932_callee`:
    # we don't need to add backedge to `pr48932_callee` from `pr48932_caller`
    # since the inference result of `pr48932_callee` is maximized and it's not inlined
    let (src, rt) = code_typed((Int,); interp=InvalidationTester()) do x
            @inline pr48932_caller_unuse(x)
        end |> only
        @test rt === Nothing
        @test any(iscall((src, pr48932_callee_inferable)), src.code)
    end

    let mi = only(Base.specializations(Base.only(Base.methods(pr48932_callee_inferable))))
        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 === InvalidationTesterToken()
        @test ci.max_world == typemax(UInt)
    end
    let mi = Base.method_instance(pr48932_caller_unuse, (Int,))
        ci = mi.cache
        @test !isdefined(ci, :next)
        @test ci.owner === InvalidationTesterToken()
        @test ci.max_world == typemax(UInt)
    end

    @test isnothing(pr48932_caller_unuse(42))
    @test "42" == String(take!(GLOBAL_BUFFER))

    # test that we didn't add the backedge from `pr48932_callee_inferable` to `pr48932_caller_unuse`:
    # this redefinition below should invalidate the cache of `pr48932_callee_inferable` but not that of `pr48932_caller_unuse`
    pr48932_callee_inferable(x) = (print(GLOBAL_BUFFER, "foo"); x)

    @test isempty(Base.specializations(Base.only(Base.methods(pr48932_callee_inferable, Tuple{Any}))))
    let mi = Base.method_instance(pr48932_caller_unuse, (Int,))
        ci = mi.cache
        @test isdefined(ci, :next)
        @test ci.owner === InvalidationTesterToken()
        @test_broken ci.max_world == typemax(UInt)
        ci = ci.next
        @test !isdefined(ci, :next)
        @test ci.owner === nothing
        @test_broken ci.max_world == typemax(UInt)
    end
    @test isnothing(pr48932_caller_unuse(42))
    @test "foo" == String(take!(GLOBAL_BUFFER))
end

# we need to add backedge when the callee is inlined
begin take!(GLOBAL_BUFFER)

    @noinline pr48932_callee_inlined(@nospecialize x) = (print(GLOBAL_BUFFER, x); Base.inferencebarrier(x))
    pr48932_caller_inlined(x) = pr48932_callee_inlined(Base.inferencebarrier(x))

    # assert that type and effects information inferred from `pr48932_callee(::Any)` are the top
    let rt = only(Base.return_types(pr48932_callee_inlined, (Any,)))
        @test rt === Any
        effects = Base.infer_effects(pr48932_callee_inlined, (Any,))
        @test effects == Compiler.Effects()
    end

    # run inference on `pr48932_caller_inlined` and `pr48932_callee_inlined`
    let (src, rt) = code_typed((Int,); interp=InvalidationTester()) do x
            @inline pr48932_caller_inlined(x)
        end |> only
        @test rt === Any
        @test any(isinvoke(:pr48932_callee_inlined), src.code)
    end

    let mi = Base.method_instance(pr48932_callee_inlined, (Int,))
        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 === InvalidationTesterToken()
        @test ci.max_world == typemax(UInt)
    end
    let mi = Base.method_instance(pr48932_caller_inlined, (Int,))
        ci = mi.cache
        @test !isdefined(ci, :next)
        @test ci.owner === InvalidationTesterToken()
        @test ci.max_world == typemax(UInt)
    end

    @test 42 == pr48932_caller_inlined(42)
    @test "42" == String(take!(GLOBAL_BUFFER))

    # test that we added the backedge from `pr48932_callee_inlined` to `pr48932_caller_inlined`:
    # this redefinition below should invalidate the cache of `pr48932_callee_inlined` but not that of `pr48932_caller_inlined`
    @noinline pr48932_callee_inlined(@nospecialize x) = (print(GLOBAL_BUFFER, x); nothing)

    @test isempty(Base.specializations(Base.only(Base.methods(pr48932_callee_inlined, Tuple{Any}))))
    let mi = Base.method_instance(pr48932_caller_inlined, (Int,))
        ci = mi.cache
        @test isdefined(ci, :next)
        @test ci.owner === InvalidationTesterToken()
        @test ci.max_world != typemax(UInt)
        ci = ci.next
        @test !isdefined(ci, :next)
        @test ci.owner === nothing
        @test ci.max_world != typemax(UInt)
    end

    @test isnothing(pr48932_caller_inlined(42))
    @test "42" == String(take!(GLOBAL_BUFFER))
end

# Issue #57696
# This test checks for invalidation of recursive backedges. However, unfortunately, the original failure
# manifestation was an unreliable segfault or an assertion failure, so we don't have a more compact test.
@test success(`$(Base.julia_cmd()) -e 'Base.typejoin(x, ::Type) = 0; exit()'`)

# Test drop_all_caches functionality
@testset "drop_all_caches" begin
    # Run in subprocess to avoid disrupting the main test process
    script = """
        # Define test functions
        drop_cache_test_f(x) = x + 1
        drop_cache_test_g(x) = drop_cache_test_f(x) * 2

        # Compile the functions and capture stderr
        drop_cache_test_g(5) == 12 || error("failure")

        println(stderr, "==DROPPING ALL CACHES==")

        # Drop all caches
        Base.drop_all_caches()

        # Functions should still work (but will be recompiled on next call)
        drop_cache_test_g(5) == 12 || error("failure")

        println(stderr, "SUCCESS: drop_all_caches test passed")
        exit(0)
    """

    io = Pipe()
    # Run the test in a subprocess because Base.drop_all_caches() is extreme
    result = run(pipeline(`$(Base.julia_cmd()[1]) --startup-file=no --trace-compile=stderr -e "$script"`, stderr=io))
    close(io.in)
    err = read(io, String)
    # println(err)
    @test success(result)
    err_before, err_after = split(err, "==DROPPING ALL CACHES==")
    @test occursin("SUCCESS: drop_all_caches test passed", err_after)
    @test occursin("precompile(Tuple{typeof(Main.drop_cache_test_g), $Int})", err_before)
    @test occursin("precompile(Tuple{typeof(Main.drop_cache_test_g), $Int}) # recompile", err_after)
end