[Picojson] Let the key of objects in json be ordered by default

Previously picojson define `object` as an alias of `std::unordered_map`. That means when parsing json, the order of keys in objects are uncertain and dependent on implementation. This makes it inconvenient for certain applications, e.g. in LLM generation output, we wish the order of keys the same as the order in the json file.

This PR implements a ordered hashmap `ordered_hashmap` that 1) maintains the order in which the elements are inserted, and 2) have the same interface as `std::unordered_map`. Picojson will define object as an alias of `ordered_hashmap`, so the order of the input json is maintained when parsing.

Macro `PICOJSON_USE_ORDERED_OBJECT` controls whether object uses the ordered version or the unordered version. It is set by default.

Author: Ubospica

3rdparty/picojson/picojson.h

@@ -26,12 +26,21 @@
  * POSSIBILITY OF SUCH DAMAGE.
  */
 #pragma once
+
 #ifndef PICOJSON_USE_INT64
 #define PICOJSON_USE_INT64
 #define __STDC_FORMAT_MACROS 1
 #endif
 
+// If PICOJSON_USE_ORDERED_OBJECT is set, picojson uses ordered dict for objects.
+// When iterating the object, the order of keys is the order they appear in the json string.
+// This macro is set by default.
+#ifndef PICOJSON_USE_ORDERED_OBJECT
+#define PICOJSON_USE_ORDERED_OBJECT 1
+#endif
+
 #include <algorithm>
+#include <cassert>
 #include <cstddef>
 #include <cstdio>
 #include <cstdlib>
@@ -137,10 +146,177 @@ enum { INDENT_WIDTH = 2 };
 
 struct null {};
 
