From b83a1fece12fbd0f07aa18e1b7198796e08c93bf Mon Sep 17 00:00:00 2001
From: Eric Salo <salo@google.com>
Date: Wed, 6 Dec 2023 07:24:31 -0800
Subject: [PATCH] upb: implement upb_MiniTable_MapKey() and
 upb_MiniTable_MapValue()

PiperOrigin-RevId: 588418734
---
 upb/mini_table/internal/message.h | 18 ++++++++++++++++++
 upb/mini_table/message.h          | 19 ++++++++++++-------
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/upb/mini_table/internal/message.h b/upb/mini_table/internal/message.h
index ff4e2b5f68..caa8f87327 100644
--- a/upb/mini_table/internal/message.h
+++ b/upb/mini_table/internal/message.h
@@ -112,6 +112,24 @@ UPB_INLINE const struct upb_MiniTableEnum* UPB_PRIVATE(
       m->UPB_PRIVATE(subs)[f->UPB_PRIVATE(submsg_index)]);
 }
 
+UPB_INLINE const struct upb_MiniTableField* UPB_PRIVATE(_upb_MiniTable_MapKey)(
+    const struct upb_MiniTable* m) {
+  UPB_ASSERT(UPB_PRIVATE(_upb_MiniTable_FieldCount)(m) == 2);
+  const struct upb_MiniTableField* f =
+      UPB_PRIVATE(_upb_MiniTable_GetFieldByIndex)(m, 0);
+  UPB_ASSERT(UPB_PRIVATE(_upb_MiniTableField_Number)(f) == 1);
+  return f;
+}
+
+UPB_INLINE const struct upb_MiniTableField* UPB_PRIVATE(
+    _upb_MiniTable_MapValue)(const struct upb_MiniTable* m) {
+  UPB_ASSERT(UPB_PRIVATE(_upb_MiniTable_FieldCount)(m) == 2);
+  const struct upb_MiniTableField* f =
+      UPB_PRIVATE(_upb_MiniTable_GetFieldByIndex)(m, 1);
+  UPB_ASSERT(UPB_PRIVATE(_upb_MiniTableField_Number)(f) == 2);
+  return f;
+}
+
 UPB_INLINE bool UPB_PRIVATE(_upb_MiniTable_MessageFieldIsLinked)(
     const struct upb_MiniTable* m, const struct upb_MiniTableField* f) {
   return UPB_PRIVATE(_upb_MiniTable_GetSubMessageTable)(m, f) != NULL;
diff --git a/upb/mini_table/message.h b/upb/mini_table/message.h
index d56c657c9b..89818fe10f 100644
--- a/upb/mini_table/message.h
+++ b/upb/mini_table/message.h
@@ -46,6 +46,18 @@ UPB_API_INLINE const upb_MiniTableEnum* upb_MiniTable_GetSubEnumTable(
   return UPB_PRIVATE(_upb_MiniTable_GetSubEnumTable)(m, f);
 }
 
+// Returns the MiniTableField for the key of a map.
+UPB_API_INLINE const upb_MiniTableField* upb_MiniTable_MapKey(
+    const upb_MiniTable* m) {
+  return UPB_PRIVATE(_upb_MiniTable_MapKey)(m);
+}
+
+// Returns the MiniTableField for the value of a map.
+UPB_API_INLINE const upb_MiniTableField* upb_MiniTable_MapValue(
+    const upb_MiniTable* m) {
+  return UPB_PRIVATE(_upb_MiniTable_MapValue)(m);
+}
+
 // Returns true if this MiniTable field is linked to a MiniTable for the
 // sub-message.
 UPB_API_INLINE bool upb_MiniTable_MessageFieldIsLinked(
@@ -53,13 +65,6 @@ UPB_API_INLINE bool upb_MiniTable_MessageFieldIsLinked(
   return UPB_PRIVATE(_upb_MiniTable_MessageFieldIsLinked)(m, f);
 }
 
-// TODO: Implement convenience getters for map entries:
-//
-// upb_MiniTable_GetMapKey()
-// upb_MiniTable_GetMapValue()
-// These could also assert that this is indeed a map entry (well, as best we
-// can. We can assert that there are two fields with field numbers 1 and 2).
-
 // If this field is in a oneof, returns the first field in the oneof.
 //
 // Otherwise returns NULL.