Reapply "Rewrite type handler to set classification immediatey (#3239)"

[wsmoses]

This reverts commit fe4b643.

I landed the wrong pr by accident

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.

diff --git a/src/compiler.jl b/src/compiler.jl
index d8e0d29d..bdfc7b8b 100644
--- a/src/compiler.jl
+++ b/src/compiler.jl
@@ -592,7 +592,7 @@ function prepare_llvm(interp, mod::LLVM.Module, job, meta)
         if !Base.isvarargtype(mi.specTypes.parameters[end])
             dl = string(LLVM.datalayout(LLVM.parent(llvmfn)))
             ctx = LLVM.context(llvmfn)
-            
+
             # Since this is prepare_llvm, no parameters can be removed yet.
             empty_parms = UInt64[]
             jlargs = classify_arguments(
@@ -681,7 +681,7 @@ function prepare_llvm(interp, mod::LLVM.Module, job, meta)
                 end
 
                 if llRT !== nothing &&
-                   LLVM.return_type(LLVM.function_type(llvmfn)) != LLVM.VoidType()
+                        LLVM.return_type(LLVM.function_type(llvmfn)) != LLVM.VoidType()
                     rest = if llRT == Ptr{RT}
                         typeTree = copy(typetree(RT, ctx, dl))
                         merge!(typeTree, TypeTree(API.DT_Pointer, ctx))
@@ -1233,7 +1233,7 @@ end
     return
 end
 
-function set_module_types!(interp, mod::LLVM.Module, primalf::Union{Nothing, LLVM.Function}, job, edges, run_enzyme, mode::API.CDerivativeMode)::Tuple{Dict{String,LLVM.API.LLVMLinkage}, HandlerState}
+    return function set_module_types!(interp, mod::LLVM.Module, primalf::Union{Nothing, LLVM.Function}, job, edges, run_enzyme, mode::API.CDerivativeMode)::Tuple{Dict{String, LLVM.API.LLVMLinkage}, HandlerState}
     custom = Dict{String,LLVM.API.LLVMLinkage}()
 
     world = job.world
@@ -1251,11 +1251,11 @@ function set_module_types!(interp, mod::LLVM.Module, primalf::Union{Nothing, LLV
     )
 
     for fname in LLVM.name.(functions(mod))
-        if !haskey(functions(mod), fname) || startswith(fname, "japi3") || startswith(fname, "japi1")
+            if !haskey(functions(mod), fname) || startswith(fname, "japi3") || startswith(fname, "japi1")
             continue
         end
         fn = functions(mod)[fname]
-        mi, RT = enzyme_custom_extract_mi(fn, false)
+            mi, RT = enzyme_custom_extract_mi(fn, false)
         if mi !== nothing && RT !== nothing
             handle_compiled(state, edges, run_enzyme, mode, world, method_table, custom, mod, mi, fname, RT)
         end
@@ -4498,7 +4498,7 @@ function lower_convention(
             if arg.arg_i in removedRoots
                 continue
             end
-            
+
             # Copy immutable attributes
             push!(
                 parameter_attributes(wrapper_f, wrapper_idx),
diff --git a/src/typeutils/jltypes.jl b/src/typeutils/jltypes.jl
index 939a7a74..51888f6c 100644
--- a/src/typeutils/jltypes.jl
+++ b/src/typeutils/jltypes.jl
@@ -274,8 +274,8 @@ function classify_arguments(
     parmsRemoved::Vector{UInt64},
     mi::Core.MethodInstance,
     world::UInt,
-    llvm_f::LLVM.Function,
-    called_at_start::Bool,
+        llvm_f::LLVM.Function,
+        called_at_start::Bool,
 )
     if called_at_start
         @assert isempty(parmsRemoved) "called_at_start=true but parmsRemoved is not empty"

[github-actions]

Benchmark Results

	main	142b666...	main / 142b666...
basics/make_zero/namedtuple	0.0525 ± 0.0023 μs	0.0521 ± 0.0016 μs	1.01 ± 0.054
basics/make_zero/struct	0.262 ± 0.0064 μs	0.261 ± 0.0062 μs	1 ± 0.034
basics/overhead	4.33 ± 0.01 ns	4.34 ± 0.01 ns	0.998 ± 0.0033
basics/remake_zero!/namedtuple	0.227 ± 0.0095 μs	0.228 ± 0.0087 μs	0.996 ± 0.056
basics/remake_zero!/struct	0.223 ± 0.01 μs	0.224 ± 0.012 μs	0.997 ± 0.07
fold_broadcast/multidim_sum_bcast/1D	10.2 ± 0.27 μs	10.2 ± 1.8 μs	0.999 ± 0.18
fold_broadcast/multidim_sum_bcast/2D	12 ± 0.24 μs	12.1 ± 0.23 μs	0.995 ± 0.028
time_to_load	1.45 ± 0.005 s	1.45 ± 0.0091 s	1 ± 0.0072

Benchmark Plots

A plot of the benchmark results has been uploaded as an artifact at https://github.com/EnzymeAD/Enzyme.jl/actions/runs/28308154242/artifacts/7930142236.