From 8014163c25e3b4c8aa61f04b4412b656f54ba495 Mon Sep 17 00:00:00 2001 From: Eugene Blikh Date: Fri, 15 May 2026 13:59:11 +0300 Subject: [PATCH] codec: precompute per-field writers for singular scalar/enum/message (~+50% small-msg encode) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pb.finalize_message now calls codec.compile_writers(desc), which attaches `f._writer` to each field whose shape we can specialize: singular scalar, singular enum, singular message — i.e. not maps, not repeated, not oneof. Each writer is a monomorphic closure that knows its tag bytes, encoder function, and default predicate. The encode_message hot loop calls writer(data, out) per field and only falls through to encode_field for shapes we haven't specialized. This eliminates the per-field `encode_field` dispatch chain (kind/proto_type branch + `is_default_scalar` call), which was the source of the remaining trace bridges in runtime mode (codec.lua:41 and codec.lua:110 in `make jit-trace` output). Effect (bench/bench.lua, hello.Person): runtime/10B encode: 5.0 → 8.1 MB/s (+62%) runtime/100B encode: 48.2 → 75.2 MB/s (+56%) runtime/1KB encode: 49.5 → 56.1 MB/s (+13%) runtime/10KB+ encode: unchanged (dominated by repeated-field iteration — not yet specialized) full mode encode: +10% across small sizes (writer closures also help when codegen calls back into the runtime) decode: unchanged alloc/op: unchanged (bench-compare clean) bridges across 10 jit-trace runs: 7 → 4 --- runtime/pb/codec.lua | 118 ++++++++++++++++++++++++++++++++++++++++++- runtime/pb/init.lua | 6 +++ 2 files changed, 123 insertions(+), 1 deletion(-) diff --git a/runtime/pb/codec.lua b/runtime/pb/codec.lua index e8c07eb304f68327bbdc5f104e7cc7831eee241f..5453043e7eaf04a858292257377428c0c695bfd3 100644 --- a/runtime/pb/codec.lua +++ b/runtime/pb/codec.lua @@ -211,6 +211,119 @@ local function encode_field(field, value, out, force) end end +-- --------------------------------------------------------------------------- +-- Per-field "writer" specialization +-- +-- For shapes we can specialize (singular scalar/enum/message without +-- presence semantics on the hot path), build a monomorphic closure at +-- finalize-time that knows its tag bytes, encoder, and default predicate. +-- The encode_message loop calls writers in order. Eliminates per-field +-- kind/proto_type dispatch and the inlined `is_default_scalar` predicate +-- — both of which create side-trace-can't-stitch-back-to-parent bridges +-- in LuaJIT 2.1 when the trace recorder sees mixed shapes across fields. +-- +-- Shapes that fall through to encode_field (no writer set): +-- - map fields (need `pairs()` over user data; can't be helped) +-- - repeated fields (TODO — common shape, worth specializing) +-- - fields inside a oneof (active-branch dispatch happens in encode_message) +-- --------------------------------------------------------------------------- + +local function build_writer(f) + -- Maps and oneof branches keep going through encode_field. + if f.kind == 'map' or f.oneof then return nil end + -- Repeated fields fall through for now. + if f.repeated then return nil end + + local fname = f.name + local kind = f.kind + local optional = f.optional + + if kind == 'scalar' then + local handler = scalar[f.proto_type] + if not handler then return nil end + local tag_bytes = wire.encode_tag(f.id, handler.wire) + local encode_value = handler.encode + local proto_type = f.proto_type + + if optional then + return function(data, out) + local v = data[fname] + if v == nil then return end + local n = #out + out[n + 1] = tag_bytes + out[n + 2] = encode_value(v) + end + end + + if proto_type == 'string' or proto_type == 'bytes' then + return function(data, out) + local v = data[fname] + if v == nil or v == '' then return end + local n = #out + out[n + 1] = tag_bytes + out[n + 2] = encode_value(v) + end + end + if proto_type == 'bool' then + return function(data, out) + local v = data[fname] + if v == nil or v == false then return end + local n = #out + out[n + 1] = tag_bytes + out[n + 2] = encode_value(v) + end + end + -- Numeric scalar (int32/uint32/int64/uint64/sint32/sint64/ + -- fixed32/sfixed32/fixed64/sfixed64/float/double). For cdata + -- 64-bit values, `v == 0` is the LuaJIT-canonical default + -- check — it works across UINT64 / INT64 because cdata-to- + -- number comparison normalizes via int64. + return function(data, out) + local v = data[fname] + if v == nil or v == 0 then return end + local n = #out + out[n + 1] = tag_bytes + out[n + 2] = encode_value(v) + end + end + + if kind == 'enum' then + local enum_desc = f.enum + local tag_bytes = wire.encode_tag(f.id, wire.WIRE_VARINT) + return function(data, out) + local v = data[fname] + if v == nil then return end + local n_enum = encode_enum_value(enum_desc, v) + if not optional and n_enum == 0 then return end + local n = #out + out[n + 1] = tag_bytes + out[n + 2] = wire.encode_varint(n_enum) + end + end + + if kind == 'message' then + local sub_desc = f.message + local tag_bytes = wire.encode_tag(f.id, wire.WIRE_LEN) + return function(data, out) + local v = data[fname] + if v == nil then return end + local n = #out + out[n + 1] = tag_bytes + out[n + 2] = wire.encode_len(encode_msg(sub_desc, v)) + end + end + + return nil +end + +-- compile_writers attaches `f._writer` to each field where the shape is +-- specialized. Called from pb.finalize_message after the oneof flatten. +function M.compile_writers(desc) + for _, f in ipairs(desc.fields) do + f._writer = build_writer(f) + end +end + encode_message = function(desc, data) if type(data) ~= 'table' then error(("expected table for message %s, got %s"):format(desc.name, type(data)), 0) @@ -238,7 +351,10 @@ encode_message = function(desc, data) for i = 1, #fields do local f = fields[i] - if f.oneof then + local writer = f._writer + if writer ~= nil then + writer(data, out) + elseif f.oneof then if active and active[f.oneof] == f.name then encode_field(f, data[f.name], out, true) -- force: emit even defaults end diff --git a/runtime/pb/init.lua b/runtime/pb/init.lua index 86ba413ec5c4ca310c211d60ceb32199bda6c465..ce4a2a9435cbeb8d51f720a0241a055c6c154523 100644 --- a/runtime/pb/init.lua +++ b/runtime/pb/init.lua @@ -111,6 +111,12 @@ return { end desc.oneofs_list = list end + -- Attach a per-field monomorphic writer function for the shapes + -- the codec can specialize (singular scalar/enum/message). The + -- encode_message hot loop calls writer(data, out) per field and + -- avoids the runtime kind/proto_type dispatch chain inside + -- encode_field. + codec.compile_writers(desc) return desc end, }