From 088dfbb912e17fe5e9a1074e602aa35895998b63 Mon Sep 17 00:00:00 2001
From: Florian Fontan <dev@florian-fontan.fr>
Date: Thu, 19 Dec 2024 21:53:31 +0100
Subject: [PATCH] Update JSON search tree output

---
 include/treesearchsolver/common.hpp           | 76 ++++++++++++++++---
 .../iterative_beam_search_2.hpp               | 12 ++-
 2 files changed, 74 insertions(+), 14 deletions(-)

diff --git a/include/treesearchsolver/common.hpp b/include/treesearchsolver/common.hpp
index 6f22414..0171ea6 100644
--- a/include/treesearchsolver/common.hpp
+++ b/include/treesearchsolver/common.hpp
@@ -501,11 +501,11 @@ void solution_write(
 }
 
 ////////////////////////////////////////////////////////////////////////////////
-///////////////////////// branching_scheme_node_export /////////////////////////
+///////////////////////// branching_scheme_json_export /////////////////////////
 ////////////////////////////////////////////////////////////////////////////////
 
 template<typename, typename T>
-struct HasNodeExportMethod
+struct HasJsonExportInitMethod
 {
     static_assert(
         std::integral_constant<T, false>::value,
@@ -513,13 +513,13 @@ struct HasNodeExportMethod
 };
 
 template<typename C, typename Ret, typename... Args>
-struct HasNodeExportMethod<C, Ret(Args...)>
+struct HasJsonExportInitMethod<C, Ret(Args...)>
 {
 
 private:
 
     template<typename T>
-    static constexpr auto check(T*) -> typename std::is_same<decltype(std::declval<T>().node_export(std::declval<Args>()...)), Ret>::type;
+    static constexpr auto check(T*) -> typename std::is_same<decltype(std::declval<T>().json_export_init(std::declval<Args>()...)), Ret>::type;
 
     template<typename>
     static constexpr std::false_type check(...);
@@ -533,7 +533,63 @@ struct HasNodeExportMethod<C, Ret(Args...)>
 };
 
 template<typename BranchingScheme>
-nlohmann::json node_export(
+nlohmann::json json_export_init(
+        const BranchingScheme&,
+        std::false_type)
+{
+    return nlohmann::json{};
+}
+
+template<typename BranchingScheme>
+nlohmann::json json_export_init(
+        const BranchingScheme& branching_scheme,
+        std::true_type)
+{
+    return branching_scheme.json_export_init();
+}
+
+template<typename BranchingScheme>
+nlohmann::json json_export_init(
+        const BranchingScheme& branching_scheme)
+{
+    return json_export_init(
+            branching_scheme,
+            std::integral_constant<
+                bool,
+                HasJsonExportInitMethod<BranchingScheme,
+                nlohmann::json()>::value>());
+}
+
+template<typename, typename T>
+struct HasJsonExportMethod
+{
+    static_assert(
+        std::integral_constant<T, false>::value,
+        "Second template parameter needs to be of function type.");
+};
+
+template<typename C, typename Ret, typename... Args>
+struct HasJsonExportMethod<C, Ret(Args...)>
+{
+
+private:
+
+    template<typename T>
+    static constexpr auto check(T*) -> typename std::is_same<decltype(std::declval<T>().json_export(std::declval<Args>()...)), Ret>::type;
+
+    template<typename>
+    static constexpr std::false_type check(...);
+
+    typedef decltype(check<C>(0)) type;
+
+public:
+
+    static constexpr bool value = type::value;
+
+};
+
+template<typename BranchingScheme>
+nlohmann::json json_export(
         const BranchingScheme&,
         const std::shared_ptr<typename BranchingScheme::Node>&,
         std::false_type)
@@ -542,25 +598,25 @@ nlohmann::json node_export(
 }
 
 template<typename BranchingScheme>
-nlohmann::json node_export(
+nlohmann::json json_export(
         const BranchingScheme& branching_scheme,
         const std::shared_ptr<typename BranchingScheme::Node>& node,
         std::true_type)
 {
-    return branching_scheme.node_export(node);
+    return branching_scheme.json_export(node);
 }
 
 template<typename BranchingScheme>
-nlohmann::json node_export(
+nlohmann::json json_export(
         const BranchingScheme& branching_scheme,
         const std::shared_ptr<typename BranchingScheme::Node>& node)
 {
-    return node_export(
+    return json_export(
             branching_scheme,
             node,
             std::integral_constant<
                 bool,
-                HasNodeExportMethod<BranchingScheme,
+                HasJsonExportMethod<BranchingScheme,
                 nlohmann::json(const std::shared_ptr<typename BranchingScheme::Node>&)>::value>());
 }
 
diff --git a/include/treesearchsolver/iterative_beam_search_2.hpp b/include/treesearchsolver/iterative_beam_search_2.hpp
index 1b26395..6bcbe78 100644
--- a/include/treesearchsolver/iterative_beam_search_2.hpp
+++ b/include/treesearchsolver/iterative_beam_search_2.hpp
@@ -133,14 +133,18 @@ inline const IterativeBeamSearch2Output<BranchingScheme> iterative_beam_search_2
             new NodeMap<BranchingScheme>(0, node_hasher, node_hasher));
     Depth number_of_queues = 2;
 
+    if (parameters.write_json_search_tree) {
+        output.json_search_tree["Init"] = json_export_init(branching_scheme);
+    }
+
     for (output.maximum_size_of_the_queue = parameters.minimum_size_of_the_queue;;) {
 
         // Initialize queue.
         bool stop = true;
         auto current_node = branching_scheme.root();
         if (parameters.write_json_search_tree) {
-            output.json_search_tree[current_node->id]
-                = node_export(branching_scheme, current_node);
+            output.json_search_tree["Nodes"][current_node->id]
+                = json_export(branching_scheme, current_node);
         }
         output.number_of_nodes_generated++;
         q[0]->insert(current_node);
@@ -192,8 +196,8 @@ inline const IterativeBeamSearch2Output<BranchingScheme> iterative_beam_search_2
                     }
 
                     if (parameters.write_json_search_tree) {
-                        output.json_search_tree[child->id]
-                            = node_export(branching_scheme, child);
+                        output.json_search_tree["Nodes"][child->id]
+                            = json_export(branching_scheme, child);
                     }
 
                     // Add child to the queue.