Skip to content

Commit 4d7bcc9

Browse files
committed
RFC: Introduce TypedCallable as a modern Function wrapper
TypedCallable provides a wrapper for callable objects, with the following benefits: 1. Enforced type-stability (for concrete AT/RT types) 2. Fast calling convention (frequently < 10 ns / call) 3. Normal Julia dispatch semantics (sees new Methods, etc.) + invoke_latest 4. Pre-compilation support (including `--trim` compatibility) It can be used like this: ```julia callbacks = @TypedCallable{(::Int,::Int)->Bool}[] register_callback!(callbacks, f::F) where {F<:Function} = push!(callbacks, @TypedCallable f(::Int,::Int)::Bool) register_callback!(callbacks, (x,y)->(x == y)) register_callback!(callbacks, (x,y)->(x != y)) @Btime callbacks[rand(1:2)](1,1) ``` This is very similar to the existing `FunctionWrappers.jl`, but there are a few key differences: - Better type support: TypedCallable supports the full range of Julia types (incl. Varargs), and it has access to all of Julia's "internal" calling conventions so calls are fast (and allocation-free) for a wider range of input types - Improved dispatch handling: The `@cfunction` functionality used by FunctionWrappers has several dispatch bugs, which cause wrappers to occasionally not see new Methods. These bugs are fixed (or soon to be fixed) for TypedCallable. - Pre-compilation support including for `juliac` / `--trim` (#55047) Many of the improvements here are actually thanks to the `OpaqueClosure` introduced by @Keno - This type just builds on top of OpaqueClosure to provide an interface with Julia's usual dispatch semantics.
1 parent 2b140ba commit 4d7bcc9

File tree

1 file changed

+193
-0
lines changed

1 file changed

+193
-0
lines changed

