From b7a531b93a0199cb502b052ea2a195136501a0b9 Mon Sep 17 00:00:00 2001
From: Arne Morten Kvarving <arne.morten.kvarving@sintef.no>
Date: Thu, 2 Feb 2023 11:52:08 +0100
Subject: [PATCH] WellState: add serialization support

---
 opm/simulators/wells/WellState.cpp  | 25 +++++++++++++++++++-
 opm/simulators/wells/WellState.hpp  | 36 ++++++++++++++++++++++++-----
 tests/test_RestartSerialization.cpp | 15 ++++++++++++
 3 files changed, 69 insertions(+), 7 deletions(-)

diff --git a/opm/simulators/wells/WellState.cpp b/opm/simulators/wells/WellState.cpp
index cd93cb2de..972ab391b 100644
--- a/opm/simulators/wells/WellState.cpp
+++ b/opm/simulators/wells/WellState.cpp
@@ -122,8 +122,24 @@ void PackUnpackXConn::unpack([[maybe_unused]] const int link,
 
 } // Anonymous namespace
 
-namespace Opm
+namespace Opm {
+
+WellState::WellState(const ParallelWellInfo& pinfo)
+    : phase_usage_{}
 {
+    wells_.add("test4",
+               SingleWellState{"dummy", pinfo, false, 0.0, {}, phase_usage_, 0.0});
+}
+
+WellState WellState::serializationTestObject(const ParallelWellInfo& pinfo)
+{
+    WellState result(PhaseUsage{});
+    result.alq_state = ALQState::serializationTestObject();
+    result.well_rates = {{"test2", {true, {1.0}}}, {"test3", {false, {2.0}}}};
+    result.wells_.add("test4", SingleWellState::serializationTestObject(pinfo));
+
+    return result;
+}
 
 void WellState::base_init(const std::vector<double>& cellPressures,
                           const std::vector<Well>& wells_ecl,
@@ -945,6 +961,13 @@ void WellState::updateWellsDefaultALQ( const std::vector<Well>& wells_ecl )
     }
 }
 
+bool WellState::operator==(const WellState& rhs) const
+{
+    return this->alq_state == rhs.alq_state &&
+           this->well_rates == rhs.well_rates &&
+           this->wells_ == rhs.wells_;
+}
+
 
 const ParallelWellInfo&
 WellState::parallelWellInfo(std::size_t well_index) const
diff --git a/opm/simulators/wells/WellState.hpp b/opm/simulators/wells/WellState.hpp
index e119ff719..23c11ac0b 100644
--- a/opm/simulators/wells/WellState.hpp
+++ b/opm/simulators/wells/WellState.hpp
@@ -21,6 +21,8 @@
 #ifndef OPM_WELLSTATEFULLYIMPLICITBLACKOIL_HEADER_INCLUDED
 #define OPM_WELLSTATEFULLYIMPLICITBLACKOIL_HEADER_INCLUDED
 
+#include <opm/common/ErrorMacros.hpp>
+
 #include <opm/core/props/BlackoilPhases.hpp>
 
 #include <opm/simulators/wells/ALQState.hpp>
@@ -58,18 +60,19 @@ class WellState
 {
 public:
     static const uint64_t event_mask = ScheduleEvents::WELL_STATUS_CHANGE + ScheduleEvents::PRODUCTION_UPDATE + ScheduleEvents::INJECTION_UPDATE;
-
-    virtual ~WellState() = default;
-
     // TODO: same definition with WellInterface, eventually they should go to a common header file.
     static const int Water = BlackoilPhases::Aqua;
     static const int Oil = BlackoilPhases::Liquid;
     static const int Gas = BlackoilPhases::Vapour;
 
+    // Only usable for testing purposes
+    explicit WellState(const ParallelWellInfo& pinfo);
+
     explicit WellState(const PhaseUsage& pu)
-    {
-        this->phase_usage_ = pu;
-    }
+        : phase_usage_(pu)
+    {}
+
+    static WellState serializationTestObject(const ParallelWellInfo& pinfo);
 
     std::size_t size() const {
         return this->wells_.size();
@@ -275,6 +278,27 @@ public:
         return this->wells_.has(well_name);
     }
 
+    bool operator==(const WellState&) const;
+
+    template<class Serializer>
+    void serializeOp(Serializer& serializer)
+    {
+        serializer(alq_state);
+        serializer(well_rates);
+        if (serializer.isSerializing()) {
+            serializer(wells_.size());
+        } else {
+            std::size_t size = 0;
+            serializer(size);
+            if (size != wells_.size()) {
+                OPM_THROW(std::runtime_error, "Error deserializing WellState: size mismatch");
+            }
+        }
+        for (auto& w : wells_) {
+            serializer(w);
+        }
+    }
+
 private:
     PhaseUsage phase_usage_;
 
diff --git a/tests/test_RestartSerialization.cpp b/tests/test_RestartSerialization.cpp
index ae538fa7d..0f71013a5 100644
--- a/tests/test_RestartSerialization.cpp
+++ b/tests/test_RestartSerialization.cpp
@@ -131,6 +131,21 @@ BOOST_AUTO_TEST_CASE(WellContainer)
     BOOST_CHECK_MESSAGE(data_out == data_in, "Deserialized WellContainer differ");
 }
 
+BOOST_AUTO_TEST_CASE(WellState)
+{
+    Opm::ParallelWellInfo dummy;
+    auto data_out = Opm::WellState::serializationTestObject(dummy);
+    Opm::Serialization::MemPacker packer;
+    Opm::Serializer ser(packer);
+    ser.pack(data_out);
+    const size_t pos1 = ser.position();
+    decltype(data_out) data_in(dummy);
+    ser.unpack(data_in);
+    const size_t pos2 = ser.position();
+    BOOST_CHECK_MESSAGE(pos1 == pos2, "Packed size differ from unpack size for WellState");
+    BOOST_CHECK_MESSAGE(data_out == data_in, "Deserialized WellState differ");
+}
+
 BOOST_AUTO_TEST_CASE(EclGenericVanguard)
 {
     auto in_params = Opm::EclGenericVanguard::serializationTestParams();