diff --git a/src/leaf_data.hpp b/src/leaf_data.hpp
index bbacbea..7139fc4 100644
--- a/src/leaf_data.hpp
+++ b/src/leaf_data.hpp
@@ -8,102 +8,34 @@
 #pragma once
 
 #include <boost/spirit/home/x3.hpp>
-#include "ast_values.hpp"
 #include "ast_handlers.hpp"
 #include "common_parsers.hpp"
+#include "leaf_data_type.hpp"
 #include "schema.hpp"
 namespace x3 = boost::spirit::x3;
 
-template <yang::LeafDataTypes TYPE>
+template <typename TYPE>
 struct leaf_data_class;
 
-template <yang::LeafDataTypes TYPE>
-struct createSetSuggestions_class {
-    std::set<std::string> getSuggestions(const ParserContext& ctx, const Schema& schema) const;
-
-    template <typename T, typename Iterator, typename Context>
-    void on_success(Iterator const& begin, Iterator const&, T&, Context const& context)
-    {
-        auto& parserContext = x3::get<parser_context_tag>(context);
-        const Schema& schema = parserContext.m_schema;
-
-        // Only generate completions if the type is correct so that we don't
-        // overwrite some other completions.
-        if (schema.leafType(parserContext.m_tmpListKeyLeafPath.m_location, parserContext.m_tmpListKeyLeafPath.m_node) == TYPE) {
-            parserContext.m_completionIterator = begin;
-            auto suggestions = getSuggestions(parserContext, schema);
-            std::set<Completion> res;
-            std::transform(suggestions.begin(), suggestions.end(), std::inserter(res, res.end()), [](auto it) { return Completion{it}; });
-            parserContext.m_suggestions = res;
-        }
-    }
-};
-
-template <>
-struct leaf_data_class<yang::LeafDataTypes::Enum> {
-    template <typename T, typename Iterator, typename Context>
-    void on_success(Iterator const&, Iterator const&, T& ast, Context const& context)
-    {
-        if (_pass(context) == false)
-            return;
-        auto& parserContext = x3::get<parser_context_tag>(context);
-        auto& schema = parserContext.m_schema;
-
-        if (!schema.leafEnumHasValue(parserContext.m_tmpListKeyLeafPath.m_location, parserContext.m_tmpListKeyLeafPath.m_node, ast.m_value)) {
-            _pass(context) = false;
-            parserContext.m_errorMsg = "leaf data type mismatch: Expected an enum here. Allowed values:";
-            for (const auto& it : schema.enumValues(parserContext.m_tmpListKeyLeafPath.m_location, parserContext.m_tmpListKeyLeafPath.m_node)) {
-                parserContext.m_errorMsg += " " + it;
-            }
-        }
-    }
-};
-
-template <>
-struct leaf_data_class<yang::LeafDataTypes::IdentityRef> {
-    template <typename T, typename Iterator, typename Context>
-    void on_success(Iterator const&, Iterator const&, T& ast, Context const& context)
-    {
-        auto& parserContext = x3::get<parser_context_tag>(context);
-        auto& schema = parserContext.m_schema;
-
-        ModuleValuePair pair;
-        if (ast.m_prefix) {
-            pair.first = ast.m_prefix.get().m_name;
-        }
-        pair.second = ast.m_value;
-
-        if (!schema.leafIdentityIsValid(parserContext.m_tmpListKeyLeafPath.m_location, parserContext.m_tmpListKeyLeafPath.m_node, pair)) {
-            _pass(context) = false;
-        }
-    }
-};
-
-x3::rule<leaf_data_class<yang::LeafDataTypes::Enum>, enum_> const leaf_data_enum = "leaf_data_enum";
-x3::rule<leaf_data_class<yang::LeafDataTypes::IdentityRef>, identityRef_> const leaf_data_identityRef = "leaf_data_identityRef";
-x3::rule<struct leaf_data_class<yang::LeafDataTypes::Binary>, binary_> const leaf_data_binary = "leaf_data_binary";
-x3::rule<struct leaf_data_class<yang::LeafDataTypes::Decimal>, double> const leaf_data_decimal = "leaf_data_decimal";
-x3::rule<struct leaf_data_class<yang::LeafDataTypes::String>, std::string> const leaf_data_string = "leaf_data_string";
+x3::rule<struct leaf_data_class<yang::Enum>, enum_> const leaf_data_enum = "leaf_data_enum";
+x3::rule<struct leaf_data_class<yang::IdentityRef>, identityRef_> const leaf_data_identityRef = "leaf_data_identityRef";
+x3::rule<struct leaf_data_class<yang::Binary>, binary_> const leaf_data_binary = "leaf_data_binary";
+x3::rule<struct leaf_data_class<yang::Decimal>, double> const leaf_data_decimal = "leaf_data_decimal";
+x3::rule<struct leaf_data_class<yang::String>, std::string> const leaf_data_string = "leaf_data_string";
 x3::rule<struct leaf_data_class_binary, std::string> const leaf_data_binary_data = "leaf_data_binary_data";
 x3::rule<struct leaf_data_identityRef_data_class, identityRef_> const leaf_data_identityRef_data = "leaf_data_identityRef_data";
 