+// The ordered version of hashmap. When iterating through the hashmap,
+// the elements maintain the order in which they were inserted.
+// Its API is the same as std::unordered_map.
+template <typename Key, typename T>
+class ordered_hashmap : private std::unordered_map<Key, T> {
+ public:
+  using value_type = std::pair<const Key, T>;
+
+ private:
+  template <bool IsConst>
+  struct iterator_base {
+   public:
+    using pointer = std::conditional_t<IsConst, const value_type*, value_type*>;
+    using reference = std::conditional_t<IsConst, const value_type&, value_type&>;
+
+    iterator_base(const iterator_base<IsConst>& other)
+        : order_index(other.order_index), obj_ptr(other.obj_ptr) {}
+
+    template <bool _IsConst = IsConst, typename = std::enable_if_t<_IsConst>>
+    iterator_base(const iterator_base<false>& other)
+        : order_index(other.order_index), obj_ptr(other.obj_ptr) {
+      static_assert(_IsConst, "This constructor should only be used for const iterators.");
+    }
+
+    iterator_base<IsConst>& operator++() {
+      ++order_index;
+      return *this;
+    }
+    iterator_base<IsConst> operator++(int) {
+      auto tmp = *this;
+      ++order_index;
+      return tmp;
+    }
+    pointer operator->() const {
+      assert(order_index >= 0 && order_index < obj_ptr->order.size());
+      return obj_ptr->std::unordered_map<Key, T>::find(obj_ptr->order[order_index]).operator->();
+    }
+    reference operator*() const { return *operator->(); }
+    bool operator==(const iterator_base<IsConst>& other) const {
+      return order_index == other.order_index && obj_ptr == other.obj_ptr;
+    }
+    bool operator!=(const iterator_base<IsConst>& other) const { return !(*this == other); }
+
+    friend class ordered_hashmap;
+
+   private:
+    using obj_pointer_type = std::conditional_t<IsConst, const ordered_hashmap*, ordered_hashmap*>;
+    iterator_base(int order_index, obj_pointer_type obj_ptr)
+        : order_index(order_index), obj_ptr(obj_ptr) {}
+
+    int order_index;
+    obj_pointer_type obj_ptr;
+  };
+
+ public:
+  using iterator = iterator_base<false>;
+  using const_iterator = iterator_base<true>;
+
+  ordered_hashmap() = default;
+  ordered_hashmap(const ordered_hashmap&) = default;
+  ordered_hashmap(ordered_hashmap&&) = default;
+  ordered_hashmap(std::initializer_list<value_type> init) : std::unordered_map<Key, T>(init) {
+    for (const auto& pair : init) {
+      order.push_back(pair.first);
+    }
+  }
+  ordered_hashmap& operator=(const ordered_hashmap&) = default;
+  ordered_hashmap& operator=(ordered_hashmap&&) = default;
+
+  iterator begin() { return {0, this}; }
+  iterator end() { return {static_cast<int>(order.size()), this}; }
+  const_iterator begin() const { return {0, this}; }
+  const_iterator end() const { return {static_cast<int>(order.size()), this}; }
+  const_iterator cbegin() const { return {0, this}; }
+  const_iterator cend() const { return {static_cast<int>(order.size()), this}; }
+
+  using std::unordered_map<Key, T>::empty;
+  using std::unordered_map<Key, T>::size;
+  using std::unordered_map<Key, T>::at;
+
+  T& operator[](const Key& key) {
+    if (count(key) == 0) {
+      order.push_back(key);
+    }
+    return std::unordered_map<Key, T>::operator[](key);
+  }
+
+  using std::unordered_map<Key, T>::count;
+
+  iterator find(const Key& key) {
+    auto it = std::find(order.begin(), order.end(), key);
+    if (it == order.end()) {
+      return end();
+    }
+    return {static_cast<int>(std::distance(order.begin(), it)), this};
+  }
+
+  const_iterator find(const Key& key) const {
+    auto it = std::find(order.begin(), order.end(), key);
+    if (it == order.end()) {
+      return end();
+    }
+    return {static_cast<int>(std::distance(order.begin(), it)), this};
+  }
+
+  void clear() {
+    std::unordered_map<Key, T>::clear();
+    order.clear();
+  }
+
+  std::pair<iterator, bool> insert(const value_type& value) {
+    if (count(value.first)) {
+      return {find(value.first), false};
+    }
+    order.push_back(value.first);
+    std::unordered_map<Key, T>::insert(value);
+    return {{static_cast<int>(order.size()) - 1, this}, true};
+  }
+
+  template <class... Args>
+  std::pair<iterator, bool> emplace(Args&&... args) {
+    return insert(value_type(std::forward<Args>(args)...));
+  }
+
+  iterator erase(const_iterator pos) {
+    assert(pos.order_index >= 0 && pos.order_index < order.size());
+    std::unordered_map<Key, T>::erase(order[pos.order_index]);
+    order.erase(order.begin() + pos.order_index);
+    return {pos.order_index, this};
+  }
+
+  iterator erase(iterator pos) {
+    assert(pos.order_index >= 0 && pos.order_index < order.size());
+    std::unordered_map<Key, T>::erase(order[pos.order_index]);
+    order.erase(order.begin() + pos.order_index);
+    return pos;
+  }
+
+  size_t erase(const Key& key) {
+    if (std::unordered_map<Key, T>::erase(key)) {
+      order.erase(std::find(order.begin(), order.end(), key));
+      return 1;
+    } else {
+      return 0;
+    }
+  }
+
+  template <typename _Key, typename _T>
+  friend bool operator==(const ordered_hashmap<_Key, _T>& lhs,
+                         const ordered_hashmap<_Key, _T>& rhs);
+
+ private:
+  std::vector<Key> order;
+};
+
+template <typename Key, typename T>
+bool operator==(const ordered_hashmap<Key, T>& lhs, const ordered_hashmap<Key, T>& rhs) {
+  return static_cast<const std::unordered_map<Key, T>&>(lhs) ==
+             static_cast<const std::unordered_map<Key, T>&>(rhs) &&
+         lhs.order == rhs.order;
+}
+
 class value {
  public:
   typedef std::vector<value> array;
+#ifdef PICOJSON_USE_ORDERED_OBJECT
+  typedef ordered_hashmap<std::string, value> object;
+#else
   typedef std::unordered_map<std::string, value> object;
+#endif
+
   union _storage {
     bool boolean_;
     double number_;

3rdparty/picojson/test_picojson.cpp

@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#include <cassert>
+#include <sstream>
+
+#include "picojson.h"
+
+using picojson::ordered_hashmap;
+
+std::string print_ordered_hashmap(const ordered_hashmap<std::string, int>& obj) {
+  std::ostringstream oss;
+  int idx = 0;
+  for (const auto& [key, value] : obj) {
+    if (idx != 0) {
+      oss << " ";
+    }
+    ++idx;
+    oss << "(" << key << ", " << value << ")";
+  }
+  return oss.str();
+}
+
+void test_constructor() {
+  ordered_hashmap<std::string, int> obj;
+  obj["foo"] = 1;
+  assert(print_ordered_hashmap(obj) == "(foo, 1)");
+
+  ordered_hashmap<std::string, int> obj1{{"foo", 1}, {"bar", 2}};
+  assert(print_ordered_hashmap(obj1) == "(foo, 1) (bar, 2)");
+
+  ordered_hashmap<std::string, int> obj2(obj1);
+  assert(print_ordered_hashmap(obj2) == "(foo, 1) (bar, 2)");
+
+  ordered_hashmap<std::string, int> obj3(std::move(obj2));
+  assert(print_ordered_hashmap(obj3) == "(foo, 1) (bar, 2)");
+
+  obj = obj3;
+  assert(print_ordered_hashmap(obj) == "(foo, 1) (bar, 2)");
+
+  const ordered_hashmap<std::string, int> obj4{{"foo", 1}, {"bar", 2}};
+  assert(print_ordered_hashmap(obj4) == "(foo, 1) (bar, 2)");
+}
+
+void test_iterator() {
+  ordered_hashmap<std::string, int> obj{{"foo", 1}, {"bar", 2}};
+  auto it = obj.begin();
+  auto past_it = it++;
+  assert(past_it == obj.begin());
+  assert(it->first == "bar");
+  assert(it->second == 2);
+  it->second = 3;
+  assert(print_ordered_hashmap(obj) == "(foo, 1) (bar, 3)");
+  (*it).second = 4;
+  assert(print_ordered_hashmap(obj) == "(foo, 1) (bar, 4)");
+  ++it;
+  assert(it == obj.end());
+  auto cit = obj.cbegin();
+  assert(cit->first == "foo");
+  assert(cit->second == 1);
+  ++cit;
+  assert(cit->first == "bar");
+  assert(cit->second == 4);
+  cit = it;
+  assert(cit == obj.cend());
+}
+
+void test_visiter() {
+  ordered_hashmap<std::string, int> obj{{"foo", 1}, {"bar", 2}};
+  auto it = obj.find("foo");
+  assert(it->first == "foo");
+  assert(it->second == 1);
+  assert(obj.find("abc") == obj.end());
+  obj["foo"] = 3;
+  assert(print_ordered_hashmap(obj) == "(foo, 3) (bar, 2)");
+  obj["abc"] = 4;
+  assert(print_ordered_hashmap(obj) == "(foo, 3) (bar, 2) (abc, 4)");
+  assert(obj.count("abc") == 1);
+  assert(obj.count("def") == 0);
+  assert(obj.at("bar") == 2);
+  assert(obj.size() == 3);
+
+  ordered_hashmap<std::string, int> obj1{{"foo", 3}, {"bar", 2}, {"abc", 4}};
+  assert(obj == obj1);
+}
+
+void test_modifier() {
+  ordered_hashmap<std::string, int> obj{{"foo", 1}, {"bar", 2}};
+
+  auto [it, successful] = obj.insert({"abc", 3});
+  assert(it->first == "abc");
+  assert(it->second == 3);
+  assert(it == obj.find("abc"));
+  ++it;
+  assert(it == obj.end());
+  assert(successful);
+
+  std::tie(it, successful) = obj.emplace("def", 4);
+  assert(it == obj.find("def"));
+  ++it;
+  assert(it == obj.end());
+  assert(successful);
+
+  std::tie(it, successful) = obj.insert({"abc", 4});
+  assert(it->second == 3);
+  ++it;
+  ++it;
+  assert(it == obj.end());
+  assert(!successful);
+
+  it = obj.find("abc");
+  it = obj.erase(it);
+  assert(it->first == "def");
+  assert(it->second == 4);
+  ++it;
+  assert(it == obj.end());
+  assert(print_ordered_hashmap(obj) == "(foo, 1) (bar, 2) (def, 4)");
+
+  assert(obj.erase("foo") == 1);
+  assert(print_ordered_hashmap(obj) == "(bar, 2) (def, 4)");
+
+  obj.clear();
+  assert(print_ordered_hashmap(obj) == "");
+  assert(obj.empty());
+}
+
+int main() {
+  test_constructor();
+  test_iterator();
+  test_visiter();
+  test_modifier();
+  return 0;
+}