From 46045da1d9330db68d0e8210ff0ba4ae2c0885db Mon Sep 17 00:00:00 2001 From: Eugene Blikh Date: Fri, 15 May 2026 14:07:56 +0300 Subject: [PATCH] =?UTF-8?q?codec:=20precompute=20per-field=20readers=20(ru?= =?UTF-8?q?ntime=20decode=20+10=E2=80=9317%)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirror of compile_writers on the decode side: pb.finalize_message now also calls codec.compile_readers(desc), attaching `f._reader` to each field whose shape can be specialized — singular scalar/enum/message and repeated scalar/enum/message (packed and unpacked). Each reader has signature (buf, pos, wt, result) -> new_pos and bakes in the field name, decode function, packed-detection, list bookkeeping, nested-message merge rules, and oneof sibling clearing. The decode_message hot loop becomes: id, wt, pos = decode_tag(buf, pos) f = fbi[id] if f and f._reader then pos = f._reader(buf, pos, wt, result) else /* existing per-kind dispatch — map fields only */ end Forward declaration `local decode_msg` so reader closures captured at finalize-time can refer to it; the later `decode_msg = function...` fills the upvalue. Effect (bench/bench.lua, hello.Person): runtime/100B decode: 246 → 272 MB/s (+11%) runtime/1KB decode: 122 → 144 MB/s (+17%) runtime/10KB decode: 176 → 200 MB/s (+13%) runtime/100KB decode: 180 → 210 MB/s (+17%) full mode decode: ~flat (already maximally inlined by codegen) encode: unchanged alloc/op: unchanged (bench-compare clean) bridges across 10 jit-trace runs: 3 → 2 Runtime mode decode is now within ~10–15% of full mode across all sizes, vs 24–28% gap before this commit. --- runtime/pb/codec.lua | 191 ++++++++++++++++++++++++++++++++++++++++++- runtime/pb/init.lua | 13 ++- 2 files changed, 198 insertions(+), 6 deletions(-) diff --git a/runtime/pb/codec.lua b/runtime/pb/codec.lua index 9b8c796fb24a5d3624a7c1df470161fff21eaa78..f254c6e4211810cf3fb7e74a3ec414600cd423c7 100644 --- a/runtime/pb/codec.lua +++ b/runtime/pb/codec.lua @@ -59,8 +59,11 @@ end -- Encode -- --------------------------------------------------------------------------- --- Forward declaration so encode_field can recurse via encode_message. +-- Forward declarations: encode_field/build_writer/build_reader need to +-- close over encode_message and decode_msg before those are assigned +-- further down (writers/readers built by compile_* at finalize-time). local encode_message +local decode_msg local function encode_enum_value(enum_desc, v) if type(v) == 'number' then return v end @@ -416,6 +419,181 @@ function M.compile_writers(desc) end end +-- --------------------------------------------------------------------------- +-- Per-field "reader" specialization, mirror of writers above. Each reader +-- has signature `(buf, pos, wt, result) -> new_pos` and bakes in the field +-- name, decode function, packed-detection, repeated-list bookkeeping, +-- nested-message merge rules, and oneof sibling clearing — so the +-- decode_message hot loop is just `pos = reader(buf, pos, wt, result)`. +-- +-- Map fields fall through to the existing in-loop dispatch (they decode +-- one entry per wire-encounter, which is awkward to express as a closure +-- without churning the result table); same as on the encoder side. +-- --------------------------------------------------------------------------- + +local function build_repeated_reader(f) + local fname = f.name + local kind = f.kind + local siblings = f.oneof_siblings -- nil if not in a oneof + + if kind == 'scalar' then + local handler = scalar[f.proto_type] + if not handler then return nil end + local decode_value = handler.decode + local packable = handler.packable and (handler.wire ~= wire.WIRE_LEN) + local WIRE_LEN = wire.WIRE_LEN + local decode_len_fn = wire.decode_len + return function(buf, pos, wt, result) + local list = result[fname] + if list == nil then list = {}; result[fname] = list end + if packable and wt == WIRE_LEN then + local payload, np = decode_len_fn(buf, pos) + local p, lim = 1, #payload + local n = #list + while p <= lim do + local v, np2 = decode_value(payload, p) + n = n + 1; list[n] = v + p = np2 + end + return np + end + local v, np = decode_value(buf, pos) + list[#list + 1] = v + return np + end + end + + if kind == 'enum' then + local WIRE_LEN = wire.WIRE_LEN + local decode_len_fn = wire.decode_len + local decode_varint_fn = wire.decode_varint + return function(buf, pos, wt, result) + local list = result[fname] + if list == nil then list = {}; result[fname] = list end + if wt == WIRE_LEN then + local payload, np = decode_len_fn(buf, pos) + local p, lim = 1, #payload + local n = #list + while p <= lim do + local u, np2 = decode_varint_fn(payload, p) + n = n + 1; list[n] = tonumber(u) + p = np2 + end + return np + end + local u, np = decode_varint_fn(buf, pos) + list[#list + 1] = tonumber(u) + return np + end + end + + if kind == 'message' then + local sub_desc = f.message + local decode_len_fn = wire.decode_len + return function(buf, pos, wt, result) + local list = result[fname] + if list == nil then list = {}; result[fname] = list end + local payload, np = decode_len_fn(buf, pos) + list[#list + 1] = decode_msg(sub_desc, payload) + return np + end + end + + return nil +end + +local function build_reader(f) + -- Maps stay on the in-loop dispatch path. + if f.kind == 'map' then return nil end + + local fname = f.name + local kind = f.kind + local siblings = f.oneof_siblings -- nil if not in a oneof + + if f.repeated then return build_repeated_reader(f) end + + if kind == 'scalar' then + local handler = scalar[f.proto_type] + if not handler then return nil end + local decode_value = handler.decode + if siblings then + return function(buf, pos, wt, result) + local v, np = decode_value(buf, pos) + result[fname] = v + for i = 1, #siblings do result[siblings[i]] = nil end + return np + end + end + return function(buf, pos, wt, result) + local v, np = decode_value(buf, pos) + result[fname] = v + return np + end + end + + if kind == 'enum' then + local decode_varint_fn = wire.decode_varint + if siblings then + return function(buf, pos, wt, result) + local u, np = decode_varint_fn(buf, pos) + result[fname] = tonumber(u) + for i = 1, #siblings do result[siblings[i]] = nil end + return np + end + end + return function(buf, pos, wt, result) + local u, np = decode_varint_fn(buf, pos) + result[fname] = tonumber(u) + return np + end + end + + if kind == 'message' then + local sub_desc = f.message + local decode_len_fn = wire.decode_len + local is_oneof = f.oneof ~= nil + local has_custom_decode = sub_desc.decode ~= nil + -- Per-spec: repeated singular-message wire entries merge into + -- existing, unless this is a oneof branch (exclusive) or the + -- nested message has a custom decode (e.g. WKT — no merge). + local always_replace = is_oneof or has_custom_decode + if always_replace then + if siblings then + return function(buf, pos, wt, result) + local payload, np = decode_len_fn(buf, pos) + result[fname] = decode_msg(sub_desc, payload) + for i = 1, #siblings do result[siblings[i]] = nil end + return np + end + end + return function(buf, pos, wt, result) + local payload, np = decode_len_fn(buf, pos) + result[fname] = decode_msg(sub_desc, payload) + return np + end + end + return function(buf, pos, wt, result) + local payload, np = decode_len_fn(buf, pos) + local decoded = decode_msg(sub_desc, payload) + local prev = result[fname] + if prev == nil then + result[fname] = decoded + else + for k, v in pairs(decoded) do prev[k] = v end + end + return np + end + end + + return nil +end + +function M.compile_readers(desc) + for _, f in ipairs(desc.fields) do + f._reader = build_reader(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) @@ -469,7 +647,9 @@ M.encode = encode_message local decode_message -- decode_msg dispatches to a descriptor's custom decode (WKT) when present. -local function decode_msg(desc, buf) +-- Assigned to the forward declaration near the top so reader closures +-- built by compile_readers can capture it. +decode_msg = function(desc, buf) if desc.decode then return desc.decode(buf) end return decode_message(desc, buf) end @@ -539,6 +719,12 @@ decode_message = function(desc, buf) if unknown == nil then unknown = {} end unknown[#unknown + 1] = buf:sub(tag_start, pos - 1) else + local reader = f._reader + if reader ~= nil then + pos = reader(buf, pos, wt, result) + else + -- Fallthrough for shapes without a specialized reader + -- (currently only map fields). local kind = f.kind if kind == 'map' then local map_t = result[f.name] @@ -629,6 +815,7 @@ decode_message = function(desc, buf) for _, s in ipairs(f.oneof_siblings) do result[s] = nil end end end + end -- end of `if reader ~= nil ... else ... end` end end if unknown ~= nil then result._unknown_fields = table.concat(unknown) end diff --git a/runtime/pb/init.lua b/runtime/pb/init.lua index ce4a2a9435cbeb8d51f720a0241a055c6c154523..7c1e38b0c2e4f6c77a87e1acff7d718ed2110c10 100644 --- a/runtime/pb/init.lua +++ b/runtime/pb/init.lua @@ -112,11 +112,16 @@ return { 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. + -- the codec can specialize (singular scalar/enum/message and + -- repeated 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) + -- Same idea for the decoder side: per-field readers handle + -- typed value extraction, list bookkeeping, message-merge + -- rules, and oneof sibling clearing. Maps fall through to + -- the existing in-loop dispatch. + codec.compile_readers(desc) return desc end, }