using Enzyme, LLVM, Test
using FileCheck
import Libdl


@testset "Partial return preservation" begin
    @test @filecheck begin
        # Both stores into the freshly allocated struct must be preserved.
        # Match without spelling out the pointer type so this works under both
        # opaque (`ptr addrspace(10)`) and typed (`{} addrspace(10)*`) pointers.
        @check_label "@inner"
        @check "store atomic"
        @check_same "%v1"
        @check "store atomic"
        @check_same "%v2"
        LLVM.Context() do ctx
            mod = parse(
                LLVM.Module, """
                source_filename = "start"
                target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128-ni:10:11:12:13"
                target triple = "x86_64-linux-gnu"

                declare noalias nonnull {} addrspace(10)* @julia.gc_alloc_obj({}**, i64, {} addrspace(10)*) local_unnamed_addr #5

                define internal fastcc nonnull {} addrspace(10)* @inner({} addrspace(10)* %v1, {} addrspace(10)* %v2) {
                top:
                  %newstruct = call noalias nonnull dereferenceable(16) {} addrspace(10)* @julia.gc_alloc_obj({}** null, i64 16, {} addrspace(10)* addrspacecast ({}* inttoptr (i64 129778359735376 to {}*) to {} addrspace(10)*)) #30
                  %a31 = addrspacecast {} addrspace(10)* %newstruct to {} addrspace(10)* addrspace(11)*
                  %a32 = getelementptr inbounds {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %a31, i64 1
                  store atomic {} addrspace(10)* %v1, {} addrspace(10)* addrspace(11)* %a31 release, align 8
                  %a33 = addrspacecast {} addrspace(10)* %newstruct to i8 addrspace(11)*
                  %a34 = getelementptr inbounds i8, i8 addrspace(11)* %a33, i64 8
                  %a35 = bitcast i8 addrspace(11)* %a34 to {} addrspace(10)* addrspace(11)*
                  store atomic {} addrspace(10)* %v2, {} addrspace(10)* addrspace(11)* %a35 release, align 8
                  ret {} addrspace(10)* %newstruct
                }

                define {} addrspace(10)* @caller({} addrspace(10)* %v1, {} addrspace(10)* %v2) {
                top:
                  %ac = call fastcc nonnull {} addrspace(10)* @inner({} addrspace(10)* %v1, {} addrspace(10)* %v2)
                  %b = addrspacecast {} addrspace(10)* %ac to {} addrspace(10)* addrspace(11)*
                  %c = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %b unordered, align 8
                  ret {} addrspace(10)* %c
                }

                attributes #5 = { inaccessiblememonly mustprogress nofree nounwind willreturn allockind("alloc,uninitialized") allocsize(1) "enzyme_no_escaping_allocation" "enzymejl_world"="31504" }
                """
            )

            Enzyme.Compiler.removeDeadArgs!(mod, Enzyme.Compiler.JIT.get_tm(), false)
            string(mod)
        end
    end
end


