consolidate map sort on write code

Summary: Consolidate the `writeMapSorted` for mutable and immutable because they're almost identical, except for call to get items from mutable dict. In the future, if we change immutable data type to dict, then they'll fully converge.

Reviewed By: createdbysk

Differential Revision: D78528675

fbshipit-source-id: 8641570fbf030ca72fd350cf22b39e4100a3cab1

Author: ahilger

thrift/lib/python/types.cpp

@@ -1021,37 +1021,6 @@ void* setIOBuf(void* objectPtr, const folly::IOBuf& value) {
   return nullptr;
 }
 
-// This helper method for `MutableMapTypeInfo::write()` sorts the map keys and
-// writes them to the wire. It is called when the `protocolSortKeys` parameter
-// of `write()` is set to `true`.
-size_t writeMapSorted(
-    const void* context,
-    const void* object,
-    size_t (*writer)(
-        const void* context, const void* keyElem, const void* valueElem)) {
-  PyObject* dict = const_cast<PyObject*>(toPyObject(object));
-  DCHECK(PyDict_Check(dict));
-  UniquePyObjectPtr listPtr =
-      UniquePyObjectPtr{PySequence_List(PyDict_Items(dict))};
-  if (!listPtr) {
-    THRIFT_PY3_CHECK_ERROR();
-  }
-  if (PyList_Sort(listPtr.get()) == -1) {
-    THRIFT_PY3_CHECK_ERROR();
-  }
-
-  size_t written = 0;
-  const Py_ssize_t size = PyList_Size(listPtr.get());
-  for (std::uint32_t i = 0; i < size; ++i) {
-    PyObject* pair = PyList_GET_ITEM(listPtr.get(), i);
-    PyObject* key = PyTuple_GET_ITEM(pair, 0);
-    PyObject* value = PyTuple_GET_ITEM(pair, 1);
-    written += writer(context, &key, &value);
-  }
-
-  return written;
-}
-
 inline UniquePyObjectPtr primitiveCppToPython(bool value) {
   PyObject* ret = value ? Py_True : Py_False;
   Py_INCREF(ret);
@@ -1647,26 +1616,43 @@ void ImmutableMapHandler::read(
   setPyObject(objectPtr, std::move(map));
 }
 
-size_t ImmutableMapHandler::write(
+size_t writeMapSorted(
+    const void* context,
+    const void* object,
+    PyObject* (*toItems)(PyObject* dict_),
+    size_t (*writer)(
+        const void* context, const void* keyElem, const void* valueElem)) {
+  PyObject* dict = const_cast<PyObject*>(toPyObject(object));
+  PyObject* keys = toItems(dict);
+  UniquePyObjectPtr listPtr = UniquePyObjectPtr{PySequence_List(keys)};
+  if (!listPtr) {
+    THRIFT_PY3_CHECK_ERROR();
+  }
+  if (PyList_Sort(listPtr.get()) == -1) {
+    THRIFT_PY3_CHECK_ERROR();
+  }
+
+  size_t written = 0;
+  const Py_ssize_t size = PyList_Size(listPtr.get());
+  for (std::uint32_t i = 0; i < size; ++i) {
+    PyObject* pair = PyList_GET_ITEM(listPtr.get(), i);
+    PyObject* key = PyTuple_GET_ITEM(pair, 0);
+    PyObject* value = PyTuple_GET_ITEM(pair, 1);
+    written += writer(context, &key, &value);
+  }
+
+  return written;
+}
+
+size_t ImmutableMapHandler::writeUnsorted(
     const void* context,
     const void* object,
-    bool protocolSortKeys,
     size_t (*writer)(
         const void* context, const void* keyElem, const void* valueElem)) {
   size_t written = 0;
   PyObject* map = const_cast<PyObject*>(toPyObject(object));
   const Py_ssize_t size = PyTuple_GET_SIZE(map);
   UniquePyObjectPtr seq;
-  if (protocolSortKeys) {
-    seq = UniquePyObjectPtr{PySequence_List(map)};
-    if (!seq) {
-      THRIFT_PY3_CHECK_ERROR();
-    }
-    if (PyList_Sort(seq.get()) == -1) {
-      THRIFT_PY3_CHECK_ERROR();
-    }
-    map = PySequence_Tuple(seq.get());
-  }
   for (std::uint32_t i = 0; i < size; ++i) {
     PyObject* pair = PyTuple_GET_ITEM(map, i);
     PyObject* key = PyTuple_GET_ITEM(pair, 0);
@@ -1723,24 +1709,21 @@ void MutableMapHandler::read(
   setPyObject(objectPtr, std::move(dict));
 }
 
-size_t MutableMapHandler::write(
+size_t MutableMapHandler::writeUnsorted(
     const void* context,
     const void* object,
-    bool protocolSortKeys,
     size_t (*writer)(
         const void* context, const void* keyElem, const void* valueElem)) {
-  if (protocolSortKeys) {
-    return writeMapSorted(context, object, writer);
-  }
-
   PyObject* dict = const_cast<PyObject*>(toPyObject(object));
   size_t written = 0;
   PyObject* key = nullptr;
   PyObject* value = nullptr;
   Py_ssize_t pos = 0;
+
   while (PyDict_Next(dict, &pos, &key, &value)) {
     written += writer(context, &key, &value);
   }
+
   return written;
 }
 

thrift/lib/python/types.h

@@ -564,6 +564,17 @@ struct MutableSetHandler {
 
 using MutableSetTypeInfo = SetTypeInfoImpl<MutableSetHandler, false>;
 
+/**
+ * This helper method sorts map keys and write them to the wire. It may  be
+ * called for both mutable and immutable maps.
+ */
+size_t writeMapSorted(
+    const void* context,
+    const void* object,
+    PyObject* (*toItems)(PyObject* dict),
+    size_t (*writer)(
+        const void* context, const void* keyElem, const void* valueElem));
+
 template <typename Handler>
 class MapTypeInfoImpl {
  public:
@@ -588,7 +599,10 @@ class MapTypeInfoImpl {
       bool protocolSortKeys,
       size_t (*writer)(
           const void* context, const void* keyElem, const void* valueElem)) {
-    return Handler::write(context, object, protocolSortKeys, writer);
+    if (UNLIKELY(protocolSortKeys)) {
+      return writeMapSorted(context, object, Handler::toItems, writer);
+    }
+    return Handler::writeUnsorted(context, object, writer);
   }
 
   static void consumeElem(
@@ -632,6 +646,12 @@ class ImmutableMapHandler {
     return PyTuple_GET_SIZE(object);
   }
 
+  // Result is used by writeMapSorted as input to PySequence_list
+  // Since immutable map representation is tuple, this is a no-op.
+  // Since maps can't have duplicate keys, the value (second value in tuple)
+  // will never be used for comparison, only the first value of each.
+  static PyObject* toItems(PyObject* dictTuple) { return dictTuple; }
+
   static void* clear(void* object) { return setContainer(object); }
 
   static void read(
@@ -641,10 +661,9 @@ class ImmutableMapHandler {
       void (*keyReader)(const void* context, void* key),
       void (*valueReader)(const void* context, void* val));
 
-  static size_t write(
+  static size_t writeUnsorted(
       const void* context,
       const void* object,
-      bool protocolSortKeys,
       size_t (*writer)(
           const void* context, const void* keyElem, const void* valueElem));
 
@@ -663,6 +682,15 @@ class MutableMapHandler {
     return PyDict_Size(const_cast<PyObject*>(object));
   }
 
+  // Result is used by writeMapSorted as input to PySequence_List
+  static PyObject* toItems(PyObject* dict) {
+    PyObject* keys = PyDict_Items(dict);
+    if (keys == nullptr) {
+      THRIFT_PY3_CHECK_ERROR();
+    }
+    return keys;
+  }
+
   static void* clear(void* objectPtr) { return setMutableMap(objectPtr); }
 
   /**
@@ -702,10 +730,9 @@ class MutableMapHandler {
    *
    * @return Total number of bytes written.
    */
-  static size_t write(
+  static size_t writeUnsorted(
       const void* context,
       const void* object,
-      bool protocolSortKeys,
       size_t (*writer)(
           const void* context, const void* keyElem, const void* valueElem));