Remove dmlc::JSONWriter

dmlc · Jul 9, 2021 · 220cc73 · 220cc73
1 parent ab621a5
commit 220cc73
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 23 deletions.
diff --git a/src/compiler/ast_native.cc b/src/compiler/ast_native.cc
@@ -8,6 +8,8 @@
 #include <treelite/compiler_param.h>
 #include <treelite/annotator.h>
 #include <fmt/format.h>
+#include <rapidjson/stringbuffer.h>
+#include <rapidjson/writer.h>
 #include <algorithm>
 #include <fstream>
 #include <unordered_map>
@@ -107,23 +109,31 @@ class ASTNativeCompiler : public Compiler {
 
     {
       /* write recipe.json */
-      std::vector<std::unordered_map<std::string, std::string>> source_list;
+      rapidjson::StringBuffer os;
+      rapidjson::Writer<rapidjson::StringBuffer> writer(os);
+
+      writer.StartObject();
+      writer.Key("target");
+      writer.String(param.native_lib_name.data(), param.native_lib_name.size());
+      writer.Key("sources");
+      writer.StartArray();
       for (const auto& kv : files_) {
         if (kv.first.compare(kv.first.length() - 2, 2, ".c") == 0) {
           const size_t line_count
             = std::count(kv.second.content.begin(), kv.second.content.end(), '\n');
-          source_list.push_back({ {"name",
-                                   kv.first.substr(0, kv.first.length() - 2)},
-                                  {"length", std::to_string(line_count)} });
+          writer.StartObject();
+          writer.Key("name");
+          std::string name = kv.first.substr(0, kv.first.length() - 2);
+          writer.String(name.data(), name.size());
+          writer.Key("length");
+          writer.Uint64(line_count);
+          writer.EndObject();
         }
       }
-      std::ostringstream oss;
-      std::unique_ptr<dmlc::JSONWriter> writer(new dmlc::JSONWriter(&oss));
-      writer->BeginObject();
-      writer->WriteObjectKeyValue("target", param.native_lib_name);
-      writer->WriteObjectKeyValue("sources", source_list);
-      writer->EndObject();
-      files_["recipe.json"] = CompiledModel::FileEntry(oss.str());
+      writer.EndArray();
+      writer.EndObject();
+
+      files_["recipe.json"] = CompiledModel::FileEntry(os.GetString());
     }
     cm.files = std::move(files_);
     return cm;

diff --git a/src/compiler/failsafe.cc b/src/compiler/failsafe.cc
@@ -10,6 +10,8 @@
 #include <treelite/compiler.h>
 #include <treelite/compiler_param.h>
 #include <fmt/format.h>
+#include <rapidjson/stringbuffer.h>
+#include <rapidjson/writer.h>
 #include <unordered_map>
 #include <set>
 #include <tuple>
@@ -360,29 +362,42 @@ class FailSafeCompiler : public Compiler {
 
     {
       /* write recipe.json */
-      std::vector<std::unordered_map<std::string, std::string>> source_list;
+      rapidjson::StringBuffer os;
+      rapidjson::Writer<rapidjson::StringBuffer> writer(os);
+
+      writer.StartObject();
+      writer.Key("target");
+      writer.String(param.native_lib_name.data(), param.native_lib_name.size());
+      writer.Key("sources");
+      writer.StartArray();
       std::vector<std::string> extra_file_list;
       for (const auto& kv : files_) {
         if (EndsWith(kv.first, ".c")) {
           const size_t line_count
             = std::count(kv.second.content.begin(), kv.second.content.end(), '\n');
-          source_list.push_back({ {"name",
-                                   kv.first.substr(0, kv.first.length() - 2)},
-                                  {"length", std::to_string(line_count)} });
+          writer.StartObject();
+          writer.Key("name");
+          std::string name = kv.first.substr(0, kv.first.length() - 2);
+          writer.String(name.data(), name.size());
+          writer.Key("length");
+          writer.Uint64(line_count);
+          writer.EndObject();
         } else if (EndsWith(kv.first, ".o")) {
           extra_file_list.push_back(kv.first);
         }
       }
-      std::ostringstream oss;
-      std::unique_ptr<dmlc::JSONWriter> writer(new dmlc::JSONWriter(&oss));
-      writer->BeginObject();
-      writer->WriteObjectKeyValue("target", param.native_lib_name);
-      writer->WriteObjectKeyValue("sources", source_list);
+      writer.EndArray();
       if (!extra_file_list.empty()) {
-        writer->WriteObjectKeyValue("extra", extra_file_list);
+        writer.Key("extra");
+        writer.StartArray();
+        for (const auto& extra_file : extra_file_list) {
+          writer.String(extra_file.data(), extra_file.size());
+        }
+        writer.EndArray();
       }
-      writer->EndObject();
-      files_["recipe.json"] = CompiledModel::FileEntry(oss.str());
+      writer.EndObject();
+
+      files_["recipe.json"] = CompiledModel::FileEntry(os.GetString());
     }
     cm.files = std::move(files_);
     return cm;