diff --git a/src/ast_handlers.hpp b/src/ast_handlers.hpp
index d365641..fede04c 100644
--- a/src/ast_handlers.hpp
+++ b/src/ast_handlers.hpp
@@ -360,6 +360,22 @@
     }
 };
 
+// This handler only checks if the module exists
+// It doesn't set any ParserContext flags (except the error message)
+struct data_module_prefix_class {
+    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);
+        const auto& schema = parserContext.m_schema;
+
+        if (!schema.isModule(parserContext.m_curPath, ast.m_name)) {
+            parserContext.m_errorMsg = "Invalid module name.";
+            _pass(context) = false;
+        }
+    }
+};
+
 struct leaf_data_class {
     template <typename Iterator, typename Exception, typename Context>
     x3::error_handler_result on_error(Iterator&, Iterator const&, Exception const&, Context const& context)
@@ -369,10 +385,10 @@
         if (parserContext.m_errorMsg.empty()) {
             leaf_ leaf = boost::get<leaf_>(parserContext.m_curPath.m_nodes.back().m_suffix);
             schemaPath_ location = pathWithoutLastNode(parserContext.m_curPath);
-            if (location.m_nodes.empty()) {
-                parserContext.m_curModule = parserContext.m_curPath.m_nodes.back().m_prefix->m_name;
-            }
-            parserContext.m_errorMsg = "Expected " + leafDataTypeToString(schema.leafType(location, {parserContext.m_curModule, leaf.m_name})) + " here:";
+            boost::optional<std::string> module;
+            if (parserContext.m_curPath.m_nodes.back().m_prefix)
+                module = parserContext.m_curPath.m_nodes.back().m_prefix.value().m_name;
+            parserContext.m_errorMsg = "Expected " + leafDataTypeToString(schema.leafType(location, {module, leaf.m_name})) + " here:";
             return x3::error_handler_result::fail;
         }
         return x3::error_handler_result::rethrow;
@@ -415,6 +431,9 @@
     void on_success(Iterator const& start, Iterator const& end, T& ast, Context const& context)
     {
         leaf_data_base_class::on_success(start, end, ast, context);
+        // Base class on_success cannot return for us, so we check if it failed the parser.
+        if (_pass(context) == false)
+            return;
         auto& parserContext = x3::get<parser_context_tag>(context);
         auto& schema = parserContext.m_schema;
         boost::optional<std::string> module;
@@ -474,6 +493,43 @@
     }
 };
 
+struct leaf_data_identityRef_data_class;
+
+struct leaf_data_identityRef_class : leaf_data_base_class {
+    leaf_data_identityRef_class()
+        : leaf_data_base_class(yang::LeafDataTypes::IdentityRef)
+    {
+    }
+
+    template <typename T, typename Iterator, typename Context>
+    void on_success(Iterator const& start, Iterator const& end, T& ast, Context const& context)
+    {
+        // FIXME: can I reuse leaf_data_enum_class somehow..?
+        leaf_data_base_class::on_success(start, end, ast, context);
+        // Base class on_success cannot return for us, so we check if it failed the parser.
+        if (_pass(context) == false)
+            return;
+        auto& parserContext = x3::get<parser_context_tag>(context);
+        auto& schema = parserContext.m_schema;
+        boost::optional<std::string> module;
+        if (parserContext.m_curPath.m_nodes.back().m_prefix)
+            module = parserContext.m_curPath.m_nodes.back().m_prefix.value().m_name;
+
+        leaf_ leaf = boost::get<leaf_>(parserContext.m_curPath.m_nodes.back().m_suffix);
+        schemaPath_ location = pathWithoutLastNode(parserContext.m_curPath);
+
+        ModuleValuePair pair;
+        if (ast.m_prefix) {
+            pair.first = ast.m_prefix.get().m_name;
+        }
+        pair.second = ast.m_value;
+
+        if (!schema.leafIdentityIsValid(location, {module, leaf.m_name}, pair)) {
+            _pass(context) = false;
+        }
+    }
+};
+
 struct set_class {
     template <typename Iterator, typename Exception, typename Context>
     x3::error_handler_result on_error(Iterator&, Iterator const&, Exception const& x, Context const& context)
@@ -616,6 +672,32 @@
         leaf_ leaf = boost::get<leaf_>(parserContext.m_curPath.m_nodes.back().m_suffix);
         schemaPath_ location = pathWithoutLastNode(parserContext.m_curPath);
 
-        parserContext.m_suggestions = schema.enumValues(location, {module, leaf.m_name});
+        // Only generate completions if the type is enum so that we don't
+        // overwrite some other completions.
+        if (schema.leafType(location, {module, leaf.m_name}) == yang::LeafDataTypes::Enum)
+            parserContext.m_suggestions = schema.enumValues(location, {module, leaf.m_name});
+    }
+};
+
+// FIXME: can I reuse createEnumSuggestions?
+struct createIdentitySuggestions_class {
+    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);
+        parserContext.m_completionIterator = begin;
+        const Schema& schema = parserContext.m_schema;
+
+        boost::optional<std::string> module;
+        if (parserContext.m_curPath.m_nodes.back().m_prefix)
+            module = parserContext.m_curPath.m_nodes.back().m_prefix.value().m_name;
+
+        leaf_ leaf = boost::get<leaf_>(parserContext.m_curPath.m_nodes.back().m_suffix);
+        schemaPath_ location = pathWithoutLastNode(parserContext.m_curPath);
+
+        // Only generate completions if the type is identityref so that we
+        // don't overwrite some other completions.
+        if (schema.leafType(location, {module, leaf.m_name}) == yang::LeafDataTypes::IdentityRef)
+            parserContext.m_suggestions = schema.validIdentities(location, {module, leaf.m_name}, Prefixes::WhenNeeded);
     }
 };