-x3::rule<createSetSuggestions_class<yang::LeafDataTypes::Enum>, x3::unused_type> const createEnumSuggestions = "createEnumSuggestions";
-x3::rule<createSetSuggestions_class<yang::LeafDataTypes::IdentityRef>, x3::unused_type> const createIdentitySuggestions = "createIdentitySuggestions";
-
 using x3::char_;
 
-auto const createEnumSuggestions_def =
-    x3::eps;
-
 auto const leaf_data_enum_def =
-    createEnumSuggestions >> +char_;
+    +char_;
 
 struct bool_symbol_table : x3::symbols<bool> {
     bool_symbol_table()
     {
-    add
-        ("true", true)
-        ("false", false);
+        add
+            ("true", true)
+            ("false", false);
     }
 } const bool_symbols;
 
@@ -124,11 +56,124 @@
 auto const leaf_data_identityRef_data_def =
     -module  >> node_identifier;
 
-auto const createIdentitySuggestions_def =
-    x3::eps;
-
+// TODO: get rid of this and use leaf_data_identityRef_data directly
 auto const leaf_data_identityRef_def =
-    createIdentitySuggestions >> leaf_data_identityRef_data;
+    leaf_data_identityRef_data;
+
+template <typename It, typename Ctx, typename RCtx, typename Attr>
+struct impl_LeafData {
+    It& first;
+    It last;
+    Ctx const& ctx;
+    RCtx& rctx;
+    Attr& attr;
+    ParserContext& parserContext;
+
+    bool operator()(const yang::Binary&) const
+    {
+        return leaf_data_binary.parse(first, last, ctx, rctx, attr);
+    }
+    bool operator()(const yang::Bool&) const
+    {
+        return bool_symbols.parse(first, last, ctx, rctx, attr);
+    }
+    bool operator()(const yang::Decimal&) const
+    {
+        return x3::double_.parse(first, last, ctx, rctx, attr);
+    }
+    bool operator()(const yang::Uint8&) const
+    {
+        return x3::uint8.parse(first, last, ctx, rctx, attr);
+    }
+    bool operator()(const yang::Uint16&) const
+    {
+        return x3::uint16.parse(first, last, ctx, rctx, attr);
+    }
+    bool operator()(const yang::Uint32&) const
+    {
+        return x3::uint32.parse(first, last, ctx, rctx, attr);
+    }
+    bool operator()(const yang::Uint64&) const
+    {
+        return x3::uint64.parse(first, last, ctx, rctx, attr);
+    }
+    bool operator()(const yang::Int8&) const
+    {
+        return x3::int8.parse(first, last, ctx, rctx, attr);
+    }
+    bool operator()(const yang::Int16&) const
+    {
+        return x3::int16.parse(first, last, ctx, rctx, attr);
+    }
+    bool operator()(const yang::Int32&) const
+    {
+        return x3::int32.parse(first, last, ctx, rctx, attr);
+    }
+    bool operator()(const yang::Int64&) const
+    {
+        return x3::int64.parse(first, last, ctx, rctx, attr);
+    }
+    bool operator()(const yang::String&) const
+    {
+        return leaf_data_string.parse(first, last, ctx, rctx, attr);
+    }
+    template <typename Type>
+    void createSetSuggestions(const Type& type) const
+    {
+        parserContext.m_suggestions.clear();
+        std::transform(type.m_allowedValues.begin(),
+                type.m_allowedValues.end(),
+                std::inserter(parserContext.m_suggestions, parserContext.m_suggestions.end()),
+                [](auto it) { return Completion{it.m_value}; });
+        parserContext.m_completionIterator = first;
+    }
+    bool operator()(const yang::Enum& type) const
+    {
+        createSetSuggestions(type);
+        // leaf_data_enum will advance the iterator if it succeeds, so I have
+        // to save the iterator here, to roll it back in case the enum is
+        // invalid.
+        auto saveIter = first;
+        auto pass = leaf_data_enum.parse(first, last, ctx, rctx, attr);
+        if (!pass) {
+            return false;
+        }
+        auto isValidEnum = type.m_allowedValues.count(boost::get<enum_>(attr)) != 0;
+        if (!isValidEnum) {
+            first = saveIter;
+            parserContext.m_errorMsg = "leaf data type mismatch: Expected an enum here. Allowed values:";
+            for (const auto& it : type.m_allowedValues) {
+                parserContext.m_errorMsg += " " + it.m_value;
+            }
+        }
+        return isValidEnum;
+    }
+    bool operator()(const yang::IdentityRef& type) const
+    {
+        createSetSuggestions(type);
+        // leaf_data_identityRef will advance the iterator if it succeeds, so I have
+        // to save the iterator here, to roll it back in case the enum is
+        // invalid.
+        auto saveIter = first;
+        auto pass = leaf_data_identityRef.parse(first, last, ctx, rctx, attr);
+        if (!pass) {
+            return false;
+        }
+        identityRef_ pair{boost::get<identityRef_>(attr)};
+        if (!pair.m_prefix) {
+            pair.m_prefix = module_{parserContext.currentSchemaPath().m_nodes.front().m_prefix.get().m_name};
+        }
+        auto isValidIdentity  = type.m_allowedValues.count(pair) != 0;
+        if (!isValidIdentity) {
+            first = saveIter;
+        }
+        return isValidIdentity;
+    }
+    bool operator()(const yang::LeafRef& leafRef) const
+    {
+        return std::visit(*this, *leafRef.m_targetType);
+    }
+};
 
 struct LeafData : x3::parser<LeafData> {
     using attribute_type = leaf_data_;
@@ -141,43 +186,7 @@
         const Schema& schema = parserContext.m_schema;
         auto type = schema.leafType(parserContext.m_tmpListKeyLeafPath.m_location, parserContext.m_tmpListKeyLeafPath.m_node);
 
-        std::function<bool(yang::LeafDataTypes)> parse_impl = [&](auto type) {
-            switch (type) {
-            case yang::LeafDataTypes::Binary:
-                return leaf_data_binary.parse(first, last, ctx, rctx, attr);
-            case yang::LeafDataTypes::Bool:
-                return bool_symbols.parse(first, last, ctx, rctx, attr);
-            case yang::LeafDataTypes::Decimal:
-                return x3::double_.parse(first, last, ctx, rctx, attr);
-            case yang::LeafDataTypes::Uint8:
-                return x3::uint8.parse(first, last, ctx, rctx, attr);
-            case yang::LeafDataTypes::Uint16:
-                return x3::uint16.parse(first, last, ctx, rctx, attr);
-            case yang::LeafDataTypes::Uint32:
-                return x3::uint32.parse(first, last, ctx, rctx, attr);
-            case yang::LeafDataTypes::Uint64:
-                return x3::uint64.parse(first, last, ctx, rctx, attr);
-            case yang::LeafDataTypes::Int8:
-                return x3::int8.parse(first, last, ctx, rctx, attr);
-            case yang::LeafDataTypes::Int16:
-                return x3::int16.parse(first, last, ctx, rctx, attr);
-            case yang::LeafDataTypes::Int32:
-                return x3::int32.parse(first, last, ctx, rctx, attr);
-            case yang::LeafDataTypes::Int64:
-                return x3::int64.parse(first, last, ctx, rctx, attr);
-            case yang::LeafDataTypes::String:
-                return leaf_data_string.parse(first, last, ctx, rctx, attr);
-            case yang::LeafDataTypes::Enum:
-                return leaf_data_enum.parse(first, last, ctx, rctx, attr);
-            case yang::LeafDataTypes::IdentityRef:
-                return leaf_data_identityRef.parse(first, last, ctx, rctx, attr);
-            case yang::LeafDataTypes::LeafRef:
-                auto actualType = schema.leafrefBaseType(parserContext.m_tmpListKeyLeafPath.m_location, parserContext.m_tmpListKeyLeafPath.m_node);
-                return parse_impl(actualType);
-            }
-            __builtin_unreachable();
-        };
-        auto pass = parse_impl(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()) {
@@ -197,5 +206,3 @@
 BOOST_SPIRIT_DEFINE(leaf_data_binary)
 BOOST_SPIRIT_DEFINE(leaf_data_identityRef_data)
 BOOST_SPIRIT_DEFINE(leaf_data_identityRef)
-BOOST_SPIRIT_DEFINE(createEnumSuggestions)
-BOOST_SPIRIT_DEFINE(createIdentitySuggestions)