[Tarantool-patches] [PATCH luajit] Consider slots used by upvalues in use-def analysis.

[Sergey Bronnikov]

Hi, Sergey

thanks for the patch! LGTM with a minor comment

On 2/7/24 17:51, Sergey Kaplun wrote:
> From: Mike Pall <mike>
>
> Reported by XmiliaH.
>
> (cherry picked from commit 3a654999c6f00de4cb9e61232d23579442e544a0)
>
> `snap_usedef()` analysis doesn't check slots for child upvalues of the
> currentlly recorded function in use-def analysis. Hence, such slots may
typo: currentlly
> be considered unused and not stored in the snapshot. So on snapshot
> restoration, values from these slots may be lost.
>
> This patch adds a marking for all such local upvalues.
>
> Sergey Kaplun:
> * added the description and the test for the problem
>
> Part of tarantool/tarantool#9595
> ---
>
> Branch: https://github.com/tarantool/luajit/tree/skaplun/lj-737-snap-usedef-upvalues
> Tarantool PR: https://github.com/tarantool/tarantool/pull/9662
> Related issues:
> * https://github.com/tarantool/tarantool/issues/9595
> * https://github.com/LuaJIT/LuaJIT/issues/737
>
>   src/lj_snap.c                                 | 35 +++++++++++-
>   .../lj-737-snap-use-def-upvalues.test.lua     | 55 +++++++++++++++++++
>   2 files changed, 87 insertions(+), 3 deletions(-)
>   create mode 100644 test/tarantool-tests/lj-737-snap-use-def-upvalues.test.lua
>
> diff --git a/src/lj_snap.c b/src/lj_snap.c
> index 26352080..1ac9296a 100644
> --- a/src/lj_snap.c
> +++ b/src/lj_snap.c
> @@ -302,6 +302,31 @@ static BCReg snap_usedef(jit_State *J, uint8_t *udf,
>     return 0;  /* unreachable */
>   }
>   
> +/* Mark slots used by upvalues of child prototypes as used. */
> +void snap_useuv(GCproto *pt, uint8_t *udf)
> +{
> +  /* This is a coarse check, because it's difficult to correlate the lifetime
> +  ** of slots and closures. But the number of false positives is quite low.
> +  ** A false positive may cause a slot not to be purged, which is just
> +  ** a missed optimization.
> +  */
> +  if ((pt->flags & PROTO_CHILD)) {
> +    ptrdiff_t i, j, n = pt->sizekgc;
> +    GCRef *kr = mref(pt->k, GCRef) - 1;
> +    for (i = 0; i < n; i++, kr--) {
> +      GCobj *o = gcref(*kr);
> +      if (o->gch.gct == ~LJ_TPROTO) {
> +	for (j = 0; j < gco2pt(o)->sizeuv; j++) {
> +	  uint32_t v = proto_uv(gco2pt(o))[j];
> +	  if ((v & PROTO_UV_LOCAL)) {
> +	    udf[(v & 0xff)] = 0;
> +	  }
> +	}
> +      }
> +    }
> +  }
> +}
> +
>   /* Purge dead slots before the next snapshot. */
>   void lj_snap_purge(jit_State *J)
>   {
> @@ -310,9 +335,12 @@ void lj_snap_purge(jit_State *J)
>     if (bc_op(*J->pc) == BC_FUNCV && maxslot > J->pt->numparams)
>       maxslot = J->pt->numparams;
>     s = snap_usedef(J, udf, J->pc, maxslot);
> -  for (; s < maxslot; s++)
> -    if (udf[s] != 0)
> -      J->base[s] = 0;  /* Purge dead slots. */
> +  if (s < maxslot) {
> +    snap_useuv(J->pt, udf);
> +    for (; s < maxslot; s++)
> +      if (udf[s] != 0)
> +	J->base[s] = 0;  /* Purge dead slots. */
> +  }
>   }
>   
>   /* Shrink last snapshot. */
> @@ -325,6 +353,7 @@ void lj_snap_shrink(jit_State *J)
>     BCReg maxslot = J->maxslot;
>     BCReg baseslot = J->baseslot;
>     BCReg minslot = snap_usedef(J, udf, snap_pc(&map[nent]), maxslot);
> +  if (minslot < maxslot) snap_useuv(J->pt, udf);
>     maxslot += baseslot;
>     minslot += baseslot;
>     snap->nslots = (uint8_t)maxslot;
> diff --git a/test/tarantool-tests/lj-737-snap-use-def-upvalues.test.lua b/test/tarantool-tests/lj-737-snap-use-def-upvalues.test.lua
> new file mode 100644
> index 00000000..8535f9f6
> --- /dev/null
> +++ b/test/tarantool-tests/lj-737-snap-use-def-upvalues.test.lua
> @@ -0,0 +1,55 @@
> +local tap = require('tap')
> +
> +-- Test file to demonstrate LuaJIT misbehaviour in use-def
> +-- snapshot analysis for local upvalues.
> +-- See also https://github.com/LuaJIT/LuaJIT/issues/737.
> +
> +local test = tap.test('lj-737-snap-use-def-upvalues'):skipcond({
> +  ['Test requires JIT enabled'] = not jit.status(),
> +})
> +
> +test:plan(1)
> +
> +-- XXX: simplify `jit.dump()` output.
> +local fmod = math.fmod
> +
> +local EXPECTED = 'expected'
> +
> +jit.opt.start('hotloop=1')
> +
> +local function wrapped_trace(create_closure)
> +  local local_upvalue, closure
> +  if create_closure then
> +    closure = function() return local_upvalue end
> +  end
> +  for i = 1, 4 do
> +    -- On the second iteration, the trace is recorded.
> +    if i == 2 then
> +      -- Before the patch, this slot was considered unused by
> +      -- use-def analysis in the `snap_usedef()` since there are
> +      -- no open unpvalues for `closure()` on recording (1st call).
> +      local_upvalue = EXPECTED
> +      -- luacheck: ignore
> +      -- Emit an additional snapshot after setting the
> +      -- upvalue.
> +      if i == 0 then end
> +      -- Stitching ends the trace here.
> +      fmod(1,1)
> +      return closure
> +    end
> +  end
> +end
> +
> +-- Compile the trace.
> +local func_with_uv = wrapped_trace(false)
> +assert(func_with_uv == nil, 'no function is returned on the first call')
> +
> +-- Now run this trace when `closure()` is defined and has an open
> +-- local upvalue.
> +func_with_uv = wrapped_trace(true)
> +assert(type(func_with_uv) == 'function',
> +       'function is returned after the second call')
> +
> +test:is(func_with_uv(), EXPECTED, 'correct result of the closure call')
> +
> +test:done(true)