NFloat: Improve handling of type parameters in ArbCall wrapper

Author: Joel-Dahne

src/ArbCall/ArbFunction.jl

@@ -201,13 +201,22 @@ function jlcode(af::ArbFunction, jl_fname = jlfname(af))
 
     returnT = returntype(af)
     cargs = arguments(af)
+    where_type_parameters = unique(reduce(vcat, type_parameters.(cargs)))
 
-    func_full_args = :(
-        function $jl_fname($(jl_full_args...))
+    func_full_args_call = :($jl_fname($(jl_full_args...)))
+
+    func_full_args_header = if isempty(where_type_parameters)
+        func_full_args_call
+    else
+        Expr(:where, func_full_args_call, where_type_parameters...)
+    end
+
+    func_full_args_body = :(
+        begin
             __ret = ccall(
                 Arblib.@libflint($(arbfname(af))),
                 $returnT,
-                $(Expr(:tuple, ctype.(cargs)...)),
+                $(Expr(:tuple, carg_expr.(cargs)...)),
                 $(name.(cargs)...),
             )
             $(
@@ -222,7 +231,11 @@ function jlcode(af::ArbFunction, jl_fname = jlfname(af))
         end
     )
 
+    func_full_args = Expr(:function, func_full_args_header, func_full_args_body)
+
     if is_series_method(af)
+        @assert isempty(where_type_parameters) # Currently not supported for series methods
+
         # Note that this currently doesn't respect any custom function
         # name given as an argument.
         jl_fname_series = jlfname_series(af)
@@ -244,9 +257,16 @@ function jlcode(af::ArbFunction, jl_fname = jlfname(af))
     if isempty(jl_kwargs)
         return code
     else
+        func_kwarg_args_call = :($jl_fname($(jl_args...); $(jl_kwargs...)))
+        func_kwarg_args_header = if isempty(where_type_parameters)
+            func_kwarg_args_call
+        else
+            Expr(:where, func_kwarg_args_call, where_type_parameters...)
+        end
+
         return quote
             $code
-            $jl_fname($(jl_args...); $(jl_kwargs...)) = $jl_fname($(name.(cargs)...))
+            $func_kwarg_args_header = $jl_fname($(name.(cargs)...))
         end
     end
 end

src/ArbCall/Carg.jl

@@ -93,6 +93,12 @@ jltype(::Carg{ArbMatrix}) = ArbMatrixLike
 # acb_mat.h
 jltype(::Carg{AcbMatrix}) = AcbMatrixLike
 
+type_parameters_from_type(type::UnionAll) =
+    [type.var.name; type_parameters_from_type(type.body)]
+type_parameters_from_type(_) = Symbol[]
+
+type_parameters(ca::Carg) = type_parameters_from_type(jltype(ca))
+
 """
     ctype(ca::Carg)
 
@@ -105,8 +111,8 @@ ctype(::Carg{T}) where {T<:Union{Mag,Arf,Acf,Arb,Acb,ArbPoly,AcbPoly,ArbMatrix,A
     Ref{cstructtype(T)}
 ctype(::Carg{T}) where {T<:Union{ArbVector,arb_vec_struct}} = Ptr{arb_struct}
 ctype(::Carg{T}) where {T<:Union{AcbVector,acb_vec_struct}} = Ptr{acb_struct}
-ctype(::Carg{T}) where {T<:NFloat} = Ref{nfloat_struct}
-ctype(::Carg{T}) where {T<:nfloat_ctx_struct} = Ref{nfloat_ctx_struct}
+ctype(::Carg{T}) where {T<:NFloat} = Ref{nfloat_struct{P,F}} where {P,F}
+ctype(::Carg{T}) where {T<:nfloat_ctx_struct} = Ref{nfloat_ctx_struct{P,F}} where {P,F}
 
 """
     jlarg(ca::Carg{T}) where {T}
@@ -119,9 +125,39 @@ julia> Arblib.ArbCall.jlarg(Arblib.ArbCall.Carg("const arb_t x"))
 :(x::ArbLike)
 julia> Arblib.ArbCall.jlarg(Arblib.ArbCall.Carg("slong prec"))
 :(prec::Integer)
+julia> Arblib.ArbCall.jlarg(Arblib.ArbCall.Carg("nfloat_ptr res"))
+:(res::(NFloatLike){P,F})
 ```
 """
-jlarg(ca::Carg) = :($(name(ca))::$(jltype(ca)))
+function jlarg(ca::Carg)
+    if jltype(ca) isa UnionAll
+        :($(name(ca))::$(jltype(ca)){$(type_parameters(ca)...)})
+    else
+        :($(name(ca))::$(jltype(ca)))
+    end
+end
+
+"""
+    carg_expr(ca::Carg{T}) where {T}
+
+Return a value for representing the argument in a `ccall`.
+
+```jldoctest
+julia> Arblib.ArbCall.carg_expr(Arblib.ArbCall.Carg("const arb_t x"))
+Ref{Arblib.arb_struct}
+julia> Arblib.ArbCall.carg_expr(Arblib.ArbCall.Carg("slong prec"))
+Int64
+julia> Arblib.ArbCall.carg_expr(Arblib.ArbCall.Carg("nfloat_ptr res"))
+:((Ref{Arblib.nfloat_struct{P, F}} where {P, F}){P, F})
+```
+"""
+function carg_expr(ca::Carg)
+    if ctype(ca) isa UnionAll
+        :($(ctype(ca)){$(type_parameters(ca)...)})
+    else
+        :($(ctype(ca)))
+    end
+end
 
 is_precision_argument(ca::Carg) = ca == Carg{Int}(:prec, false)
 

src/Arblib.jl

@@ -112,4 +112,6 @@ include("arbcalls/arb_fpwrap.jl")
 include("arbcalls/fmpz_extras.jl")
 include("arbcalls/eigen.jl")
 
+include("nfloat_init_contexts.jl")
+
 end # module

src/nfloat.jl

@@ -28,7 +28,7 @@ mutable struct nfloat_struct{P,F}
     function nfloat_struct{P,F}() where {P,F}
         @assert P isa Int && F isa Int
         res = new{P,F}()
-        init!(res, nfloat_ctx_struct{P,F}())
+        init!(res)
         return res
     end
 end
@@ -49,9 +49,11 @@ const NFloatLike{P,F} = Union{NFloat{P,F},NFloatRef{P,F},nfloat_struct{P,F}}
 nfloat_ctx_struct(::NFloatLike{P,F}) where {P,F} = nfloat_ctx_struct{P,F}()
 nfloat_ctx_struct(::Type{NFloatLike{P,F}}) where {P,F} = nfloat_ctx_struct{P,F}()
 
-# TODO: Precompute these for all(?) possible values
-function _get_nfloat_ctx_struct(::Union{Type{NFloatLike{P,F}},NFloatLike{P,F}}) where {P,F}
-    return nfloat_ctx_struct{P,F}()
+# The contexts are precomputed in nfloat_late.jl
+@generated function _get_nfloat_ctx_struct(
+    ::Union{Type{<:NFloatLike{P,F}},NFloatLike{P,F}},
+) where {P,F}
+    Symbol(:_nfloat_ctx_struct_, P, :_, F)
 end
 
 # Helper function for constructing a flag argument for NFloat
@@ -179,6 +181,13 @@ end
 
 Base.show(io::IO, ::Type{NFloatLike}) = print(io, :NFloatLike)
 
+# As in promotion.jl
+
+Base.promote_rule(
+    ::Type{<:NFloatOrRef{P1,F1}},
+    ::Type{<:Union{NFloatOrRef{P2,F2}}},
+) where {P1,P2,F1,F2} = NFloat{max(P1, P2),F1 | F2}
+
 # As in arithmetic.jl
 
 for (jf, af) in [(:+, :add!), (:-, :sub!), (:*, :mul!), (:/, :div!)]

src/nfloat_init_contexts.jl

@@ -0,0 +1,5 @@
+for P = 1:66
+    for F = 0:7
+        @eval const $(Symbol(:_nfloat_ctx_struct_, P, :_, F)) = nfloat_ctx_struct{$P,$F}()
+    end
+end