From 482c32b117dadadab33c31dac18e69147cc5446a 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] SegmentState: add serialization support

---
 opm/simulators/wells/SegmentState.cpp | 35 +++++++++++++++++++++++++++
 opm/simulators/wells/SegmentState.hpp | 21 ++++++++++++++++
 tests/test_RestartSerialization.cpp   |  2 ++
 3 files changed, 58 insertions(+)

diff --git a/opm/simulators/wells/SegmentState.cpp b/opm/simulators/wells/SegmentState.cpp
index 6984f6b2f..52a2177aa 100644
--- a/opm/simulators/wells/SegmentState.cpp
+++ b/opm/simulators/wells/SegmentState.cpp
@@ -68,6 +68,25 @@ SegmentState::SegmentState(int num_phases, const WellSegments& segments)
     , m_segment_number         (make_segment_number(segments))
 {}
 
+SegmentState SegmentState::serializationTestObject()
+{
+    SegmentState result;
+    result.rates = {1.0, 2.0};
+    result.dissolved_gas_rate = {3.0, 4.0, 5.0};
+    result.vaporized_oil_rate = {6.0};
+    result.phase_resv_rates = {7.0, 8.0};
+    result.phase_velocity = {9.0};
+    result.phase_holdup = {10.0, 11.0};
+    result.phase_viscosity = {12.0};
+    result.pressure = {13.0, 14.0};
+    result.pressure_drop_friction = {15.0};
+    result.pressure_drop_hydrostatic = {16.0, 17.0};
+    result.pressure_drop_accel = {18.0};
+    result.m_segment_number = {19, 20};
+
+    return result;
+}
+
 double SegmentState::pressure_drop(std::size_t index) const {
     return this->pressure_drop_friction[index] + this->pressure_drop_hydrostatic[index] + this->pressure_drop_accel[index];
 }
@@ -96,4 +115,20 @@ const std::vector<int>& SegmentState::segment_number() const {
     return this->m_segment_number;
 }
 
+bool SegmentState::operator==(const SegmentState& rhs) const
+{
+    return this->rates == rhs.rates &&
+           this->dissolved_gas_rate == rhs.dissolved_gas_rate &&
+           this->vaporized_oil_rate == rhs.vaporized_oil_rate &&
+           this->phase_resv_rates == rhs.phase_resv_rates &&
+           this->phase_velocity == rhs.phase_velocity &&
+           this->phase_holdup == rhs.phase_holdup &&
+           this->phase_viscosity == rhs.phase_viscosity &&
+           this->pressure == rhs.pressure &&
+           this->pressure_drop_friction == rhs.pressure_drop_friction &&
+           this->pressure_drop_hydrostatic == rhs.pressure_drop_hydrostatic &&
+           this->pressure_drop_accel == rhs.pressure_drop_accel &&
+           this->m_segment_number == rhs.m_segment_number;
+}
+
 } // namespace Opm
diff --git a/opm/simulators/wells/SegmentState.hpp b/opm/simulators/wells/SegmentState.hpp
index 00d7049c4..71cde3dec 100644
--- a/opm/simulators/wells/SegmentState.hpp
+++ b/opm/simulators/wells/SegmentState.hpp
@@ -37,6 +37,8 @@ public:
     SegmentState() = default;
     SegmentState(int num_phases, const WellSegments& segments);
 
+    static SegmentState serializationTestObject();
+
     double pressure_drop(std::size_t index) const;
     bool empty() const;
     void scale_pressure(double bhp);
@@ -44,6 +46,25 @@ public:
     const std::vector<int>& segment_number() const;
     std::size_t size() const;
 
+    template<class Serializer>
+    void serializeOp(Serializer& serializer)
+    {
+        serializer(rates);
+        serializer(dissolved_gas_rate);
+        serializer(vaporized_oil_rate);
+        serializer(phase_resv_rates);
+        serializer(phase_velocity);
+        serializer(phase_holdup);
+        serializer(phase_viscosity);
+        serializer(pressure);
+        serializer(pressure_drop_friction);
+        serializer(pressure_drop_hydrostatic);
+        serializer(pressure_drop_accel);
+        serializer(m_segment_number);
+    }
+
+    bool operator==(const SegmentState&) const;
+
     std::vector<double> rates;
     std::vector<double> dissolved_gas_rate;
     std::vector<double> vaporized_oil_rate;
diff --git a/tests/test_RestartSerialization.cpp b/tests/test_RestartSerialization.cpp
index fb718e800..f74b7269a 100644
--- a/tests/test_RestartSerialization.cpp
+++ b/tests/test_RestartSerialization.cpp
@@ -35,6 +35,7 @@
 #include <opm/simulators/timestepping/SimulatorTimer.hpp>
 #include <opm/simulators/timestepping/TimeStepControl.hpp>
 #include <opm/simulators/utils/SerializationPackers.hpp>
+#include <opm/simulators/wells/SegmentState.hpp>
 
 #define BOOST_TEST_MODULE TestRestartSerialization
 #define BOOST_TEST_NO_MAIN
@@ -76,6 +77,7 @@ BOOST_AUTO_TEST_CASE(NAME) \
 TEST_FOR_TYPE(HardcodedTimeStepControl)
 TEST_FOR_TYPE(PIDAndIterationCountTimeStepControl)
 TEST_FOR_TYPE(PIDTimeStepControl)
+TEST_FOR_TYPE(SegmentState)
 TEST_FOR_TYPE(SimpleIterationCountTimeStepControl)
 TEST_FOR_TYPE(SimulatorTimer)