2 files changed, 123 insertions(+), 1 deletions(-)

M runtime/pb/codec.lua
M runtime/pb/init.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

@@ 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,
 }