@testset "Dead return removal" begin
    @test @filecheck begin
        # The call to the (now dead) callee is removed, leaving only `ret void`,
        # and the callee itself is deleted from the module.
        @check_label "define void @caller"
        @check_next "top:"
        @check_next "ret void"
        @check_not "@julia_MyPrognosticVars_161"
        LLVM.Context() do ctx
            mod = parse(
                LLVM.Module, """
                source_filename = "start"
                target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128-ni:10:11:12:13"
                target triple = "x86_64-linux-gnu"

                declare noalias nonnull {} addrspace(10)* @julia.gc_alloc_obj({}**, i64, {} addrspace(10)*) local_unnamed_addr #5

                define internal fastcc nonnull {} addrspace(10)* @julia_MyPrognosticVars_161({} addrspace(10)* %v1, {} addrspace(10)* %v2) {
                top:
                  %newstruct = call noalias nonnull dereferenceable(16) {} addrspace(10)* @julia.gc_alloc_obj({}** null, i64 16, {} addrspace(10)* addrspacecast ({}* inttoptr (i64 129778359735376 to {}*) to {} addrspace(10)*)) #30
                  %a31 = addrspacecast {} addrspace(10)* %newstruct to {} addrspace(10)* addrspace(11)*
                  %a32 = getelementptr inbounds {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %a31, i64 1
                  store atomic {} addrspace(10)* %v1, {} addrspace(10)* addrspace(11)* %a31 release, align 8
                  %a33 = addrspacecast {} addrspace(10)* %newstruct to i8 addrspace(11)*
                  %a34 = getelementptr inbounds i8, i8 addrspace(11)* %a33, i64 8
                  %a35 = bitcast i8 addrspace(11)* %a34 to {} addrspace(10)* addrspace(11)*
                  store atomic {} addrspace(10)* %v2, {} addrspace(10)* addrspace(11)* %a35 release, align 8
                  ret {} addrspace(10)* %newstruct
                }

                define void @caller({} addrspace(10)* %v1, {} addrspace(10)* %v2) {
                top:
                  %ac = call fastcc nonnull {} addrspace(10)* @julia_MyPrognosticVars_161({} addrspace(10)* %v1, {} addrspace(10)* %v2)
                  ret void
                }

                attributes #5 = { inaccessiblememonly mustprogress nofree nounwind willreturn allockind("alloc,uninitialized") allocsize(1) "enzyme_no_escaping_allocation" "enzymejl_world"="31504" }
                """
            )

            Enzyme.Compiler.removeDeadArgs!(mod, Enzyme.Compiler.JIT.get_tm(), false)
            string(mod)
        end
    end
end

@testset "Return roots preservation" begin
    @test @filecheck begin
        # The dead `%data` argument is dropped from the call, but the `sret` and
        # `enzymejl_returnRoots` arguments must be preserved (roots last).
        @check_label "define void @caller"
        @check "call void @julia_dims_4189"
        @check_same "sret"
        @check_same "enzymejl_returnRoots"
        @check_same "%roots)"
        LLVM.Context() do ctx
            mod = parse(
                LLVM.Module, """
                define private void @julia_dims_4189({ double, {} addrspace(10)*, {} addrspace(10)* }* sret({ double, {} addrspace(10)*, {} addrspace(10)* }) %res, [2 x {} addrspace(10)*]* "enzymejl_returnRoots"="2", double addrspace(11)* %data) #0 {
                top:
                  %val = load double, double addrspace(11)* %data, align 8
                  store { double, {} addrspace(10)*, {} addrspace(10)* } zeroinitializer, { double, {} addrspace(10)*, {} addrspace(10)* }* %res
                  ret void
                }

                define void @caller({} addrspace(10)* %v1, {} addrspace(10)* %v2, double addrspace(11)* %data) {
                top:
                  %sret = alloca { double, {} addrspace(10)*, {} addrspace(10)* }
                  %roots = alloca [2 x {} addrspace(10)*]
                  call void @julia_dims_4189({ double, {} addrspace(10)*, {} addrspace(10)* }* sret({ double, {} addrspace(10)*, {} addrspace(10)* }) %sret, [2 x {} addrspace(10)*]* "enzymejl_returnRoots"="2" %roots, double addrspace(11)* %data)
                  ret void
                }

                attributes #0 = { nofree nosync nounwind willreturn noinline "enzyme_inactive" }
                """
            )

            Enzyme.Compiler.removeDeadArgs!(mod, Enzyme.Compiler.JIT.get_tm(), true)
            string(mod)
        end
    end
end

@testset "Recursively dead function removal" begin
    @test @filecheck begin
        # Both the recursive function and its callee are dead and must be removed.
        @check_not "@dead_recursive_fn"
        @check_not "@dead_callee"
        LLVM.Context() do ctx
            mod = parse(
                LLVM.Module, """
                source_filename = "start"
                target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128-ni:10:11:12:13"
                target triple = "x86_64-linux-gnu"

                define internal fastcc void @dead_callee(i32* nocapture %arg) {
                top:
                  %val = load i32, i32* %arg, align 4
                  ret void
                }

                define internal fastcc void @dead_recursive_fn(i32* %arg) {
                top:
                  call fastcc void @dead_recursive_fn(i32* %arg)
                  call fastcc void @dead_callee(i32* %arg)
                  ret void
                }
                """
            )

            Enzyme.Compiler.removeDeadArgs!(mod, Enzyme.Compiler.JIT.get_tm(), false)
            string(mod)
        end
    end
end

