2 files changed, 198 insertions(+), 6 deletions(-)

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

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