Rework Schema::units into Schema::leafType

Change-Id: I9c1f039e7e054f84559a3d57812d0254db183a99
diff --git a/src/interpreter.cpp b/src/interpreter.cpp
index c7f6644..22c1824 100644
--- a/src/interpreter.cpp
+++ b/src/interpreter.cpp
@@ -100,12 +100,12 @@
         auto leafType = m_datastore.schema()->leafType(path);
         auto typedefName = m_datastore.schema()->leafTypeName(path);
         std::string baseTypeStr;
-        if (std::holds_alternative<yang::LeafRef>(leafType)) {
+        if (std::holds_alternative<yang::LeafRef>(leafType.m_type)) {
             ss << "-> ";
             ss << m_datastore.schema()->leafrefPath(path) << " ";
-            baseTypeStr = leafDataTypeToString(*std::get<yang::LeafRef>(leafType).m_targetType);
+            baseTypeStr = leafDataTypeToString(std::get<yang::LeafRef>(leafType.m_type).m_targetType->m_type);
         } else {
-            baseTypeStr = leafDataTypeToString(leafType);
+            baseTypeStr = leafDataTypeToString(leafType.m_type);
         }
 
         if (typedefName) {
@@ -114,8 +114,8 @@
             ss << baseTypeStr;
         }
 
-        if (auto units = m_datastore.schema()->units(path)) {
-            ss << " [" + *units + "]";
+        if (leafType.m_units) {
+            ss << " [" + *leafType.m_units + "]";
         }
 
         if (m_datastore.schema()->leafIsKey(path)) {
diff --git a/src/leaf_data.hpp b/src/leaf_data.hpp
index 1c4176e..68705d0 100644
--- a/src/leaf_data.hpp
+++ b/src/leaf_data.hpp
@@ -171,12 +171,12 @@
     }
     bool operator()(const yang::LeafRef& leafRef) const
     {
-        return std::visit(*this, *leafRef.m_targetType);
+        return std::visit(*this, leafRef.m_targetType->m_type);
     }
     bool operator()(const yang::Union& unionInfo) const
     {
         return std::any_of(unionInfo.m_unionTypes.begin(), unionInfo.m_unionTypes.end(), [this](const auto& type) {
-            return std::visit(*this, type);
+            return std::visit(*this, type.m_type);
         });
     }
 };
@@ -190,14 +190,13 @@
     {
         ParserContext& parserContext = x3::get<parser_context_tag>(ctx);
         const Schema& schema = parserContext.m_schema;
-        auto type = schema.leafType(parserContext.m_tmpListKeyLeafPath.m_location, parserContext.m_tmpListKeyLeafPath.m_node);
+        auto type = schema.leafType(parserContext.m_tmpListKeyLeafPath.m_location, parserContext.m_tmpListKeyLeafPath.m_node).m_type;
 
         auto pass = std::visit(impl_LeafData<It, Ctx, RCtx, Attr>{first, last, ctx, rctx, attr, parserContext}, type);
 
         if (!pass) {
             if (parserContext.m_errorMsg.empty()) {
-                parserContext.m_errorMsg = "leaf data type mismatch: Expected " +
-                    leafDataTypeToString(schema.leafType(parserContext.m_tmpListKeyLeafPath.m_location, parserContext.m_tmpListKeyLeafPath.m_node)) + " here:";
+                parserContext.m_errorMsg = "leaf data type mismatch: Expected " + leafDataTypeToString(type) + " here:";
             }
         }
         return pass;
diff --git a/src/leaf_data_type.cpp b/src/leaf_data_type.cpp
index 8048efa..29fb355 100644
--- a/src/leaf_data_type.cpp
+++ b/src/leaf_data_type.cpp
@@ -8,6 +8,15 @@
 #include "leaf_data_type.hpp"
 
 namespace yang {
+bool TypeInfo::operator==(const TypeInfo& other) const
+{
+    return this->m_type == other.m_type && this->m_units == other.m_units;
+}
+TypeInfo::TypeInfo(const yang::LeafDataType& type, const std::optional<std::string> units)
+    : m_type(type)
+    , m_units(units)
+{
+}
 Enum::Enum(std::set<enum_>&& values)
     : m_allowedValues(std::move(values))
 {
@@ -27,10 +36,10 @@
 // Copy constructor needed, because unique_ptr is not copy-constructible
 LeafRef::LeafRef(const LeafRef& src)
     : m_targetXPath(src.m_targetXPath)
-    , m_targetType(std::make_unique<LeafDataType>(*src.m_targetType))
+    , m_targetType(std::make_unique<TypeInfo>(*src.m_targetType))
 {
 }
-LeafRef::LeafRef(const std::string& xpath, std::unique_ptr<LeafDataType>&& type)
+LeafRef::LeafRef(const std::string& xpath, std::unique_ptr<TypeInfo>&& type)
     : m_targetXPath(xpath)
     , m_targetType(std::move(type))
 {
diff --git a/src/leaf_data_type.hpp b/src/leaf_data_type.hpp
index a86691b..470c435 100644
--- a/src/leaf_data_type.hpp
+++ b/src/leaf_data_type.hpp
@@ -82,16 +82,23 @@
     yang::LeafRef,
     yang::Union
 >;
+struct TypeInfo;
 struct LeafRef {
     LeafRef(const LeafRef& src);
-    LeafRef(const std::string& xpath, std::unique_ptr<LeafDataType>&& type);
+    LeafRef(const std::string& xpath, std::unique_ptr<TypeInfo>&& type);
     bool operator==(const LeafRef& other) const;
     std::string m_targetXPath;
-    std::unique_ptr<LeafDataType> m_targetType;
+    std::unique_ptr<TypeInfo> m_targetType;
 };
 
 struct Union {
     bool operator==(const Union& other) const;
-    std::vector<LeafDataType> m_unionTypes;
+    std::vector<TypeInfo> m_unionTypes;
+};
+struct TypeInfo {
+    TypeInfo(const yang::LeafDataType& type, const std::optional<std::string> units = std::nullopt);
+    bool operator==(const TypeInfo& other) const;
+    yang::LeafDataType m_type;
+    std::optional<std::string> m_units;
 };
 }
diff --git a/src/schema.hpp b/src/schema.hpp
index 9a86045..8cdffa1 100644
--- a/src/schema.hpp
+++ b/src/schema.hpp
@@ -55,12 +55,11 @@
     virtual bool listHasKey(const schemaPath_& location, const ModuleNodePair& node, const std::string& key) const = 0;
     virtual bool leafIsKey(const std::string& leafPath) const = 0;
     virtual const std::set<std::string> listKeys(const schemaPath_& location, const ModuleNodePair& node) const = 0;
-    virtual yang::LeafDataType leafType(const schemaPath_& location, const ModuleNodePair& node) const = 0;
-    virtual yang::LeafDataType leafType(const std::string& path) const = 0;
+    virtual yang::TypeInfo leafType(const schemaPath_& location, const ModuleNodePair& node) const = 0;
+    virtual yang::TypeInfo leafType(const std::string& path) const = 0;
     virtual std::optional<std::string> leafTypeName(const std::string& path) const = 0;
     virtual std::string leafrefPath(const std::string& leafrefPath) const = 0;
     virtual std::optional<std::string> description(const std::string& location) const = 0;
-    virtual std::optional<std::string> units(const std::string& location) const = 0;
 
     virtual std::set<std::string> childNodes(const schemaPath_& path, const Recursion recursion) const = 0;
     virtual std::set<std::string> moduleNodes(const module_& module, const Recursion recursion) const = 0;
diff --git a/src/static_schema.cpp b/src/static_schema.cpp
index 3eb42f6..6815015 100644
--- a/src/static_schema.cpp
+++ b/src/static_schema.cpp
@@ -86,7 +86,7 @@
 
 void StaticSchema::addLeaf(const std::string& location, const std::string& name, const yang::LeafDataType& type)
 {
-    m_nodes.at(location).emplace(name, yang::leaf{type});
+    m_nodes.at(location).emplace(name, yang::leaf{yang::TypeInfo{type, std::nullopt}});
     std::string key = joinPaths(location, name);
     m_nodes.emplace(key, std::unordered_map<std::string, NodeType>());
 }
@@ -122,13 +122,13 @@
     return res;
 }
 
-yang::LeafDataType StaticSchema::leafType(const schemaPath_& location, const ModuleNodePair& node) const
+yang::TypeInfo StaticSchema::leafType(const schemaPath_& location, const ModuleNodePair& node) const
 {
     std::string locationString = pathToSchemaString(location, Prefixes::Always);
     return boost::get<yang::leaf>(children(locationString).at(fullNodeName(location, node))).m_type;
 }
 
-yang::LeafDataType StaticSchema::leafType(const std::string& path) const
+yang::TypeInfo StaticSchema::leafType(const std::string& path) const
 {
     auto locationString = stripLastNodeFromPath(path);
     auto node = lastNodeOfSchemaPath(path);
@@ -197,11 +197,6 @@
     throw std::runtime_error{"StaticSchema::description not implemented"};
 }
 
-std::optional<std::string> StaticSchema::units([[maybe_unused]] const std::string& path) const
-{
-    throw std::runtime_error{"StaticSchema::units not implemented"};
-}
-
 yang::NodeTypes StaticSchema::nodeType([[maybe_unused]] const std::string& path) const
 {
     throw std::runtime_error{"Internal error: StaticSchema::nodeType(std::string) not implemented. The tests should not have called this overload."};
diff --git a/src/static_schema.hpp b/src/static_schema.hpp
index a94c78d..8b160b0 100644
--- a/src/static_schema.hpp
+++ b/src/static_schema.hpp
@@ -27,7 +27,7 @@
 };
 
 struct leaf {
-    yang::LeafDataType m_type;
+    yang::TypeInfo m_type;
 };
 
 struct module {
@@ -51,14 +51,13 @@
     bool listHasKey(const schemaPath_& location, const ModuleNodePair& node, const std::string& key) const override;
     bool leafIsKey(const std::string& leafPath) const override;
     const std::set<std::string> listKeys(const schemaPath_& location, const ModuleNodePair& node) const override;
-    yang::LeafDataType leafType(const schemaPath_& location, const ModuleNodePair& node) const override;
-    yang::LeafDataType leafType(const std::string& path) const override;
+    yang::TypeInfo leafType(const schemaPath_& location, const ModuleNodePair& node) const override;
+    yang::TypeInfo leafType(const std::string& path) const override;
     std::optional<std::string> leafTypeName(const std::string& path) const override;
     std::string leafrefPath(const std::string& leafrefPath) const override;
     std::set<std::string> childNodes(const schemaPath_& path, const Recursion) const override;
     std::set<std::string> moduleNodes(const module_& module, const Recursion recursion) const override;
     std::optional<std::string> description(const std::string& path) const override;
-    std::optional<std::string> units(const std::string& path) const override;
 
     /** A helper for making tests a little bit easier. It returns all
      * identities which are based on the argument passed and which can then be
diff --git a/src/utils.cpp b/src/utils.cpp
index 31854b2..9496db3 100644
--- a/src/utils.cpp
+++ b/src/utils.cpp
@@ -118,7 +118,7 @@
     {
         std::ostringstream ss;
         std::transform(type.m_unionTypes.begin(), type.m_unionTypes.end(), std::experimental::make_ostream_joiner(ss, ", "), [this](const auto& unionType) {
-            return std::visit(*this, unionType);
+            return std::visit(*this, unionType.m_type);
         });
         return ss.str();
     }
diff --git a/src/yang_schema.cpp b/src/yang_schema.cpp
index a453b50..5a34e4f 100644
--- a/src/yang_schema.cpp
+++ b/src/yang_schema.cpp
@@ -201,65 +201,99 @@
 }
 }
 
-yang::LeafDataType YangSchema::impl_leafType(const libyang::S_Schema_Node& node) const
+yang::TypeInfo YangSchema::impl_leafType(const libyang::S_Schema_Node& node) const
 {
     using namespace std::string_literals;
     auto leaf = std::make_shared<libyang::Schema_Node_Leaf>(node);
-    std::function<yang::LeafDataType(std::shared_ptr<libyang::Type>)> resolveType;
-    resolveType = [this, &resolveType, leaf] (auto type) -> yang::LeafDataType {
+    auto leafUnits = leaf->units();
+    std::function<yang::TypeInfo(std::shared_ptr<libyang::Type>)> resolveType;
+    resolveType = [this, &resolveType, leaf, leafUnits] (std::shared_ptr<libyang::Type> type) -> yang::TypeInfo {
+        yang::LeafDataType resType;
         switch (type->base()) {
         case LY_TYPE_STRING:
-            return yang::String{};
+            resType.emplace<yang::String>();
+            break;
         case LY_TYPE_DEC64:
-            return yang::Decimal{};
+            resType.emplace<yang::Decimal>();
+            break;
         case LY_TYPE_BOOL:
-            return yang::Bool{};
+            resType.emplace<yang::Bool>();
+            break;
         case LY_TYPE_INT8:
-            return yang::Int8{};
+            resType.emplace<yang::Int8>();
+            break;
         case LY_TYPE_INT16:
-            return yang::Int16{};
+            resType.emplace<yang::Int16>();
+            break;
         case LY_TYPE_INT32:
-            return yang::Int32{};
+            resType.emplace<yang::Int32>();
+            break;
         case LY_TYPE_INT64:
-            return yang::Int64{};
+            resType.emplace<yang::Int64>();
+            break;
         case LY_TYPE_UINT8:
-            return yang::Uint8{};
+            resType.emplace<yang::Uint8>();
+            break;
         case LY_TYPE_UINT16:
-            return yang::Uint16{};
+            resType.emplace<yang::Uint16>();
+            break;
         case LY_TYPE_UINT32:
-            return yang::Uint32{};
+            resType.emplace<yang::Uint32>();
+            break;
         case LY_TYPE_UINT64:
-            return yang::Uint64{};
+            resType.emplace<yang::Uint64>();
+            break;
         case LY_TYPE_BINARY:
-            return yang::Binary{};
+            resType.emplace<yang::Binary>();
+            break;
         case LY_TYPE_ENUM:
-            return yang::Enum{enumValues(type)};
+            resType.emplace<yang::Enum>(enumValues(type));
+            break;
         case LY_TYPE_IDENT:
-            return yang::IdentityRef{validIdentities(type)};
+            resType.emplace<yang::IdentityRef>(validIdentities(type));
+            break;
         case LY_TYPE_LEAFREF:
-            return yang::LeafRef{::leafrefPath(type), std::make_unique<yang::LeafDataType>(leafType(::leafrefPath(type)))};
+            resType.emplace<yang::LeafRef>(::leafrefPath(type), std::make_unique<yang::TypeInfo>(leafType(::leafrefPath(type))));
+            break;
         case LY_TYPE_UNION:
             {
-                auto res = yang::Union{};
+                auto resUnion = yang::Union{};
                 for (auto unionType : type->info()->uni()->types()) {
-                    res.m_unionTypes.push_back(resolveType(unionType));
+                    resUnion.m_unionTypes.push_back(resolveType(unionType));
                 }
-                return res;
+                resType.emplace<yang::Union>(std::move(resUnion));
+                break;
             }
         default:
             using namespace std::string_literals;
             throw UnsupportedYangTypeException("the type of "s + leaf->name() + " is not supported: " + std::to_string(leaf->type()->base()));
         }
-        };
+
+        std::optional<std::string> resUnits;
+
+        if (leafUnits) {
+            resUnits = leafUnits;
+        } else {
+            for (auto parentTypedef = type->der(); parentTypedef; parentTypedef = parentTypedef->type()->der()) {
+                auto units = parentTypedef->units();
+                if (units) {
+                    resUnits = units;
+                    break;
+                }
+            }
+        }
+
+        return yang::TypeInfo(resType, resUnits);
+    };
     return resolveType(leaf->type());
 }
 
-yang::LeafDataType YangSchema::leafType(const schemaPath_& location, const ModuleNodePair& node) const
+yang::TypeInfo YangSchema::leafType(const schemaPath_& location, const ModuleNodePair& node) const
 {
     return impl_leafType(getSchemaNode(location, node));
 }
 
-yang::LeafDataType YangSchema::leafType(const std::string& path) const
+yang::TypeInfo YangSchema::leafType(const std::string& path) const
 {
     return impl_leafType(getSchemaNode(path));
 }
@@ -421,28 +455,3 @@
     auto node = getSchemaNode(path.c_str());
     return node->dsc() ? std::optional{node->dsc()} : std::nullopt;
 }
-
-std::optional<std::string> YangSchema::units(const std::string& path) const
-{
-    auto node = getSchemaNode(path.c_str());
-    if (node->nodetype() != LYS_LEAF) {
-        return std::nullopt;
-    }
-    libyang::Schema_Node_Leaf leaf{node};
-    auto units = leaf.units();
-
-    // A leaf can specify units as part of its definition.
-    if (units) {
-        return units;
-    }
-
-    // A typedef (or its parent typedefs) can specify units too. We'll use the first `units` we find.
-    for (auto parentTypedef = leaf.type()->der(); parentTypedef; parentTypedef = parentTypedef->type()->der()) {
-        units = parentTypedef->units();
-        if (units) {
-            return units;
-        }
-    }
-
-    return std::nullopt;
-}
diff --git a/src/yang_schema.hpp b/src/yang_schema.hpp
index e3a9025..a2981ed 100644
--- a/src/yang_schema.hpp
+++ b/src/yang_schema.hpp
@@ -37,14 +37,13 @@
     bool listHasKey(const schemaPath_& location, const ModuleNodePair& node, const std::string& key) const override;
     bool leafIsKey(const std::string& leafPath) const override;
     const std::set<std::string> listKeys(const schemaPath_& location, const ModuleNodePair& node) const override;
-    yang::LeafDataType leafType(const schemaPath_& location, const ModuleNodePair& node) const override;
-    yang::LeafDataType leafType(const std::string& path) const override;
+    yang::TypeInfo leafType(const schemaPath_& location, const ModuleNodePair& node) const override;
+    yang::TypeInfo leafType(const std::string& path) const override;
     std::optional<std::string> leafTypeName(const std::string& path) const override;
     std::string leafrefPath(const std::string& leafrefPath) const override;
     std::set<std::string> childNodes(const schemaPath_& path, const Recursion recursion) const override;
     std::set<std::string> moduleNodes(const module_& module, const Recursion recursion) const override;
     std::optional<std::string> description(const std::string& path) const override;
-    std::optional<std::string> units(const std::string& path) const override;
 
     void registerModuleCallback(const std::function<std::string(const char*, const char*, const char*, const char*)>& clb);
 
@@ -68,7 +67,7 @@
     std::shared_ptr<libyang::Module> getYangModule(const std::string& name);
 
 private:
-    yang::LeafDataType impl_leafType(const std::shared_ptr<libyang::Schema_Node>& node) const;
+    yang::TypeInfo impl_leafType(const std::shared_ptr<libyang::Schema_Node>& node) const;
     std::set<std::string> modules() const;
 
     /** @short Returns a set of nodes, that match the location and name criteria. */