From 888d897a3eda962e6b765085bb40d115b05a3491 Mon Sep 17 00:00:00 2001
From: Jason Turner <jason@emptycrate.com>
Date: Sat, 2 Jan 2016 19:59:54 -0700
Subject: [PATCH] Simplify use of enum helper

---
 include/chaiscript/utility/utility.hpp |  8 +++++---
 unittests/compiled_tests.cpp           | 10 +++++-----
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/include/chaiscript/utility/utility.hpp b/include/chaiscript/utility/utility.hpp
index 05c49c5..381afda 100644
--- a/include/chaiscript/utility/utility.hpp
+++ b/include/chaiscript/utility/utility.hpp
@@ -69,8 +69,7 @@ namespace chaiscript
       typename std::enable_if<std::is_enum<Enum>::value, void>::type
       add_class(ModuleType &t_module,
         const std::string &t_class_name,
-        const std::vector<chaiscript::Proxy_Function> &t_constructors,
-        const std::vector<std::pair<chaiscript::Boxed_Value, std::string>> &t_constants)
+        const std::vector<std::pair<typename std::underlying_type<Enum>::type, std::string>> &t_constants)
       {
         t_module.add(chaiscript::user_type<Enum>(), t_class_name);
 
@@ -83,9 +82,12 @@ namespace chaiscript
               return assign<Enum>(not_equal<Enum>(equal<Enum>()));
             }());
 
+        t_module.add(chaiscript::fun([](const Enum &e, const typename std::underlying_type<Enum>::type &i) { return e == i; }), "==");
+        t_module.add(chaiscript::fun([](const typename std::underlying_type<Enum>::type &i, const Enum &e) { return i == e; }), "==");
+
         for (const auto &constant : t_constants)
         {
-          t_module.add_global_const(constant.first, constant.second);
+          t_module.add_global_const(chaiscript::const_var(Enum(constant.first)), constant.second);
         }
       }
   }
diff --git a/unittests/compiled_tests.cpp b/unittests/compiled_tests.cpp
index a028526..6724973 100644
--- a/unittests/compiled_tests.cpp
+++ b/unittests/compiled_tests.cpp
@@ -539,11 +539,9 @@ TEST_CASE("Utility_Test utility class wrapper for enum")
 
   chaiscript::utility::add_class<Utility_Test_Numbers>(*m,
       "Utility_Test_Numbers",
-      {
-      },
-      { { const_var(ONE), "ONE" },
-        { const_var(TWO), "TWO" },
-        { const_var(THREE), "THREE" }
+      { { ONE, "ONE" },
+        { TWO, "TWO" },
+        { THREE, "THREE" }
 
         }
       );
@@ -556,6 +554,8 @@ TEST_CASE("Utility_Test utility class wrapper for enum")
   CHECK(chai.eval<Utility_Test_Numbers>("TWO ") == 1);
   CHECK(chai.eval<Utility_Test_Numbers>("THREE ") == 2);
 
+  CHECK(chai.eval<bool>("ONE == 0"));
+
   chai.add(chaiscript::fun(&do_something_with_enum_vector), "do_something_with_enum_vector");
   chai.add(chaiscript::vector_conversion<std::vector<Utility_Test_Numbers>>());
   CHECK_NOTHROW(chai.eval("var a = [ONE, TWO, THREE]"));