JIT-friendly version of unpack()

[Roman Tsisyk]

Hi!

I found one interesting pattern often used by Lua developers:

local t = { some Lua table with 3 elements }
local v1, v2, v3 = unpack(t)

A trick with unpack() is very nice, but it breaks all JIT traces
because unpack() is not compiled.
For plain Lua tables the example above can be easily rewritten:

local v1, v2, v3 = t[1], t[2], t[3]

Now let's imagine if we have an iterator without random access instead
of a plain Lua table:

local v1, v2, v3 = obj:iterate_and_unpack_all_elements()

where iterate_and_unpack_all_elements() should be logically equivalent to:

local t = {}
for state, v in ipairs(obj) do
table.insert(t, v)
end
return unpack(t)

Is it possible to write JIT-friendly version of such function?

I experimented with unpack() itself:

1. Tail recursion to build multi-return stack - doesn't compile

local function unpack2_r(gen, init, state)
local v
state, v = gen(init, state)
if state == nil then
return
end
return v, unpack2_r(gen, init, state)
end

local function unpack2(t)
return unpack2_r(ipairs(t))
end

2. Specialized versions for particular sizes - do compile, but looks ugly

local function unpack3(t)
if #t == 1 then
return t[1]
else if #t == 2 then
return t[1], t[2]
else if #t == 3 then
return t[1], t[2], t[3]
else
return unpack(t) -- goodbye, JIT
end
end

Any ideas?
Thanks!

[Craig Barnes]

Post by Roman Tsisyk
2. Specialized versions for particular sizes - do compile, but looks ugly
local function unpack3(t)
if #t == 1 then
return t[1]
else if #t == 2 then
return t[1], t[2]
else if #t == 3 then
return t[1], t[2], t[3]
else
return unpack(t) -- goodbye, JIT
end
end

It doesn't have to be that ugly. The following is more or less equivalent:

local function unpack3(t)
if #t <= 3 then
return t[1], t[2], t[3]
else
return unpack(t)
end
end

You could also make an unpack5 without any extra redundant branches or
length operators and only a little extra code.

[Stefano]

Please notice that if it's meant as a replacement of the standard
library unpack(), then the signature should be:

unpack (list [, i [, j]])

Anyway, you could use template preprocessors for a limited number of
cases, not ugly:

http://colberg.org/lua-templet/

Stefano

Post by Roman Tsisyk

Post by Roman Tsisyk
2. Specialized versions for particular sizes - do compile, but looks ugly
local function unpack3(t)
if #t == 1 then
return t[1]
else if #t == 2 then
return t[1], t[2]
else if #t == 3 then
return t[1], t[2], t[3]
else
return unpack(t) -- goodbye, JIT
end
end

local function unpack3(t)
if #t <= 3 then
return t[1], t[2], t[3]
else
return unpack(t)
end
end
You could also make an unpack5 without any extra redundant branches or
length operators and only a little extra code.

[Roman Tsisyk]

Post by Roman Tsisyk
local function unpack3(t)
if #t <= 3 then
return t[1], t[2], t[3]
else
return unpack(t)
end
end

Yeah, it is clear for me.
I posted that example just to describe the problem with multi-return and JIT.
Let's look at a more complex example:

local state, v1, v2, v3, ...
state, v1 = obj:next()
if not state then return end
state, v2 = obj:next(state)
if not state then return v1 end
state, v3 = obj:next(state)
if not state then return v1, v2 end
...
state, vn = obj:next(state)
if not state then return v1, v2, ..., vn end
...

How to implement this pattern in Lua without breaking JIT?
Lua/C API can do that using sequence of lua_pushxxx() calls.

[lex pops]

Post by Roman Tsisyk

Post by Roman Tsisyk
local function unpack3(t)
if #t <= 3 then
return t[1], t[2], t[3]
else
return unpack(t)
end
end

Yeah, it is clear for me.
I posted that example just to describe the problem with multi-return and JIT.
local state, v1, v2, v3, ...
state, v1 = obj:next()
if not state then return end
state, v2 = obj:next(state)
if not state then return v1 end
state, v3 = obj:next(state)
if not state then return v1, v2 end
...
state, vn = obj:next(state)
if not state then return v1, v2, ..., vn end
...
How to implement this pattern in Lua without breaking JIT?

Firstly, what you are doing is unidiomatic - using multiple return
values for a homogeneous list of items. Often multiple return makes
sense for heterogeneous tuples - e.g. (status, result) or (host, port)
etc. What's your exact use case for trying to unpack an iterator like
that?

Secondly, note that it shouldn't matter how many items are returned by
the iterator - instead of exhausting it you could just check how many
items are expected. I don't see how you could have a generic jittable
unpack do what you want. You need specialized versions based on the
number of items. But if you're doing multiple assign, you already know
the number of items. It's just redundant info you can insert into the
upack call.

I think the trick I presented earlier would still work for the iterator case.

[lex pops]

Post by Roman Tsisyk
Hi!
local t = { some Lua table with 3 elements }
local v1, v2, v3 = unpack(t)
A trick with unpack() is very nice, but it breaks all JIT traces
because unpack() is not compiled.

<snip>

Post by Roman Tsisyk
2. Specialized versions for particular sizes - do compile, but looks ugly
local function unpack3(t)
if #t == 1 then
return t[1]
else if #t == 2 then
return t[1], t[2]
else if #t == 3 then
return t[1], t[2], t[3]
else
return unpack(t) -- goodbye, JIT
end
end
Any ideas?

Another untested idea:

local a, b,c = unpackn[3](x)
local a, b, c, d = unpackn[4](x)

unpackn's metatable would generate+eval the correct lua code and then
insert it at unpackn[n]. The first invocation wouldn't be jittable of
course. Subsequent invocations are one table lookup and one function
call to fully jittable code.