@testset "Mismatched calling convention/function type DAE safety" begin
    @test @filecheck begin
        # The call site bitcasts the callee to a mismatched function type, so DAE
        # must leave the callee untouched: both of its arguments are preserved.
        @check "define internal fastcc void @callee("
        @check_same "%arg1"
        @check_same "%arg2"
        LLVM.Context() do ctx
            mod = parse(
                LLVM.Module, """
                source_filename = "start"
                target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128-ni:10:11:12:13"
                target triple = "x86_64-linux-gnu"

                define internal fastcc void @callee(i32* %arg1, i32* %arg2) {
                top:
                  store i32 42, i32* %arg1, align 4
                  ret void
                }

                define void @caller(i32 addrspace(10)* %arg1, i32 addrspace(10)* %arg2) {
                top:
                  call fastcc void (i32 addrspace(10)*, i32 addrspace(10)*) bitcast (void (i32*, i32*)* @callee to void (i32 addrspace(10)*, i32 addrspace(10)*)*)(i32 addrspace(10)* %arg1, i32 addrspace(10)* %arg2)
                  ret void
                }
                """
            )

            Enzyme.Compiler.removeDeadArgs!(mod, Enzyme.Compiler.JIT.get_tm(), false)
            string(mod)
        end
    end
end

@testset "link_split_existing!" begin
    LLVM.Context() do ctx
        # Test 1: Differing definitions are preserved as distinct internal specializations
        dst = parse(
            LLVM.Module,
            """
            define i64 @julia___dup(i64 %x) {
              %r = add i64 %x, 1
              ret i64 %r
            }
            define i64 @only_in_dst(i64 %x) {
              ret i64 %x
            }
            """,
        )
        src = parse(
            LLVM.Module,
            """
            define i64 @julia___dup(i64 %x) {
              %r = add i64 %x, 2
              ret i64 %r
            }
            define i64 @uses_dup(i64 %x) {
              %r = call i64 @julia___dup(i64 %x)
              ret i64 %r
            }
            """,
        )

        Enzyme.Compiler.link_split_existing!(dst, src)

        fns = LLVM.functions(dst)
        @test haskey(fns, "julia___dup")
        @test !LLVM.isdeclaration(fns["julia___dup"])
        @test LLVM.linkage(fns["julia___dup"]) == LLVM.API.LLVMExternalLinkage
        @test haskey(fns, "only_in_dst")
        @test haskey(fns, "uses_dup")
        usesfn = fns["uses_dup"]
        callinst = first(
            filter(
                Base.Fix2(isa, LLVM.CallInst),
                collect(instructions(first(blocks(usesfn)))),
            ),
        )
        called_fn = last(collect(operands(callinst)))
        @test LLVM.linkage(called_fn) == LLVM.API.LLVMInternalLinkage

        # Test 2: Identical definitions are folded by MergeFunctionsPass
        dst2 = parse(
            LLVM.Module,
            """
            define i64 @julia___dup(i64 %x) {
              %r = add i64 %x, 1
              ret i64 %r
            }
            """,
        )
        src2 = parse(
            LLVM.Module,
            """
            define i64 @julia___dup(i64 %x) {
              %r = add i64 %x, 1
              ret i64 %r
            }
            define i64 @uses_dup2(i64 %x) {
              %r = call i64 @julia___dup(i64 %x)
              ret i64 %r
            }
            """,
        )
        Enzyme.Compiler.link_split_existing!(dst2, src2)
        fns2 = LLVM.functions(dst2)
        @test haskey(fns2, "julia___dup")
        @test haskey(fns2, "uses_dup2")
        usesfn2 = fns2["uses_dup2"]
        callinst2 = first(
            filter(
                Base.Fix2(isa, LLVM.CallInst),
                collect(instructions(first(blocks(usesfn2)))),
            ),
        )
        @test LLVM.name(last(collect(operands(callinst2)))) == "julia___dup"
    end
end

