forked from JuliaLang/julia
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathstaged.jl
61 lines (57 loc) · 1.68 KB
/
staged.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
# This file is a part of Julia. License is MIT: http://julialang.org/license
function add_method(gf, an, at, body)
argexs = [Expr(symbol("::"), an[i], at[i]) for i=1:length(an)]
def = quote
let __F__=($gf)
function __F__($(argexs...))
$body
end
end
end
eval(def)
end
macro staged(fdef)
if !isa(fdef,Expr) || !is(fdef.head,:function)
error("@staged: expected method definition")
end
fname = fdef.args[1].args[1]
argspec = fdef.args[1].args[2:end]
argnames = map(x->(isa(x,Expr) ? x.args[1] : x), argspec)
qargnames = map(x->Expr(:quote,x), argnames)
fbody = fdef.args[2]
@gensym gengf argtypes expander genbody
quote
let ($gengf)
global ($fname) # should be "outer"
local ($expander)
function ($expander)($(argnames...))
$fbody
end
($gengf)() = 0 # should be initially empty GF
function ($fname)($(argspec...))
($argtypes) = typeof(tuple($(argnames...)))
if !method_exists($gengf, $argtypes)
($genbody) = apply(($expander), ($argtypes))
add_method($gengf, Any[$(qargnames...)],
$argtypes, $genbody)
end
return ($gengf)($(argnames...))
end
end
end
end
# example
@staged function nloops(dims::Tuple)
names = map(x->gensym(), dims)
ex = quote
println([$(names...)])
end
for i = 1:length(dims)
ex = quote
for $(names[i]) in dims[$i]
$ex
end
end
end
ex
end