From 7044b51105d268cc56f74fe5c352db3a2fdb7cdc Mon Sep 17 00:00:00 2001
From: Simeon David Schaub <simeon@schaub.rocks>
Date: Tue, 24 Sep 2024 11:43:16 +0200
Subject: [PATCH] add some JET tests

---
 src/AztecDiamonds.jl |  4 ++--
 src/ka.jl            |  4 ++--
 test/Project.toml    |  2 ++
 test/core.jl         | 10 ++++++++++
 test/makie.jl        |  3 +++
 5 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/src/AztecDiamonds.jl b/src/AztecDiamonds.jl
index 0a7c027..0a017fb 100644
--- a/src/AztecDiamonds.jl
+++ b/src/AztecDiamonds.jl
@@ -52,11 +52,11 @@ function Base.checkbounds(t::Tiling, i, j)
     return nothing
 end
 
-Base.@propagate_inbounds function Base.getindex(t::Tiling, i, j)
+Base.@propagate_inbounds function Base.getindex(t::Tiling, i::Integer, j::Integer)
     @boundscheck checkbounds(t, i, j)
     return t.x[i, j]
 end
-Base.@propagate_inbounds function Base.setindex!(t::Tiling, x, i, j)
+Base.@propagate_inbounds function Base.setindex!(t::Tiling, x, i::Integer, j::Integer)
     @boundscheck checkbounds(t, i, j)
     return setindex!(t.x, x, i, j)
 end
diff --git a/src/ka.jl b/src/ka.jl
index e1a0da8..159cef5 100644
--- a/src/ka.jl
+++ b/src/ka.jl
@@ -79,13 +79,13 @@ end
     end
 end
 
-@kernel function zero_kernel!(t::Tiling, N)  # COV_EXCL_LINE
+@kernel function zero_kernel!(t::Tiling, N::Int)  # COV_EXCL_LINE
     I = @index(Global, NTuple)  # COV_EXCL_LINE
     i, j = I .- N
     @inbounds t.x[i, j] = NONE
 end
 
-function ka_diamond!(t, t′, N; backend)
+function ka_diamond!(t::Tiling, t′::Tiling, N::Int; backend)
     zero! = zero_kernel!(backend)
     remove_bad_blocks! = remove_bad_blocks_kernel!(backend)
     slide_tiles! = slide_tiles_kernel!(backend)
diff --git a/test/Project.toml b/test/Project.toml
index 508db1f..7e4d0c6 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -6,6 +6,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
@@ -17,6 +18,7 @@ CUDA = "5.4.3"
 CairoMakie = "0.12.11"
 Colors = "0.12.11"
 Images = "0.26.1"
+JET = "0.9.9"
 Pkg = "1"
 Test = "1"
 TestItemRunner = "1.0.5"
diff --git a/test/core.jl b/test/core.jl
index 77d9016..ba58407 100644
--- a/test/core.jl
+++ b/test/core.jl
@@ -82,3 +82,13 @@ end
         @test _to_img(D′) == replace(rot(_to_img(D)), replacements...)
     end
 end
+
+@testitem "JET" begin
+    using JET
+
+    @test_opt diamond(10)
+    D = diamond(10)
+    @test_opt dr_path(D)
+    @test_opt AztecDiamonds.to_img(D)
+    @test_call show(stdout, MIME("text/plain"), D)
+end
diff --git a/test/makie.jl b/test/makie.jl
index a96668c..8088676 100644
--- a/test/makie.jl
+++ b/test/makie.jl
@@ -12,4 +12,7 @@
     save(path, f)
     @test isfile(path)
     @test filesize(path) > 1024 # 1 kiB
+
+    using JET
+    @test_opt Base.get_extension(AztecDiamonds, :MakieExtension).prepare_plot(D)
 end