@testset "Literal-pointer symbol resolution" begin
    # `jl_lookup_code_address` can attribute two distinct pointers to the same
    # (nearest) symbol name on platforms with export-only symbol info. The
    # helpers below back the guards in `check_ir!` that keep such call sites
    # from being merged onto one restored address.
    libmpfr = Libdl.dlpath(Base.MPFR.libmpfr)
    hnd = Libdl.dlopen(libmpfr)
    p_add = Libdl.dlsym(hnd, :mpfr_add)
    p_sub = Libdl.dlsym(hnd, :mpfr_sub)
    @test Enzyme.Compiler.resolve_symbol_name("mpfr_add", libmpfr, p_add) == :match
    # A nearest-symbol misattribution must be detected, not trusted.
    @test Enzyme.Compiler.resolve_symbol_name("mpfr_add", libmpfr, p_sub) == :mismatch
    # The containing module is found via the loader, so an unloadable `file` (as
    # reported on Linux/macOS, where it is a source path) does not degrade the answer.
    @test Enzyme.Compiler.resolve_symbol_name("mpfr_add", "/not/a/library/add.c", p_add) ==
        :match
    @test Enzyme.Compiler.resolve_symbol_name("mpfr_add", "/not/a/library/add.c", p_sub) ==
        :mismatch
    @test Enzyme.Compiler.resolve_symbol_name("not_a_real_symbol_abcxyz", libmpfr, p_add) ==
        :unknown
    # A pointer outside any loaded module cannot be attributed at all.
    heapptr = Libc.malloc(8)
    @test Enzyme.Compiler.resolve_symbol_name("mpfr_add", "", heapptr) == :unknown
    Libc.free(heapptr)

    LLVM.Context() do ctx
        mod = parse(
            LLVM.Module,
            """
            declare void @foo() #0
            declare void @bar()
            attributes #0 = { "enzymejl_needs_restoration"="12345" }
            """,
        )
        @test Enzyme.Compiler.restoration_ptr(functions(mod)["foo"]) == UInt(12345)
        @test Enzyme.Compiler.restoration_ptr(functions(mod)["bar"]) === nothing
    end
end

# https://github.com/EnzymeAD/Enzyme.jl/issues/3284
# The `jl_get_abi_converter` lowering of `@cfunction` only exists on 1.12+, and
# older Julia's LLVM cannot parse the opaque-pointer IR below.
@static if VERSION >= v"1.12"

    @testset "Rewrite abi converter calls (1.12 pattern)" begin
        LLVM.Context() do ctx
            # Shape of the `@cfunction` dispatch site emitted by Julia 1.12's
            # codegen: `jl_get_abi_converter(ct, fptr, last_world, cfuncdata)` with
            # a six-slot cfuncdata whose third slot holds the in-module
            # `unspecialized` apply-generic thunk.
            mod = parse(
                LLVM.Module,
                """
                @jl_world_counter = external global i64
                @fptr = private global ptr @gfthunk
                @last_world = private global i64 0
                @declrt = private global ptr null
                @sigt = private global ptr null
                @cfuncdata = private global [6 x ptr] [ptr null, ptr null, ptr @gfthunk, ptr @declrt, ptr @sigt, ptr inttoptr (i64 1 to ptr)]

                declare ptr @jl_get_abi_converter(ptr, ptr, ptr, ptr)

                define internal double @gfthunk(double %0) {
                top:
                  ret double %0
                }

                define double @trampoline(ptr %ct, double %x) {
                top:
                  %last_world = load atomic i64, ptr @last_world acquire, align 8
                  %fptr = load atomic ptr, ptr @fptr monotonic, align 8
                  %world = load atomic i64, ptr @jl_world_counter acquire, align 8
                  %stale = icmp ne i64 %last_world, %world
                  br i1 %stale, label %guard_pass, label %guard_exit

                guard_pass:
                  %cw = call ptr @jl_get_abi_converter(ptr %ct, ptr @fptr, ptr @last_world, ptr @cfuncdata)
                  br label %guard_exit

                guard_exit:
                  %target = phi ptr [ %fptr, %top ], [ %cw, %guard_pass ]
                  %res = call double %target(double %x)
                  ret double %res
                }
                """,
            )
            Enzyme.Compiler.rewrite_abi_converter_calls!(mod)
            @test @filecheck begin
                @check_label "define double @trampoline"
                @check_not "call ptr @jl_get_abi_converter"
                @check "phi ptr [ %fptr, %top ], [ @gfthunk, %guard_pass ]"
                string(mod)
            end
        end
    end

    @testset "Rewrite abi converter calls (1.13 pattern)" begin
        LLVM.Context() do ctx
            # Shape of the dispatch site emitted by Julia 1.13+:
            # `jl_get_abi_converter(ct, cfuncdata)` with an eight-slot cfuncdata
            # ([fptr, last_world, plast_codeinst, last_codeinst, unspecialized,
            # declrt, sigt, flags]) whose fifth slot holds the thunk.
            mod = parse(
                LLVM.Module,
                """
                @jl_world_counter = external global i64
                @declrt = private global ptr null
                @sigt = private global ptr null
                @cfuncdata = private global [8 x ptr] [ptr @gfthunk, ptr null, ptr null, ptr null, ptr @gfthunk, ptr @declrt, ptr @sigt, ptr inttoptr (i64 1 to ptr)]

                declare ptr @ijl_get_abi_converter(ptr, ptr)

                define internal double @gfthunk(double %0) {
                top:
                  ret double %0
                }

                define double @trampoline(ptr %ct, double %x) {
                top:
                  %last_world_p = getelementptr inbounds i64, ptr @cfuncdata, i32 1
                  %last_world = load atomic i64, ptr %last_world_p acquire, align 8
                  %fptr = load atomic ptr, ptr @cfuncdata monotonic, align 8
                  %world = load atomic i64, ptr @jl_world_counter acquire, align 8
                  %stale = icmp ne i64 %last_world, %world
                  br i1 %stale, label %guard_pass, label %guard_exit

                guard_pass:
                  %cw = call ptr @ijl_get_abi_converter(ptr %ct, ptr @cfuncdata)
                  br label %guard_exit

                guard_exit:
                  %target = phi ptr [ %fptr, %top ], [ %cw, %guard_pass ]
                  %res = call double %target(double %x)
                  ret double %res
                }
                """,
            )
            Enzyme.Compiler.rewrite_abi_converter_calls!(mod)
            @test @filecheck begin
                @check_label "define double @trampoline"
                @check_not "call ptr @ijl_get_abi_converter"
                @check "phi ptr [ %fptr, %top ], [ @gfthunk, %guard_pass ]"
                string(mod)
            end
        end
    end

