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 @@ error("unknown field kind " .. tostring(kind), 0) 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 @@ end 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 @@ end 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, }