From 1a79e811dd80befe8139dd65d69117d53c75252f Mon Sep 17 00:00:00 2001
From: Joakim Hove <joakim.hove@gmail.com>
Date: Mon, 29 Jul 2019 10:50:28 +0200
Subject: [PATCH] Add operator[] to UDQConfig class

---
 .../EclipseState/Schedule/UDQ/UDQConfig.hpp   |  2 ++
 .../EclipseState/Schedule/UDQ/UDQConfig.cpp   | 21 ++++++++++++++++++-
 tests/parser/UDQTests.cpp                     |  6 ++++++
 3 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/opm/parser/eclipse/EclipseState/Schedule/UDQ/UDQConfig.hpp b/opm/parser/eclipse/EclipseState/Schedule/UDQ/UDQConfig.hpp
index ec9e717be..d84b06c96 100644
--- a/opm/parser/eclipse/EclipseState/Schedule/UDQ/UDQConfig.hpp
+++ b/opm/parser/eclipse/EclipseState/Schedule/UDQ/UDQConfig.hpp
@@ -56,6 +56,8 @@ namespace Opm {
         // from input().
         size_t size() const;
 
+        const UDQInput operator[](const std::string& keyword) const;
+
         std::vector<UDQAssign> assignments() const;
         std::vector<UDQAssign> assignments(UDQVarType var_type) const;
         const UDQParams& params() const;
diff --git a/src/opm/parser/eclipse/EclipseState/Schedule/UDQ/UDQConfig.cpp b/src/opm/parser/eclipse/EclipseState/Schedule/UDQ/UDQConfig.cpp
index 502153f1a..4e1106596 100644
--- a/src/opm/parser/eclipse/EclipseState/Schedule/UDQ/UDQConfig.cpp
+++ b/src/opm/parser/eclipse/EclipseState/Schedule/UDQ/UDQConfig.cpp
@@ -203,7 +203,7 @@ namespace Opm {
 
         /*
           That a keyword is mentioned with UNITS is enough to consider it
-           as a keyword which is present.
+          as a keyword which is present.
         */
         if (this->units.count(keyword) > 0)
             return true;
@@ -212,6 +212,25 @@ namespace Opm {
     }
 
 
+    const UDQInput UDQConfig::operator[](const std::string& keyword) const {
+        const auto index_iter = this->input_index.find(keyword);
+        if (index_iter == this->input_index.end())
+            throw std::invalid_argument("Keyword: " + keyword + " not recognized as ASSIGN/DEFINE UDQ");
+
+        std::string u;
+        if (this->has_unit(keyword))
+            u = this->unit(keyword);
+
+        if (index_iter->second.action == UDQAction::ASSIGN)
+            return UDQInput(this->input_index.at(keyword), this->m_assignments.at(keyword), u);
+
+        if (index_iter->second.action == UDQAction::DEFINE)
+            return UDQInput(this->input_index.at(keyword), this->m_definitions.at(keyword), u);
+
+        throw std::logic_error("Internal error - should not be here");
+    }
+
+
     const UDQFunctionTable& UDQConfig::function_table() const {
         return this->udqft;
     }
diff --git a/tests/parser/UDQTests.cpp b/tests/parser/UDQTests.cpp
index b03065d10..32fcdb6a8 100644
--- a/tests/parser/UDQTests.cpp
+++ b/tests/parser/UDQTests.cpp
@@ -1198,6 +1198,8 @@ UDQ
     BOOST_CHECK_EQUAL( input[6].get<UDQDefine>().keyword(), "FUOPRX");
 
     BOOST_CHECK( udq.has_keyword("FUXXX") );
+    const auto wubhp1 = udq["WUBHP1"];
+    BOOST_CHECK( wubhp1.is<UDQAssign>() );
 }
 
 
@@ -1232,6 +1234,10 @@ UDQ
     BOOST_CHECK_EQUAL(udq.size(), 5);
 
     BOOST_CHECK( input[0].is<UDQDefine>());
+    BOOST_CHECK_EQUAL( input[0].keyword(), "WUBHP1");
+
+    const auto wubhp1 = udq["WUBHP1"];
+    BOOST_CHECK( wubhp1.is<UDQDefine>() );
 }