end # VERSION >= v"1.12"

# An addrspace(11) phi with argument-derived and poison incomings must be skipped by
# the `all_args` prefilter; the rewriter cannot handle a bare addrspace(11) argument.
@testset "nodecayed_phis! addrspace(11) phi with poison incoming" begin
    @test @filecheck begin
        # incomings: gep(arg, 8) and poison
        @check_label "@kernel_gep"
        @check_not "nodecayed"
        @check "phi"
        @check_same "addrspace(11)"
        @check_same "%gep"
        @check_same "poison"
        # zero-offset variant: the argument itself and poison
        @check_label "@kernel_direct"
        @check_not "nodecayed"
        @check "phi"
        @check_same "addrspace(11)"
        @check_same "%q"
        @check_same "poison"
        LLVM.Context() do ctx
            mod = parse(
                LLVM.Module, """
                source_filename = "start"
                target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128-ni:10:11:12:13"
                target triple = "x86_64-linux-gnu"

                define i8 @kernel_gep(i1 %cond, i8 addrspace(11)* %g) #0 {
                top:
                  br i1 %cond, label %ok, label %guard

                ok:
                  %gep = getelementptr inbounds i8, i8 addrspace(11)* %g, i64 8
                  br label %merge

                guard:
                  br label %merge

                merge:
                  %p = phi i8 addrspace(11)* [ %gep, %ok ], [ poison, %guard ]
                  %ld = load i8, i8 addrspace(11)* %p, align 1
                  ret i8 %ld
                }

                define i8 @kernel_direct(i1 %cond, i8 addrspace(11)* %q) #0 {
                top:
                  br i1 %cond, label %ok2, label %guard2

                ok2:
                  br label %merge2

                guard2:
                  br label %merge2

                merge2:
                  %p2 = phi i8 addrspace(11)* [ %q, %ok2 ], [ poison, %guard2 ]
                  %ld2 = load i8, i8 addrspace(11)* %p2, align 1
                  ret i8 %ld2
                }

                attributes #0 = { "enzymejl_world"="1" }
                """
            )

            Enzyme.Compiler.nodecayed_phis!(mod)
            string(mod)
        end
    end
end
