I’ve been experimenting with methods for transforming arbitrary Julia code for lazy evaluation (repo at https://github.com/tbenst/Thunks.jl). Here’s an example of what I’d like to be able to do in a fraction of a second:
add1(x) = x + 1
abc = @thunk begin
b = identity(10)
z = sleep(10)
a = identity(collect(1:b))[8:end]
c,d,e,f = [x-1 for x in 1:4]
abc = sum(add1.(a) .* [c,d,f])
end
@time @assert reify(abc) == 43
However, transforming arbitrary code including dot broadcasting of functions or infix operators, variadic return values, and indexing is a bit tricky.
The core implementation is very terse:
mutable struct Thunk
f::Any # usually a Function, but could be any callable
args::Tuple # args will be passed to f
kwargs::Dict # kwargs will be passed to f
evaluated::Bool # false until computed, then true
result::Any # cache result once computed
Thunk(f, args, kwargs) = new(f, args, kwargs, false, nothing)
end
function thunk(f)
(args...; kwargs...) -> Thunk(f, args, kwargs)
end
"""
reify(thunk::Thunk)
reify(value::Any)
Reify a thunk into a value.
In other words, compute the value of the expression.
We walk through the thunk's arguments and keywords, recursively evaluating each one,
and then evaluating the thunk's function with the evaluated arguments.
"""
function reify(thunk::Thunk)
if thunk.evaluated
return thunk.result
else
args = [reify(x) for x in thunk.args]
kwargs = Dict(k => reify(v) for (k,v) in thunk.kwargs)
thunk.result = thunk.f(args...; kwargs...)
thunk.evaluated = true
return thunk.result
end
end
function reify(value)
value
end
This does allow for any arbitrary Julia expression be written, albeit in ugly & manual fashion:
b = thunk(identity)(10)
z = thunk(sleep)(10)
a = thunk(b->identity(collect(1:b))[8:end])(b)
res = thunk(()->[x-1 for x in 1:4])()
c,d,e,f = [thunk(getindex)(res,1), thunk(getindex)(res,2),
thunk(getindex)(res,3), thunk(getindex)(res,4)]
abc = thunk((a,c,d,f)->sum(add1.(a) .* [c,d,f]))(a,c,d,f)
@time @assert reify(abc) == 43
As a programmer writing this transform, I have to parse each line, and look for the existence of a symbol that is a Thunk, and then I use each of these symbols as a variable when wrapping in a function.
In theory, this should be very fast and easy to do thanks to the homoiconic nature of Julia.
"Walk an AST and find all unique symbols."
function find_symbols_in_ast(ex)
# accumulate sub-expressions for recursion
expressions = [ex]
# accumulate found symbols
symbols = []
# since no Tail Call Optimization in Julia,
# we write recursion in a while loop
while length(expressions) > 0
ex = pop!(expressions)
first, rest = _find_symbols_in_ast(ex)
if typeof(first) == Symbol
# got value, no more recursion
push!(symbols, first)
end
if ~isnothing(rest)
# recur
expressions = vcat(expressions, rest)
end
end
return unique(symbols)
end
"Helper function following `cons` and `nil` pattern from Lisp."
function _find_symbols_in_ast(ex)
if typeof(ex) == Expr
head, args = ex.head, ex.args
return nothing, args
elseif typeof(ex) == Symbol
return ex, nothing
else
return nothing, nothing
end
end
isthunk(x) = typeof(x) == Thunk
"Safely evaluate symbols that may not have an assignment."
function safe_eval_isthunk(ex)
try
return isthunk(eval(ex))
catch
return false
end
end
"Return array of symbols that are assigned to a Thunk."
function find_thunks_in_ast(ex)
symbols = find_symbols_in_ast(ex)
filter(safe_eval_isthunk, symbols)
end
Unfortunately, find_thunks_in_ast does not work if Thunks are assigned to variables in the local scope because “Julia’s eval always evaluates code [in] the current module’s scope, not your local scope.”
Any ideas how to sidestep this problem? I think this is an example where allowing eval to access variables in a local scope has massive ergonomic benefit and virtually no performance cost, and I’m not sure there are any easy alternatives…
Many thanks for your help!