base/opaque_closure.jl

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,3 +110,196 @@ function generate_opaque_closure(@nospecialize(sig), @nospecialize(rt_lb), @nosp
110110
return ccall(:jl_new_opaque_closure_from_code_info, Any, (Any, Any, Any, Any, Any, Cint, Any, Cint, Cint, Any, Cint, Cint),
111111
sig, rt_lb, rt_ub, mod, src, lineno, file, nargs, isva, env, do_compile, isinferred)
112112
end
113+
114+
struct Slot{T} end
115+
struct Splat{T}
116+
value::T
117+
end
118+
119+
# Args is a Tuple{Vararg{Union{Slot{T},Some{T}}}} where Slot{T} represents
120+
# an uncurried argument slot, and Some{T} represents an argument to curry.
121+
@noinline @generated function Core.OpaqueClosure(Args::Tuple, ::Slot{RT}) where RT
122+
AT = Any[]
123+
call = Expr(:call)
124+
extracted = Expr[]
125+
closure_args = Expr(:tuple)
126+
for (i, T) in enumerate(Args.parameters)
127+
v = Symbol("arg", i)
128+
is_splat = T <: Splat
129+
if is_splat # TODO: check position
130+
push!(call.args, :($v...))
131+
T = T.parameters[1]
132+
else
133+
push!(call.args, v)
134+
end
135+
if T <: Some
136+
push!(extracted, :($v = something(Args[$i])))
137+
elseif T <: Slot
138+
SlotT = T.parameters[1]
139+
push!(AT, is_splat ? Vararg{SlotT} : SlotT)
140+
push!(closure_args.args, call.args[end])
141+
else @assert false end
142+
end
143+
AT = Tuple{AT...}
144+
return Base.remove_linenums!(quote
145+
$(extracted...)
146+
@opaque allow_partial=false $AT->$RT $(closure_args)->@inline $(call)
147+
end)
148+
end
149+
150+
"""
151+
TypedCallable{AT,RT}
152+
153+
TypedCallable provides a wrapper for callable objects, with the following benefits:
154+
1. Enforced type-stability (for concrete AT/RT types)
155+
2. Fast calling convention (frequently < 10 ns / call)
156+
3. Normal Julia dispatch semantics (sees new Methods, etc.) + invoke_latest
157+
158+
This was specifically introduced to provide a form of Function-like that will be
159+
compatible with the newly-introduced `--trim` without requiring fully specializing
160+
on (singleton) `Function` types.
161+
162+
This is similar to the existing FunctionWrappers.jl, with the advantages that
163+
TypedCallable supports the full range of Julia types/arguments, has a faster calling
164+
convention for a wider range of types, and has better pre-compilation support
165+
(including support for `--trim`).
166+
167+
Examples:
168+
169+
170+
If `invoke_latest` is `false`, the provide TypedCallable has normal Julia semantics.
171+
Otherwise it always attempts to make the call in the current world.
172+
173+
# Extended help
174+
175+
### As an invalidation barrier
176+
177+
TypedCallable can also be used as an "invalidation barrier", since the caller of a
178+
TypedCallable is not affected by any invalidations of its callee(s). This doesn't
179+
completely cure the original invalidation, but it stops it from propagating all the
180+
way through your code.
181+
182+
This can be especially helpful, e.g., when calling back to user-provided functions
183+
whose invalidations you may have no control over.
184+
"""
185+
mutable struct TypedCallable{AT,RT}
186+
@atomic oc::Base.RefValue{Core.OpaqueClosure{AT,RT}}
187+
const task::Union{Task,Nothing}
188+
const build_oc::Function
189+
end
190+
191+
function Base.show(io::IO, tc::Base.Experimental.TypedCallable)
192+
A, R = typeof(tc).parameters
193+
Base.print(io, "@TypedCallable{")
194+
Base.show_tuple_as_call(io, Symbol(""), A; hasfirst=false)
195+
Base.print(io, "->◌::", R, "}()")
196+
end
197+
198+
function rebuild_in_world!(@nospecialize(self::TypedCallable), world::UInt)
199+
oc = Base.invoke_in_world(world, self.build_oc)
200+
@atomic :release self.oc = Base.Ref(oc)
201+
return oc
202+
end
203+
204+
@inline function (self::TypedCallable{AT,RT})(args...) where {AT,RT}
205+
invoke_world = if self.task === nothing
206+
Base.get_world_counter() # Base.unsafe_load(cglobal(:jl_world_counter, UInt), :acquire) ?
207+
elseif self.task === Base.current_task()
208+
Base.tls_world_age()
209+
else
210+
error("TypedCallable{...} was called from a different task than it was created in.")
211+
end
212+
oc = (@atomic :acquire self.oc)[]
213+
if oc.world != invoke_world
214+
oc = @noinline rebuild_in_world!(self, invoke_world)::Core.OpaqueClosure{AT,RT}
215+
end
216+
return oc(args...)
217+
end
218+
219+
function _TypedCallable_type(ex)
220+
type_err = "Invalid @TypedCallable expression: $(ex)\nExpected \"@TypedCallable{(::T,::U,...)->RT}\""
221+
222+
# Unwrap {...}
223+
(length(ex.args) != 1) && error(type_err)
224+
ex = ex.args[1]
225+
226+
# Unwrap (...)->RT
227+
!(Base.isexpr(ex, :->) && length(ex.args) == 2) && error(type_err)
228+
tuple_, rt = ex.args
229+
if !(Base.isexpr(tuple_, :tuple) && all((x)->Base.isexpr(x, :(::)), tuple_.args))
230+
# note: (arg::T, ...) is specifically allowed (the "arg" part is unused)
231+
error(type_err)
232+
end
233+
!Base.isexpr(rt, :block) && error(type_err)
234+
235+
# Remove any LineNumberNodes inserted by lowering
236+
filter!((x)->!isa(x,Core.LineNumberNode), rt.args)
237+
(length(rt.args) != 1) && error(type_err)
238+
239+
# Build args
240+
AT = Expr[esc(last(x.args)) for x in tuple_.args]
241+
RT = rt.args[1]
242+
243+
# Unwrap ◌::T to T
244+
if Base.isexpr(RT, :(::)) && length(RT.args) == 2 && RT.args[1] == :◌
245+
RT = RT.args[2]
246+
end
247+
248+
return :($TypedCallable{Tuple{$(AT...)}, $(esc(RT))})
249+
end
250+
251+
function _TypedCallable_closure(ex)
252+
if Base.isexpr(ex, :call)
253+
error("""
254+
Invalid @TypedCallable expression: $(ex)
255+
An explicit return type assert is required (e.g. "@TypedCallable f(...)::RT")
256+
""")
257+
end
258+
259+
call_, RT = ex.args
260+
if !Base.isexpr(call_, :call)
261+
error("""Invalid @TypedCallable expression: $(ex)
262+
The supported syntax is:
263+
@TypedCallable{(::T,::U,...)->RT} (to construct the type)
264+
@TypedCallable f(x,::T,...)::RT (to construct the TypedCallable)
265+
""")
266+
end
267+
oc_args = map(call_.args) do arg
268+
is_splat = Base.isexpr(arg, :(...))
269+
arg = is_splat ? arg.args[1] : arg
270+
transformed = if Base.isexpr(arg, :(::))
271+
if length(arg.args) == 1 # it's a "slot"
272+
slot_ty = esc(only(arg.args))
273+
:(Slot{$slot_ty}())
274+
elseif length(arg.args) == 2
275+
(arg, ty) = arg.args
276+
:(Some{$(esc(ty))}($(esc(arg))))
277+
else @assert false end
278+
else
279+
:(Some($(esc(arg))))
280+
end
281+
return is_splat ? Expr(:call, Splat, transformed) : transformed
282+
end
283+
# TODO: kwargs support
284+
RT = :(Slot{$(esc(RT))}())
285+
invoke_latest = true # expose as flag?
286+
task = invoke_latest ? nothing : :(Base.current_task())
287+
return quote
288+
build_oc = ()->Core.OpaqueClosure(($(oc_args...),), $(RT))
289+
$(TypedCallable)(Ref(build_oc()), $task, build_oc)
290+
end
291+
end
292+
293+
macro TypedCallable(ex)
294+
if Base.isexpr(ex, :braces)
295+
return _TypedCallable_type(ex)
296+
elseif Base.isexpr(ex, :call) || (Base.isexpr(ex, :(::)) && length(ex.args) == 2)
297+
return _TypedCallable_closure(ex)
298+
else
299+
error("""Invalid @TypedCallable expression: $(ex)
300+
The supported syntax is:
301+
@TypedCallable{(::T,::U,...)->RT} (to construct the type)
302+
@TypedCallable f(x,::T,...)::RT (to construct the TypedCallable)
303+
""")
304+
end
305+
end

0 commit comments

Comments
 (0)