From 8a7e89cb6368b74916ef74797fa44baae66f9eff Mon Sep 17 00:00:00 2001
From: Bartek Kryza <bkryza@gmail.com>
Date: Mon, 18 Jul 2022 23:41:37 +0200
Subject: [PATCH] WIP

---
 .../plantuml/class_diagram_generator.cc       |  18 +-
 src/class_diagram/model/class.h               |  13 +-
 src/class_diagram/model/class_parent.h        |   7 +
 src/class_diagram/model/diagram.cc            |  64 +++--
 src/class_diagram/model/diagram.h             |  30 ++-
 .../visitor/translation_unit_visitor.cc       | 254 ++++++++++++++----
 .../visitor/translation_unit_visitor.h        |  37 +--
 src/common/generators/plantuml/generator.h    |  48 ++--
 src/common/model/diagram.h                    |  13 +-
 src/common/model/diagram_element.cc           |  22 +-
 src/common/model/diagram_element.h            |  10 +-
 src/common/model/diagram_filter.cc            |  42 ++-
 src/common/model/diagram_filter.h             |  26 +-
 src/common/model/element.cc                   |   4 -
 src/common/model/element.h                    |   7 +-
 src/common/model/package.h                    |  11 +-
 src/common/model/relationship.cc              |   9 +-
 src/common/model/relationship.h               |   9 +-
 src/common/model/source_file.h                |  15 +-
 .../types.h}                                  |  35 ++-
 .../plantuml/include_diagram_generator.cc     |  38 +--
 src/include_diagram/model/diagram.cc          |   9 +-
 src/include_diagram/model/diagram.h           |  13 +-
 .../visitor/translation_unit_context.cc.bak   |  63 -----
 .../plantuml/package_diagram_generator.cc     |   4 +-
 src/package_diagram/model/diagram.cc          |  14 +-
 src/package_diagram/model/diagram.h           |  17 +-
 .../visitor/element_visitor_context.cc.bak    |  43 ---
 .../visitor/translation_unit_context.cc.bak   | 194 -------------
 src/sequence_diagram/model/diagram.cc         |   2 +-
 src/sequence_diagram/model/diagram.h          |   3 +-
 .../visitor/translation_unit_context.cc.bak   |  65 -----
 32 files changed, 491 insertions(+), 648 deletions(-)
 rename src/{include_diagram/visitor/element_visitor_context.cc.bak => common/types.h} (54%)
 delete mode 100644 src/include_diagram/visitor/translation_unit_context.cc.bak
 delete mode 100644 src/package_diagram/visitor/element_visitor_context.cc.bak
 delete mode 100644 src/package_diagram/visitor/translation_unit_context.cc.bak
 delete mode 100644 src/sequence_diagram/visitor/translation_unit_context.cc.bak

diff --git a/src/class_diagram/generators/plantuml/class_diagram_generator.cc b/src/class_diagram/generators/plantuml/class_diagram_generator.cc
index b0922237..041b0ec9 100644
--- a/src/class_diagram/generators/plantuml/class_diagram_generator.cc
+++ b/src/class_diagram/generators/plantuml/class_diagram_generator.cc
@@ -246,7 +246,7 @@ void generator::generate_relationships(
 {
     namespace plantuml_common = clanguml::common::generators::plantuml;
 
-    const auto &uns = m_config.using_namespace();
+//    const auto &uns = m_config.using_namespace();
 
     //
     // Process relationships
@@ -264,14 +264,14 @@ void generator::generate_relationships(
             plantuml_common::to_plantuml(r.type(), r.style()));
 
         std::stringstream relstr;
-        std::string destination;
+        clanguml::common::id_t destination;
         try {
             destination = r.destination();
 
             // TODO: Refactor destination to a namespace qualified entity
             //       name
-            if (util::starts_with(destination, std::string{"::"}))
-                destination = destination.substr(2, destination.size());
+//            if (util::starts_with(destination, std::string{"::"}))
+//                destination = destination.substr(2, destination.size());
 
             LOG_DBG("=== Destination is: {}", destination);
 
@@ -284,8 +284,7 @@ void generator::generate_relationships(
             if (!r.multiplicity_destination().empty())
                 puml_relation += " \"" + r.multiplicity_destination() + "\"";
 
-            auto target_alias = m_model.to_alias(
-                m_config.using_namespace().relative(destination));
+            auto target_alias = m_model.to_alias(destination);
 
             if (m_generated_aliases.find(target_alias) ==
                 m_generated_aliases.end())
@@ -321,7 +320,7 @@ void generator::generate_relationships(
         for (const auto &b : c.parents()) {
             std::stringstream relstr;
             try {
-                auto target_alias = m_model.to_alias(uns.relative(b.name()));
+                auto target_alias = m_model.to_alias(b.id());
 
                 if (m_generated_aliases.find(target_alias) ==
                     m_generated_aliases.end())
@@ -369,13 +368,12 @@ void generator::generate_relationships(const enum_ &e, std::ostream &ostr) const
         if (!m_model.should_include(r.type()))
             continue;
 
-        std::string destination;
+        clanguml::common::id_t destination;
         std::stringstream relstr;
         try {
             destination = r.destination();
 
-            auto target_alias = m_model.to_alias(
-                m_config.using_namespace().relative(destination));
+            auto target_alias = m_model.to_alias(destination);
 
             if (m_generated_aliases.find(target_alias) ==
                 m_generated_aliases.end())
diff --git a/src/class_diagram/model/class.h b/src/class_diagram/model/class.h
index 757c5850..87086b20 100644
--- a/src/class_diagram/model/class.h
+++ b/src/class_diagram/model/class.h
@@ -23,11 +23,10 @@
 #include "common/model/element.h"
 #include "common/model/enums.h"
 #include "common/model/stylable_element.h"
+#include "common/types.h"
 #include "template_parameter.h"
 #include "type_alias.h"
 
-#include <type_safe/reference.hpp>
-
 #include <string>
 #include <vector>
 
@@ -108,14 +107,14 @@ private:
 
 namespace std {
 template <>
-struct hash<
-    type_safe::object_ref<const clanguml::class_diagram::model::class_>> {
-    std::size_t operator()(const type_safe::object_ref<
+struct hash<std::reference_wrapper<
+    const clanguml::class_diagram::model::class_>> {
+    std::size_t operator()(const std::reference_wrapper<
         const clanguml::class_diagram::model::class_> &key) const
     {
-        using clanguml::class_diagram::model::class_;
+        using clanguml::common::id_t;
 
-        return std::hash<std::string>{}(key.get().full_name(false));
+        return std::hash<id_t>{}(key.get().id());
     }
 };
 }
diff --git a/src/class_diagram/model/class_parent.h b/src/class_diagram/model/class_parent.h
index f51953d2..0ef20766 100644
--- a/src/class_diagram/model/class_parent.h
+++ b/src/class_diagram/model/class_parent.h
@@ -18,6 +18,7 @@
 #pragma once
 
 #include "common/model/enums.h"
+#include "common/types.h"
 
 #include <string>
 
@@ -28,6 +29,11 @@ public:
     void set_name(const std::string &name);
     std::string name() const;
 
+    clanguml::common::id_t id() const noexcept { return id_; }
+    void set_id(clanguml::common::id_t id) { id_ = id; }
+
+    void set_id(id_t id);
+
     void is_virtual(bool is_virtual);
     bool is_virtual() const;
 
@@ -35,6 +41,7 @@ public:
     common::model::access_t access() const;
 
 private:
+    clanguml::common::id_t id_;
     std::string name_;
     bool is_virtual_{false};
     common::model::access_t access_;
diff --git a/src/class_diagram/model/diagram.cc b/src/class_diagram/model/diagram.cc
index 725f3670..479c2ec5 100644
--- a/src/class_diagram/model/diagram.cc
+++ b/src/class_diagram/model/diagram.cc
@@ -25,12 +25,13 @@
 
 namespace clanguml::class_diagram::model {
 
-const std::vector<type_safe::object_ref<const class_>> &diagram::classes() const
+const std::vector<std::reference_wrapper<const class_>> &
+diagram::classes() const
 {
     return classes_;
 }
 
-const std::vector<type_safe::object_ref<const enum_>> &diagram::enums() const
+const std::vector<std::reference_wrapper<const enum_>> &diagram::enums() const
 {
     return enums_;
 }
@@ -40,10 +41,15 @@ common::model::diagram_t diagram::type() const
     return common::model::diagram_t::kClass;
 }
 
-type_safe::optional_ref<const clanguml::common::model::diagram_element>
+std::optional<
+    std::reference_wrapper<const clanguml::common::model::diagram_element>>
 diagram::get(const std::string &full_name) const
 {
-    type_safe::optional_ref<const clanguml::common::model::diagram_element> res;
+    // type_safe::optional_ref<const clanguml::common::model::diagram_element>
+    // res;
+    std::optional<
+        std::reference_wrapper<const clanguml::common::model::diagram_element>>
+        res;
 
     res = get_class(full_name);
 
@@ -67,7 +73,7 @@ bool diagram::has_enum(const enum_ &e) const
         [&e](const auto &ee) { return ee.get().full_name() == e.full_name(); });
 }
 
-type_safe::optional_ref<const class_> diagram::get_class(
+std::optional<std::reference_wrapper<const class_>> diagram::get_class(
     const std::string &name) const
 {
     for (const auto &c : classes_) {
@@ -76,10 +82,10 @@ type_safe::optional_ref<const class_> diagram::get_class(
         }
     }
 
-    return type_safe::nullopt;
+    return {};
 }
 
-type_safe::optional_ref<const enum_> diagram::get_enum(
+std::optional<std::reference_wrapper<const enum_>> diagram::get_enum(
     const std::string &name) const
 {
     for (const auto &e : enums_) {
@@ -88,7 +94,7 @@ type_safe::optional_ref<const enum_> diagram::get_enum(
         }
     }
 
-    return type_safe::nullopt;
+    return {};
 }
 
 void diagram::add_type_alias(std::unique_ptr<type_alias> &&ta)
@@ -131,7 +137,7 @@ bool diagram::add_class(std::unique_ptr<class_> &&c)
     auto name_and_ns = ns | name;
     const auto &cc = *c;
 
-    auto cc_ref = type_safe::ref(cc);
+    auto cc_ref = std::ref(cc);
 
     if (!has_class(cc)) {
         if (add_element(ns, std::move(c)))
@@ -159,7 +165,7 @@ bool diagram::add_enum(std::unique_ptr<enum_> &&e)
 
     assert(!util::contains(e->name(), "::"));
 
-    auto e_ref = type_safe::ref(*e);
+    auto e_ref = std::ref(*e);
     auto ns = e->get_relative_namespace();
 
     if (!has_enum(*e)) {
@@ -175,9 +181,9 @@ bool diagram::add_enum(std::unique_ptr<enum_> &&e)
 }
 
 void diagram::get_parents(
-    std::unordered_set<type_safe::object_ref<const class_>> &parents) const
+    clanguml::common::reference_set<class_> &parents) const
 {
-    bool found_new = false;
+    bool found_new{false};
     for (const auto &parent : parents) {
         for (const auto &pp : parent.get().parents()) {
             const auto p = get_class(pp.name());
@@ -194,25 +200,39 @@ void diagram::get_parents(
     }
 }
 
-std::string diagram::to_alias(const std::string &full_name) const
+bool diagram::has_element(
+    clanguml::common::model::diagram_element::id_t id) const
 {
-    LOG_DBG("Looking for alias for {}", full_name);
+    for (const auto &c : classes_) {
+        if (c.get().id() == id)
+            return true;
+    }
+
+    for (const auto &c : enums_) {
+        if (c.get().id() == id)
+            return true;
+    }
+
+    return false;
+}
+
+std::string diagram::to_alias(
+    clanguml::common::model::diagram_element::id_t id) const
+{
+    LOG_DBG("Looking for alias for {}", id);
 
     for (const auto &c : classes_) {
-        const auto &cc = c.get();
-        if (cc.full_name() == full_name) {
-            return c->alias();
+        if (c.get().id() == id) {
+            return c.get().alias();
         }
     }
 
     for (const auto &e : enums_) {
-        if (e.get().full_name() == full_name) {
-            return e->alias();
-        }
+        if (e.get().id() == id)
+            return e.get().alias();
     }
 
-    throw error::uml_alias_missing(
-        fmt::format("Missing alias for {}", full_name));
+    throw error::uml_alias_missing(fmt::format("Missing alias for {}", id));
 }
 
 }
diff --git a/src/class_diagram/model/diagram.h b/src/class_diagram/model/diagram.h
index 48d3887b..0a11476b 100644
--- a/src/class_diagram/model/diagram.h
+++ b/src/class_diagram/model/diagram.h
@@ -21,6 +21,7 @@
 #include "common/model/diagram.h"
 #include "common/model/nested_trait.h"
 #include "common/model/package.h"
+#include "common/types.h"
 #include "enum.h"
 #include "type_alias.h"
 
@@ -44,21 +45,22 @@ public:
 
     common::model::diagram_t type() const override;
 
-    type_safe::optional_ref<const clanguml::common::model::diagram_element> get(
-        const std::string &full_name) const override;
+    std::optional<
+        std::reference_wrapper<const clanguml::common::model::diagram_element>>
+    get(const std::string &full_name) const override;
 
-    const std::vector<type_safe::object_ref<const class_>> &classes() const;
+    const std::vector<std::reference_wrapper<const class_>> &classes() const;
 
-    const std::vector<type_safe::object_ref<const enum_>> &enums() const;
+    const std::vector<std::reference_wrapper<const enum_>> &enums() const;
 
     bool has_class(const class_ &c) const;
 
     bool has_enum(const enum_ &e) const;
 
-    type_safe::optional_ref<const class_> get_class(
+    std::optional<std::reference_wrapper<const class_>> get_class(
         const std::string &name) const;
 
-    type_safe::optional_ref<const enum_> get_enum(
+    std::optional<std::reference_wrapper<const enum_>> get_enum(
         const std::string &name) const;
 
     void add_type_alias(std::unique_ptr<type_alias> &&ta);
@@ -69,16 +71,22 @@ public:
 
     bool add_package(std::unique_ptr<common::model::package> &&p);
 
-    std::string to_alias(const std::string &full_name) const;
+    std::string to_alias(
+        clanguml::common::model::diagram_element::id_t id) const;
 
-    void get_parents(
-        std::unordered_set<type_safe::object_ref<const class_>> &parents) const;
+    void get_parents(clanguml::common::reference_set<class_> &parents) const;
 
     friend void print_diagram_tree(const diagram &d, const int level);
 
+    bool has_element(
+        const clanguml::common::model::diagram_element::id_t id) const override;
+
 private:
-    std::vector<type_safe::object_ref<const class_, false>> classes_;
-    std::vector<type_safe::object_ref<const enum_, false>> enums_;
+    std::vector<std::reference_wrapper<const class_>> classes_;
+    std::vector<std::reference_wrapper<const enum_>> enums_;
+
+    //    std::vector<type_safe::object_ref<const class_, false>> classes_;
+    //    std::vector<type_safe::object_ref<const enum_, false>> enums_;
     std::map<std::string, std::unique_ptr<type_alias>> type_aliases_;
 };
 }
diff --git a/src/class_diagram/visitor/translation_unit_visitor.cc b/src/class_diagram/visitor/translation_unit_visitor.cc
index 3fdd2129..e62f5c14 100644
--- a/src/class_diagram/visitor/translation_unit_visitor.cc
+++ b/src/class_diagram/visitor/translation_unit_visitor.cc
@@ -17,23 +17,8 @@
  */
 
 #include "translation_unit_visitor.h"
-
-//#include "cppast/cpp_function_type.hpp"
 #include "cx/util.h"
 
-//#include <cppast/cpp_alias_template.hpp>
-//#include <cppast/cpp_array_type.hpp>
-//#include <cppast/cpp_class_template.hpp>
-//#include <cppast/cpp_entity_kind.hpp>
-//#include <cppast/cpp_enum.hpp>
-//#include <cppast/cpp_friend.hpp>
-//#include <cppast/cpp_function_type.hpp>
-//#include <cppast/cpp_member_function.hpp>
-//#include <cppast/cpp_member_variable.hpp>
-//#include <cppast/cpp_namespace.hpp>
-//#include <cppast/cpp_template.hpp>
-//#include <cppast/cpp_type_alias.hpp>
-//#include <cppast/cpp_variable.hpp>
 #include <clang/Basic/FileManager.h>
 #include <spdlog/spdlog.h>
 
@@ -78,9 +63,11 @@ access_t access_specifier_to_access_t(clang::AccessSpecifier access_specifier)
 std::optional<clanguml::common::model::namespace_> get_enclosing_namespace(
     const clang::DeclContext *decl)
 {
+    if (!decl->getEnclosingNamespaceContext()->isNamespace())
+        return {};
+
     const auto *namespace_declaration =
-        static_cast<const clang::NamespaceDecl *>(
-            decl->getEnclosingNamespaceContext());
+        clang::cast<clang::NamespaceDecl>(decl->getEnclosingNamespaceContext());
 
     if (namespace_declaration == nullptr) {
         return {};
@@ -92,8 +79,22 @@ std::optional<clanguml::common::model::namespace_> get_enclosing_namespace(
 
 std::string to_string(const clang::QualType &type, const clang::ASTContext &ctx)
 {
-    clang::PrintingPolicy print_policy(ctx.getLangOpts());
-    return type.getAsString(print_policy);
+    const clang::PrintingPolicy print_policy(ctx.getLangOpts());
+
+    auto result{type.getAsString(print_policy)};
+
+    if (result.find('<') != std::string::npos) {
+        auto canonical_type_name =
+            type.getCanonicalType().getAsString(print_policy);
+
+        auto canonical_qualified_name =
+            canonical_type_name.substr(0, canonical_type_name.find('<'));
+        auto result_template_arguments = result.substr(result.find('<'));
+
+        result = canonical_qualified_name + result_template_arguments;
+    }
+
+    return result;
 }
 
 translation_unit_visitor::translation_unit_visitor(clang::SourceManager &sm,
@@ -115,6 +116,7 @@ bool translation_unit_visitor::VisitEnumDecl(clang::EnumDecl *enm)
     ns.pop_back();
     e.set_name(enm->getNameAsString());
     e.set_namespace(ns);
+    e.set_id(enm->getID());
 
     process_comment(*enm, e);
     set_source_location(*enm, e);
@@ -129,7 +131,7 @@ bool translation_unit_visitor::VisitEnumDecl(clang::EnumDecl *enm)
     }
 
     if (enm->getParent()->isRecord()) {
-        // process_record_containment(*enm, e);
+        process_record_containment(*enm, e);
     }
 
     auto namespace_declaration = detail::get_enclosing_namespace(enm);
@@ -143,15 +145,56 @@ bool translation_unit_visitor::VisitEnumDecl(clang::EnumDecl *enm)
     return true;
 }
 
+bool translation_unit_visitor::VisitClassTemplateDecl(
+    clang::ClassTemplateDecl *cls)
+{
+    if (source_manager_.isInSystemHeader(cls->getSourceRange().getBegin()))
+        return true;
+
+    auto c_ptr = process_class_declaration(cls->getTemplatedDecl());
+
+    c_ptr->set_id(cls->getID());
+
+    if (!c_ptr)
+        return true;
+
+    process_template_parameters(*cls, *c_ptr);
+
+    if (diagram_.should_include(*c_ptr)) {
+        diagram_.add_class(std::move(c_ptr));
+    }
+
+    return true;
+}
+
 bool translation_unit_visitor::VisitCXXRecordDecl(clang::CXXRecordDecl *cls)
 {
     if (source_manager_.isInSystemHeader(cls->getSourceRange().getBegin()))
         return true;
 
+    // Check if the class was already processed within VisitClassTemplateDecl()
+    if (diagram_.has_element(cls->getID()))
+        return true;
+
     // TODO: Add support for classes defined in function/method bodies
     if (cls->isLocalClass())
         return true;
 
+    auto c_ptr = process_class_declaration(cls);
+
+    if (!c_ptr)
+        return true;
+
+    if (diagram_.should_include(*c_ptr)) {
+        diagram_.add_class(std::move(c_ptr));
+    }
+
+    return true;
+}
+
+std::unique_ptr<class_> translation_unit_visitor::process_class_declaration(
+    clang::CXXRecordDecl *cls)
+{
     auto c_ptr = std::make_unique<class_>(config_.using_namespace());
     auto &c = *c_ptr;
 
@@ -161,12 +204,13 @@ bool translation_unit_visitor::VisitCXXRecordDecl(clang::CXXRecordDecl *cls)
     ns.pop_back();
     c.set_name(cls->getNameAsString());
     c.set_namespace(ns);
+    c.set_id(cls->getID());
 
     process_comment(*cls, c);
     set_source_location(*cls, c);
 
     if (c.skip())
-        return true;
+        return {};
 
     c.set_style(c.style_spec());
 
@@ -176,22 +220,82 @@ bool translation_unit_visitor::VisitCXXRecordDecl(clang::CXXRecordDecl *cls)
     // Process class bases
     process_class_bases(cls, c);
 
-    // Process class template arguments
-    //    if (cls->isTemplateDecl()) {
-    //        bool skip = process_template_parameters(cls, c, tspec);
-    //        if (skip)
-    //            return;
-    //    }
-
     if (cls->getParent()->isRecord()) {
         process_record_containment(*cls, c);
     }
 
-    if (diagram_.should_include(c)) {
-        diagram_.add_class(std::move(c_ptr));
+    return c_ptr;
+}
+
+bool translation_unit_visitor::process_template_parameters(
+    const clang::ClassTemplateDecl &template_declaration, class_ &c)
+{
+    LOG_DBG("Processing class {} template parameters...",
+        template_declaration.getQualifiedNameAsString());
+
+    if (template_declaration.getTemplateParameters() == nullptr)
+        return false;
+
+    for (const auto *parameter :
+        *template_declaration.getTemplateParameters()) {
+        if (clang::dyn_cast_or_null<clang::TemplateTypeParmDecl>(parameter)) {
+            const auto *template_type_parameter =
+                clang::dyn_cast_or_null<clang::TemplateTypeParmDecl>(parameter);
+            template_parameter ct;
+            ct.set_type("");
+            ct.is_template_parameter(true);
+            ct.set_name(template_type_parameter->getNameAsString());
+            ct.set_default_value("");
+            ct.is_variadic(template_type_parameter->isParameterPack());
+
+            c.add_template(std::move(ct));
+        }
+        else if (clang::dyn_cast_or_null<clang::NonTypeTemplateParmDecl>(
+                     parameter)) {
+            const auto *template_nontype_parameter =
+                clang::dyn_cast_or_null<clang::NonTypeTemplateParmDecl>(
+                    parameter);
+            template_parameter ct;
+            ct.set_type("");
+            ct.is_template_parameter(false);
+            ct.set_name(template_nontype_parameter->getNameAsString());
+            ct.set_default_value("");
+            ct.is_variadic(template_nontype_parameter->isParameterPack());
+
+            c.add_template(std::move(ct));
+        }
     }
 
-    return true;
+    /*
+    auto scope = cppast::cpp_scope_name(type_safe::ref(cls));
+    // Even if this is a template the scope.is_templated() returns
+    // false when the template parameter list is empty
+    if (scope.is_templated()) {
+        process_scope_template_parameters(c, scope);
+    }
+    else {
+        LOG_DBG("Class {} is templated but it's scope {} is not - "
+                "probably this is a specialization",
+            cls.name(), scope.name());
+
+        // Add specialization arguments
+        if (tspec) {
+            if (!tspec.value().arguments_exposed()) {
+                process_unexposed_template_specialization_parameters(tspec, c);
+            }
+            else {
+                process_exposed_template_specialization_parameters(tspec, c);
+            }
+        }
+        else {
+            LOG_DBG("Skipping template class declaration which has only "
+                    "unexposed arguments but no tspec provided");
+            return true;
+        }
+    }
+     */
+
+    return false;
 }
 
 void translation_unit_visitor::process_record_containment(
@@ -210,9 +314,10 @@ void translation_unit_visitor::process_record_containment(
         element.set_namespace(namespace_declaration.value());
     }
 
-    static_cast<const clang::RecordDecl *>(record.getParent())->getID();
+    const auto id =
+        static_cast<const clang::RecordDecl *>(record.getParent())->getID();
 
-    element.add_relationship({relationship_t::kContainment, parent_name});
+    element.add_relationship({relationship_t::kContainment, id});
 }
 
 void translation_unit_visitor::process_class_bases(
@@ -224,6 +329,8 @@ void translation_unit_visitor::process_class_bases(
             to_string(base.getType(), cls->getASTContext())};
 
         cp.set_name(name_and_ns.to_string());
+        cp.set_id(
+            base.getType()->getAs<clang::RecordType>()->getDecl()->getID());
         cp.is_virtual(base.isVirtual());
 
         cp.set_access(
@@ -304,6 +411,9 @@ bool translation_unit_visitor::find_relationships(const clang::QualType &type,
 {
     bool result = false;
 
+    std::string type_name = type.getAsString();
+    (void)type_name;
+
     if (type->isPointerType()) {
         relationship_hint = relationship_t::kAssociation;
         find_relationships(
@@ -412,13 +522,9 @@ void translation_unit_visitor::add_relationships(class_ &c,
     std::optional<std::string> label)
 {
     for (const auto &[target, relationship_type] : relationships) {
-        if (target.empty())
-            continue;
-
-        auto [target_ns, target_name] = cx::util::split_ns(target);
-        if (diagram().should_include(target_ns, target_name) &&
-            (relationship_type != relationship_t::kNone) &&
-            (target != c.name_and_ns())) {
+        if (diagram().has_element(target) &&
+            (relationship_type != relationship_t::kNone) /*&&
+            (target != c.name_and_ns())*/) {
             relationship r{relationship_type, target};
             if (label)
                 r.set_label(label.value());
@@ -453,9 +559,8 @@ void translation_unit_visitor::process_static_field(
         return;
 
     if (!field.skip_relationship()) {
-        std::vector<
-            std::pair<std::string, clanguml::common::model::relationship_t>>
-            relationships;
+        found_relationships_t relationships;
+
         // find relationship for the type
         find_relationships(field_declaration.getType(), relationships,
             relationship_t::kAssociation);
@@ -468,6 +573,50 @@ void translation_unit_visitor::process_static_field(
     c.add_member(std::move(field));
 }
 
+std::unique_ptr<class_> translation_unit_visitor::build_template_instantiation(
+    const clang::TemplateSpecializationType &template_type,
+    std::optional<clanguml::class_diagram::model::class_ *> parent)
+{
+    //
+    // Create class_ instance to hold the template instantiation
+    //
+    auto template_instantiation_ptr =
+        std::make_unique<class_>(config_.using_namespace());
+    auto &template_instantiation = *template_instantiation_ptr;
+    std::string full_template_specialization_name =
+        to_string(template_type.desugar(),
+            template_type.getAsCXXRecordDecl()->getASTContext());
+
+    const auto *template_record{
+        template_type.getTemplateName().getAsTemplateDecl()};
+    std::string qualified_name = template_record->getQualifiedNameAsString();
+    namespace_ ns{qualified_name};
+    ns.pop_back();
+    template_instantiation.set_name(template_record->getNameAsString());
+    template_instantiation.set_namespace(ns);
+    template_instantiation.set_id(template_record->getID() +
+        std::hash<std::string>{}(full_template_specialization_name));
+
+    for (const auto &arg : template_type) {
+        auto argument_kind = arg.getKind();
+        if (argument_kind == clang::TemplateArgument::ArgKind::Type) {
+            const auto *argument_record_decl =
+                arg.getAsType()->getAsRecordDecl();
+            template_parameter argument;
+            argument.is_template_parameter(false);
+            argument.set_name(to_string(
+                arg.getAsType(), argument_record_decl->getASTContext()));
+
+            template_instantiation.add_template(std::move(argument));
+
+            template_instantiation.add_relationship(
+                {relationship_t::kInstantiation, template_record->getID()});
+        }
+    }
+
+    return template_instantiation_ptr;
+}
+
 void translation_unit_visitor::process_field(
     const clang::FieldDecl &field_declaration, class_ &c)
 {
@@ -481,6 +630,10 @@ void translation_unit_visitor::process_field(
         detail::access_specifier_to_access_t(field_declaration.getAccess()),
         field_declaration.getNameAsString(), type_name};
 
+    if (field.name() == "e") {
+        LOG_DBG("EEEEEEEEE");
+    }
+
     process_comment(field_declaration, field);
     set_source_location(field_declaration, field);
 
@@ -488,9 +641,7 @@ void translation_unit_visitor::process_field(
         return;
 
     if (!field.skip_relationship()) {
-        std::vector<
-            std::pair<std::string, clanguml::common::model::relationship_t>>
-            relationships;
+        found_relationships_t relationships;
         // find relationship for the type
         find_relationships(field_declaration.getType(), relationships,
             relationship_t::kAggregation);
@@ -500,6 +651,17 @@ void translation_unit_visitor::process_field(
             field_declaration.getNameAsString());
     }
 
+    if (field_type->getAs<clang::TemplateSpecializationType>() &&
+        diagram().should_include(type_name)) {
+
+        auto template_specialization_ptr = build_template_instantiation(
+            *field_type->getAs<clang::TemplateSpecializationType>());
+
+        diagram().add_class(std::move(template_specialization_ptr));
+
+        LOG_DBG("ADDED TEMPLATE SPECIALIZATION TO DIAGRAM");
+    }
+
     /*
     if (field_type->getAsCXXRecordDecl()) {
 
diff --git a/src/class_diagram/visitor/translation_unit_visitor.h b/src/class_diagram/visitor/translation_unit_visitor.h
index b82596b6..3d59c92e 100644
--- a/src/class_diagram/visitor/translation_unit_visitor.h
+++ b/src/class_diagram/visitor/translation_unit_visitor.h
@@ -17,28 +17,14 @@
  */
 #pragma once
 
+#include "class_diagram/model/class.h"
 #include "class_diagram/model/diagram.h"
-//#include "class_diagram/visitor/translation_unit_context.h"
 #include "common/model/enums.h"
 #include "config/config.h"
 
-//#include <clang-c/CXCompilationDatabase.h>
-//#include <clang-c/Index.h>
-//#include <cppast/cpp_friend.hpp>
-//#include <cppast/cpp_function_template.hpp>
-//#include <cppast/cpp_member_function.hpp>
-//#include <cppast/cpp_member_variable.hpp>
-//#include <cppast/cpp_template.hpp>
-//#include <cppast/cpp_template_parameter.hpp>
-//#include <cppast/cpp_type.hpp>
-//#include <cppast/visitor.hpp>
 #include <clang/AST/RecursiveASTVisitor.h>
 #include <clang/Basic/SourceManager.h>
-#include <class_diagram/model/class.h>
-#include <common/model/enums.h>
-#include <type_safe/reference.hpp>
-//#include <cppast/cpp_alias_template.hpp>
-//#include <cppast/cpp_type_alias.hpp>
+
 #include <deque>
 #include <functional>
 #include <map>
@@ -48,7 +34,7 @@
 namespace clanguml::class_diagram::visitor {
 
 using found_relationships_t =
-    std::vector<std::pair<clanguml::common::model::element::id_t,
+    std::vector<std::pair<clanguml::common::model::diagram_element::id_t,
         common::model::relationship_t>>;
 
 class translation_unit_visitor
@@ -62,17 +48,27 @@ public:
 
     virtual bool VisitEnumDecl(clang::EnumDecl *e);
 
+    virtual bool VisitClassTemplateDecl(
+        clang::ClassTemplateDecl *class_template_declaration);
+
     //    virtual bool VisitVarDecl(clang::VarDecl *variable_declaration);
     clanguml::class_diagram::model::diagram &diagram() { return diagram_; }
     //    void operator()();
 
 private:
+    std::unique_ptr<clanguml::class_diagram::model::class_>
+    process_class_declaration(clang::CXXRecordDecl *cls);
+
     void process_class_bases(const clang::CXXRecordDecl *cls,
         clanguml::class_diagram::model::class_ &c) const;
 
     void process_class_children(const clang::CXXRecordDecl *cls,
         clanguml::class_diagram::model::class_ &c);
 
+    bool process_template_parameters(
+        const clang::ClassTemplateDecl &template_declaration,
+        clanguml::class_diagram::model::class_ &c);
+
     void process_record_containment(const clang::TagDecl &record,
         clanguml::common::model::element &c) const;
 
@@ -90,7 +86,8 @@ private:
         clanguml::class_diagram::model::class_ &c,
         const std::set<std::string> &template_parameter_names = {});
 
-    bool find_relationships(const clang::QualType &type, &,
+    bool find_relationships(const clang::QualType &type,
+        found_relationships_t &,
         clanguml::common::model::relationship_t relationship_hint);
 
     void add_relationships(clanguml::class_diagram::model::class_ &c,
@@ -101,6 +98,10 @@ private:
     void set_source_location(const clang::Decl &decl,
         clanguml::common::model::source_location &element);
 
+    std::unique_ptr<clanguml::class_diagram::model::class_>
+    build_template_instantiation(const clang::TemplateSpecializationType& template_type,
+        std::optional<clanguml::class_diagram::model::class_ *> parent = {});
+
     template <typename ClangDecl>
     void process_comment(
         const ClangDecl &decl, clanguml::common::model::decorated_element &e)
diff --git a/src/common/generators/plantuml/generator.h b/src/common/generators/plantuml/generator.h
index 20743787..04b2ff43 100644
--- a/src/common/generators/plantuml/generator.h
+++ b/src/common/generators/plantuml/generator.h
@@ -23,9 +23,9 @@
 #include "util/util.h"
 
 //#include <cppast/libclang_parser.hpp>
+#include <clang/Frontend/CompilerInstance.h>
 #include <clang/Tooling/CompilationDatabase.h>
 #include <clang/Tooling/Tooling.h>
-#include <clang/Frontend/CompilerInstance.h>
 #include <glob/glob.hpp>
 #include <inja/inja.hpp>
 
@@ -146,23 +146,28 @@ void generator<C, D>::generate_config_layout_hints(std::ostream &ostr) const
 {
     using namespace clanguml::util;
 
-    const auto &uns = m_config.using_namespace();
+    //    const auto &uns = m_config.using_namespace();
 
     // Generate layout hints
-    for (const auto &[entity, hints] : m_config.layout()) {
+    for (const auto &[entity_name, hints] : m_config.layout()) {
         for (const auto &hint : hints) {
             std::stringstream hint_str;
             try {
-                hint_str << m_model.to_alias(uns.relative(entity))
-                         << " -[hidden]"
+                auto element_opt = m_model.get(
+                    m_config.using_namespace().relative(entity_name));
+                auto hint_element_opt = m_model.get(
+                    m_config.using_namespace().relative(hint.entity));
+                if (!element_opt || !hint_element_opt)
+                    continue;
+                hint_str << element_opt.value().get().alias() << " -[hidden]"
                          << clanguml::config::to_string(hint.hint) << "- "
-                         << m_model.to_alias(uns.relative(hint.entity)) << '\n';
+                         << hint_element_opt.value().get().alias() << '\n';
                 ostr << hint_str.str();
             }
             catch (clanguml::error::uml_alias_missing &e) {
                 LOG_DBG("=== Skipping layout hint from {} to {} due "
                         "to: {}",
-                    entity, hint.entity, e.what());
+                    entity_name, hint.entity, e.what());
             }
         }
     }
@@ -181,10 +186,17 @@ void generator<C, D>::generate_plantuml_directives(
         // Now search for alias @A() directives in the text
         std::tuple<std::string, size_t, size_t> alias_match;
         while (util::find_element_alias(directive, alias_match)) {
-            auto alias = m_model.to_alias(
-                m_config.using_namespace().relative(std::get<0>(alias_match)));
-            directive.replace(
-                std::get<1>(alias_match), std::get<2>(alias_match), alias);
+            const auto full_name =
+                m_config.using_namespace() | std::get<0>(alias_match);
+            auto element_opt = m_model.get(full_name.to_string());
+
+            if (element_opt)
+                directive.replace(std::get<1>(alias_match),
+                    std::get<2>(alias_match),
+                    element_opt.value().get().alias());
+            else
+                directive.replace(std::get<1>(alias_match),
+                    std::get<2>(alias_match), "UNKNOWN_ALIAS");
         }
         ostr << directive << '\n';
     }
@@ -242,8 +254,8 @@ public:
 
     virtual void HandleTranslationUnit(clang::ASTContext &ast_context)
     {
-//        const auto* tud = ast_context.getTranslationUnitDecl();
-////        tud->dump();
+        //        const auto* tud = ast_context.getTranslationUnitDecl();
+        ////        tud->dump();
         visitor_.TraverseDecl(ast_context.getTranslationUnitDecl());
     }
 };
@@ -402,9 +414,11 @@ template <typename C, typename D> void generator<C, D>::init_env()
     // is equivalent to the old syntax:
     //   "note left of @A(ClassA): This is a note"
     m_env.add_callback("alias", 1, [this](inja::Arguments &args) {
-        auto alias_match = args[0]->get<std::string>();
-        return m_model.to_alias(
-            m_config.using_namespace().relative(alias_match));
+        auto alias_match =
+            m_config.using_namespace() | args[0]->get<std::string>();
+        auto element_opt = m_model.get(alias_match.to_string());
+
+        return element_opt.value().get().alias();
     });
 
     m_env.add_callback("comment", 1, [this](inja::Arguments &args) {
@@ -419,7 +433,7 @@ template <typename C, typename D> void generator<C, D>::init_env()
         }
 
         if (element.has_value()) {
-            auto comment = element.value().comment();
+            auto comment = element.value().get().comment();
 
             if (comment.has_value())
                 res = comment.value();
diff --git a/src/common/model/diagram.h b/src/common/model/diagram.h
index b49bd953..8e48d99a 100644
--- a/src/common/model/diagram.h
+++ b/src/common/model/diagram.h
@@ -22,8 +22,6 @@
 #include "namespace.h"
 #include "source_file.h"
 
-#include <type_safe/optional_ref.hpp>
-
 #include <memory>
 #include <string>
 
@@ -40,8 +38,9 @@ public:
 
     virtual diagram_t type() const = 0;
 
-    virtual type_safe::optional_ref<const diagram_element> get(
-        const std::string &full_name) const = 0;
+    virtual std::optional<
+        std::reference_wrapper<const clanguml::common::model::diagram_element>>
+    get(const std::string &full_name) const = 0;
 
     diagram(const diagram &) = delete;
     diagram(diagram &&);
@@ -63,8 +62,12 @@ public:
     bool should_include(const relationship r) const;
     bool should_include(const relationship_t r) const;
     bool should_include(const access_t s) const;
+    virtual bool has_element(const diagram_element::id_t id) const
+    {
+        return false;
+    }
 
-    bool should_include(const namespace_ &ns, const std::string &name) const;
+    virtual bool should_include(const namespace_ &ns, const std::string &name) const;
 
 private:
     std::string name_;
diff --git a/src/common/model/diagram_element.cc b/src/common/model/diagram_element.cc
index bd3179f3..6dbb1161 100644
--- a/src/common/model/diagram_element.cc
+++ b/src/common/model/diagram_element.cc
@@ -27,26 +27,25 @@ namespace clanguml::common::model {
 std::atomic_uint64_t diagram_element::m_nextId = 1;
 
 diagram_element::diagram_element()
-    : m_id{m_nextId++}
+    : id_{0}
 {
 }
 
+diagram_element::id_t diagram_element::id() const { return id_; }
+
+void diagram_element::set_id(diagram_element::id_t id) { id_ = id; }
+
 std::string diagram_element::alias() const
 {
-    return fmt::format("C_{:010}", m_id);
+    assert(id_ >= 0);
+
+    return fmt::format("C_{:022}", id_);
 }
 
 void diagram_element::add_relationship(relationship &&cr)
 {
-    if (cr.destination().empty()) {
-        LOG_DBG("Skipping relationship '{}' - {} - '{}' due empty "
-                "destination",
-            cr.destination(), to_string(cr.type()), full_name(true));
-        return;
-    }
-
     if ((cr.type() == relationship_t::kInstantiation) &&
-        (cr.destination() == full_name(true))) {
+        (cr.destination() == id())) {
         LOG_DBG("Skipping self instantiation relationship for {}",
             cr.destination());
         return;
@@ -86,7 +85,8 @@ inja::json diagram_element::context() const
 
 bool operator==(const diagram_element &l, const diagram_element &r)
 {
-    return l.full_name(false) == r.full_name(false);
+    return l.id() == r.id();
+    //return l.full_name(false) == r.full_name(false);
 }
 
 std::ostream &operator<<(std::ostream &out, const diagram_element &rhs)
diff --git a/src/common/model/diagram_element.h b/src/common/model/diagram_element.h
index 5fa49a70..547f6a4f 100644
--- a/src/common/model/diagram_element.h
+++ b/src/common/model/diagram_element.h
@@ -32,10 +32,16 @@ namespace clanguml::common::model {
 
 class diagram_element : public decorated_element {
 public:
+    using id_t = int64_t;
+
     diagram_element();
 
     virtual ~diagram_element() = default;
 
+    id_t id() const;
+
+    void set_id(id_t id);
+
     std::string alias() const;
 
     void set_name(const std::string &name) { name_ = name; }
@@ -59,10 +65,8 @@ public:
 
     virtual inja::json context() const;
 
-protected:
-    const uint64_t m_id{0};
-
 private:
+    id_t id_;
     std::string name_;
     std::vector<relationship> relationships_;
 
diff --git a/src/common/model/diagram_filter.cc b/src/common/model/diagram_filter.cc
index e37392ff..65612561 100644
--- a/src/common/model/diagram_filter.cc
+++ b/src/common/model/diagram_filter.cc
@@ -28,59 +28,59 @@ namespace clanguml::common::model {
 namespace detail {
 
 template <>
-const std::vector<type_safe::object_ref<const class_diagram::model::class_>> &
-view(const class_diagram::model::diagram &d)
+const clanguml::common::reference_vector<class_diagram::model::class_> &view(
+    const class_diagram::model::diagram &d)
 {
     return d.classes();
 }
 
 template <>
-const std::vector<type_safe::object_ref<const class_diagram::model::enum_>> &
-view(const class_diagram::model::diagram &d)
+const clanguml::common::reference_vector<class_diagram::model::enum_> &view(
+    const class_diagram::model::diagram &d)
 {
     return d.enums();
 }
 
 template <>
-const std::vector<type_safe::object_ref<const common::model::package>> &view(
+const clanguml::common::reference_vector<common::model::package> &view(
     const package_diagram::model::diagram &d)
 {
     return d.packages();
 }
 
 template <>
-const std::vector<type_safe::object_ref<const common::model::source_file>> &
-view(const include_diagram::model::diagram &d)
+const clanguml::common::reference_vector<common::model::source_file> &view(
+    const include_diagram::model::diagram &d)
 {
     return d.files();
 }
 
 template <>
-const type_safe::optional_ref<const class_diagram::model::class_> get(
+const clanguml::common::optional_ref<class_diagram::model::class_> get(
     const class_diagram::model::diagram &d, const std::string &full_name)
 {
     return d.get_class(full_name);
 }
 
 template <>
-const type_safe::optional_ref<const common::model::package> get(
+const clanguml::common::optional_ref<const common::model::package> get(
     const package_diagram::model::diagram &d, const std::string &full_name)
 {
     return d.get_package(full_name);
 }
 
 template <>
-const type_safe::optional_ref<const common::model::source_file> get(
+const clanguml::common::optional_ref<const common::model::source_file> get(
     const include_diagram::model::diagram &d, const std::string &full_name)
 {
     return d.get_file(full_name);
 }
 
 template <>
-std::string destination_comparator<common::model::source_file>(
+clanguml::common::id_t destination_comparator<common::model::source_file>(
     const common::model::source_file &f)
 {
-    return f.alias();
+    return f.id();
 }
 } // namespace detail
 
@@ -221,9 +221,7 @@ tvl::value_t subclass_filter::match(const diagram &d, const element &e) const
     const auto &cd = dynamic_cast<const class_diagram::model::diagram &>(d);
 
     // First get all parents of element e
-    std::unordered_set<
-        type_safe::object_ref<const class_diagram::model::class_, false>>
-        parents;
+    clanguml::common::reference_set<class_diagram::model::class_> parents;
 
     const auto &fn = e.full_name(false);
     auto class_ref = cd.get_class(fn);
@@ -296,26 +294,25 @@ tvl::value_t context_filter::match(const diagram &d, const element &e) const
 
             if (context_root.has_value()) {
                 // This is a direct match to the context root
-                if (context_root.value().full_name(false) == e.full_name(false))
+                if (context_root.value().get().id() == e.id())
                     return true;
 
                 // Return a positive match if the element e is in a direct
                 // relationship with any of the context_root's
                 for (const relationship &rel :
-                    context_root.value().relationships()) {
-                    if (rel.destination() == e.full_name(false))
+                    context_root.value().get().relationships()) {
+                    if (rel.destination() == e.id())
                         return true;
                 }
                 for (const relationship &rel : e.relationships()) {
-                    if (rel.destination() ==
-                        context_root.value().full_name(false))
+                    if (rel.destination() == context_root.value().get().id())
                         return true;
                 }
 
                 // Return a positive match if the context_root is a parent
                 // of the element
                 for (const class_diagram::model::class_parent &p :
-                    context_root.value().parents()) {
+                    context_root.value().get().parents()) {
                     if (p.name() == e.full_name(false))
                         return true;
                 }
@@ -324,7 +321,8 @@ tvl::value_t context_filter::match(const diagram &d, const element &e) const
                     for (const class_diagram::model::class_parent &p :
                         static_cast<const class_diagram::model::class_ &>(e)
                             .parents()) {
-                        if (p.name() == context_root.value().full_name(false))
+                        if (p.name() ==
+                            context_root.value().get().full_name(false))
                             return true;
                     }
                 }
diff --git a/src/common/model/diagram_filter.h b/src/common/model/diagram_filter.h
index 2e27815f..4e084d6e 100644
--- a/src/common/model/diagram_filter.h
+++ b/src/common/model/diagram_filter.h
@@ -37,21 +37,20 @@ enum filter_t { kInclusive, kExclusive };
 
 namespace detail {
 template <typename ElementT, typename DiagramT>
-const std::vector<type_safe::object_ref<const ElementT>> &view(
-    const DiagramT &d);
+const clanguml::common::reference_vector<ElementT> &view(const DiagramT &d);
 
 template <typename ElementT, typename DiagramT>
-const type_safe::optional_ref<const ElementT> get(
+const clanguml::common::optional_ref<ElementT> get(
     const DiagramT &d, const std::string &full_name);
 
-template <typename ElementT>
-std::string destination_comparator(const ElementT &e)
+template <typename ElementT> int64_t destination_comparator(const ElementT &e)
 {
-    return e.full_name(false);
+    return e.id();
 }
 
 template <>
-std::string destination_comparator(const common::model::source_file &f);
+clanguml::common::id_t destination_comparator(
+    const common::model::source_file &f);
 } // namespace detail
 
 class filter_visitor {
@@ -177,7 +176,7 @@ struct edge_traversal_filter : public filter_visitor {
         // Now check if the e element is contained in the calculated set
         return std::any_of(matching_elements_.begin(), matching_elements_.end(),
             [&e](const auto &te) {
-                return te->full_name(false) == e.full_name(false);
+                return te.get().full_name(false) == e.full_name(false);
             });
     }
 
@@ -192,12 +191,12 @@ private:
             //  Check if any of its relationships of type relationship_
             //  points to an element already in the matching_elements_
             //  set
-            for (const auto &rel : from_el->relationships()) {
+            for (const auto &rel : from_el.get().relationships()) {
                 // Consider only if connected by one of specified relationships
                 if (util::contains(relationships, rel.type())) {
                     for (const auto &to_el : to) {
                         if (rel.destination() ==
-                            detail::destination_comparator(*to_el)) {
+                            detail::destination_comparator(to_el.get())) {
                             const auto &to_add = forward_ ? to_el : from_el;
                             if (matching_elements_.insert(to_add).second)
                                 added_new_element = true;
@@ -220,9 +219,9 @@ private:
                     cd, element.get().path().to_string());
 
                 while (parent.has_value()) {
-                    parents.emplace(type_safe::ref(parent.value()));
+                    parents.emplace(std::ref(parent.value()));
                     parent = detail::get<ElementT, DiagramT>(
-                        cd, parent.value().path().to_string());
+                        cd, parent.value().get().path().to_string());
                 }
             });
 
@@ -272,8 +271,7 @@ private:
     std::vector<std::string> roots_;
     relationship_t relationship_;
     mutable bool initialized_{false};
-    mutable std::unordered_set<type_safe::object_ref<const ElementT, false>>
-        matching_elements_;
+    mutable clanguml::common::reference_set<ElementT> matching_elements_;
     bool forward_;
 };
 
diff --git a/src/common/model/element.cc b/src/common/model/element.cc
index d67eb43a..dd789d92 100644
--- a/src/common/model/element.cc
+++ b/src/common/model/element.cc
@@ -29,10 +29,6 @@ element::element(const namespace_ &using_namespace)
 {
 }
 
-element::id_t element::id() const { return id_; }
-
-void element::set_id(element::id_t id) { id_ = id; }
-
 void element::set_using_namespaces(const namespace_ &un)
 {
     using_namespace_ = un;
diff --git a/src/common/model/element.h b/src/common/model/element.h
index ef2e4a0e..00db43b5 100644
--- a/src/common/model/element.h
+++ b/src/common/model/element.h
@@ -34,16 +34,12 @@ namespace clanguml::common::model {
 
 class element : public diagram_element, public source_location {
 public:
-    using id_t = int64_t;
+
 
     element(const namespace_ &using_namespace);
 
     virtual ~element() = default;
 
-    id_t id() const;
-
-    void set_id(id_t id);
-
     std::string name_and_ns() const
     {
         auto ns = ns_ | name();
@@ -77,7 +73,6 @@ public:
     inja::json context() const override;
 
 private:
-    id_t id_;
     namespace_ ns_;
     namespace_ using_namespace_;
 };
diff --git a/src/common/model/package.h b/src/common/model/package.h
index 1de5e675..6ae4c702 100644
--- a/src/common/model/package.h
+++ b/src/common/model/package.h
@@ -20,6 +20,7 @@
 #include "common/model/element.h"
 #include "common/model/nested_trait.h"
 #include "common/model/stylable_element.h"
+#include "common/types.h"
 #include "util/util.h"
 
 #include <spdlog/spdlog.h>
@@ -57,14 +58,14 @@ private:
 
 namespace std {
 template <>
-struct hash<type_safe::object_ref<const clanguml::common::model::package>> {
+struct hash<std::reference_wrapper<const clanguml::common::model::package>> {
     std::size_t operator()(
-        const type_safe::object_ref<const clanguml::common::model::package>
+        const std::reference_wrapper<const clanguml::common::model::package>
             &key) const
     {
-        using clanguml::common::model::package;
+        using clanguml::common::id_t;
 
-        return std::hash<std::string>{}(key.get().full_name(false));
+        return std::hash<id_t>{}(key.get().id());
     }
 };
-}
+}
\ No newline at end of file
diff --git a/src/common/model/relationship.cc b/src/common/model/relationship.cc
index 9ead05e9..f0d8ec86 100644
--- a/src/common/model/relationship.cc
+++ b/src/common/model/relationship.cc
@@ -20,7 +20,7 @@
 
 namespace clanguml::common::model {
 
-relationship::relationship(relationship_t type, const std::string &destination,
+relationship::relationship(relationship_t type, int64_t destination,
     access_t access, const std::string &label,
     const std::string &multiplicity_source,
     const std::string &multiplicity_destination)
@@ -37,12 +37,15 @@ void relationship::set_type(relationship_t type) noexcept { type_ = type; }
 
 relationship_t relationship::type() const noexcept { return type_; }
 
-void relationship::set_destination(const std::string &destination)
+void relationship::set_destination(int64_t destination)
 {
     destination_ = destination;
 }
 
-std::string relationship::destination() const { return destination_; }
+clanguml::common::id_t relationship::destination() const
+{
+    return destination_;
+}
 
 void relationship::set_multiplicity_source(
     const std::string &multiplicity_source)
diff --git a/src/common/model/relationship.h b/src/common/model/relationship.h
index 3b154223..d85fbc8a 100644
--- a/src/common/model/relationship.h
+++ b/src/common/model/relationship.h
@@ -19,6 +19,7 @@
 
 #include "common/model/decorated_element.h"
 #include "common/model/stylable_element.h"
+#include "common/types.h"
 
 #include <string>
 
@@ -27,7 +28,7 @@ namespace clanguml::common::model {
 class relationship : public common::model::decorated_element,
                      public common::model::stylable_element {
 public:
-    relationship(relationship_t type, const std::string &destination,
+    relationship(relationship_t type, int64_t destination,
         access_t access = access_t::kPublic, const std::string &label = "",
         const std::string &multiplicity_source = "",
         const std::string &multiplicity_destination = "");
@@ -37,8 +38,8 @@ public:
     void set_type(relationship_t type) noexcept;
     relationship_t type() const noexcept;
 
-    void set_destination(const std::string &destination);
-    std::string destination() const;
+    void set_destination(int64_t destination);
+    clanguml::common::id_t destination() const;
 
     void set_multiplicity_source(const std::string &multiplicity_source);
     std::string multiplicity_source() const;
@@ -57,7 +58,7 @@ public:
 
 private:
     relationship_t type_;
-    std::string destination_;
+    int64_t destination_;
     std::string multiplicity_source_;
     std::string multiplicity_destination_;
     std::string label_;
diff --git a/src/common/model/source_file.h b/src/common/model/source_file.h
index be99de07..a99599af 100644
--- a/src/common/model/source_file.h
+++ b/src/common/model/source_file.h
@@ -22,6 +22,7 @@
 #include "common/model/path.h"
 #include "common/model/source_location.h"
 #include "common/model/stylable_element.h"
+#include "common/types.h"
 #include "util/util.h"
 
 #include <spdlog/spdlog.h>
@@ -129,15 +130,19 @@ template <> struct hash<clanguml::common::model::filesystem_path> {
     }
 };
 
+}
+
+namespace std {
 template <>
-struct hash<type_safe::object_ref<const clanguml::common::model::source_file>> {
+struct hash<
+    std::reference_wrapper<const clanguml::common::model::source_file>> {
     std::size_t operator()(
-        const type_safe::object_ref<const clanguml::common::model::source_file>
+        const std::reference_wrapper<const clanguml::common::model::source_file>
             &key) const
     {
-        using clanguml::common::model::source_file;
+        using clanguml::common::id_t;
 
-        return std::hash<std::string>{}(key.get().full_name(false));
+        return std::hash<id_t>{}(key.get().id());
     }
 };
-}
+}
\ No newline at end of file
diff --git a/src/include_diagram/visitor/element_visitor_context.cc.bak b/src/common/types.h
similarity index 54%
rename from src/include_diagram/visitor/element_visitor_context.cc.bak
rename to src/common/types.h
index 05118ab7..8d3329e6 100644
--- a/src/include_diagram/visitor/element_visitor_context.cc.bak
+++ b/src/common/types.h
@@ -1,5 +1,5 @@
 /**
- * src/include_diagram/model/visitor/element_visitor_context.cc
+ * src/class_diagram/visitor/translation_unit_visitor.h
  *
  * Copyright (c) 2021-2022 Bartek Kryza <bkryza@gmail.com>
  *
@@ -15,29 +15,24 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#pragma once
 
-#include "element_visitor_context.h"
+#include <cstdint>
+#include <optional>
+#include <unordered_set>
+#include <vector>
 
-#include "translation_unit_context.h"
+namespace clanguml::common {
 
-namespace clanguml::include_diagram::visitor {
+using id_t = int64_t;
 
 template <typename T>
-element_visitor_context<T>::element_visitor_context(
-    clanguml::include_diagram::model::diagram &diagram, T &element)
-    : element_{element}
-    , diagram_{diagram}
-{
-}
-
-template <typename T> T &element_visitor_context<T>::element()
-{
-    return element_;
-}
+using optional_ref = std::optional<std::reference_wrapper<const T>>;
 
 template <typename T>
-clanguml::include_diagram::model::diagram &element_visitor_context<T>::diagram()
-{
-    return diagram_;
-}
-}
+using reference_vector = std::vector<std::reference_wrapper<const T>>;
+
+template <typename T>
+using reference_set = std::unordered_set<std::reference_wrapper<const T>>;
+
+} // namespace clang::common
\ No newline at end of file
diff --git a/src/include_diagram/generators/plantuml/include_diagram_generator.cc b/src/include_diagram/generators/plantuml/include_diagram_generator.cc
index b670a55b..6a0ee957 100644
--- a/src/include_diagram/generators/plantuml/include_diagram_generator.cc
+++ b/src/include_diagram/generators/plantuml/include_diagram_generator.cc
@@ -37,25 +37,25 @@ void generator::generate_relationships(
 
     namespace plantuml_common = clanguml::common::generators::plantuml;
 
-    if (f.type() == common::model::source_file_t::kDirectory) {
-        util::for_each(f, [this, &ostr](const auto &file) {
-            generate_relationships(
-                dynamic_cast<const source_file &>(*file), ostr);
-        });
-    }
-    else {
-        util::for_each_if(
-            f.relationships(),
-            [this](const auto &r) {
-                return m_model.should_include(r.type()) &&
-                    util::contains(m_generated_aliases, r.destination());
-            },
-            [&f, &ostr](const auto &r) {
-                ostr << f.alias() << " "
-                     << plantuml_common::to_plantuml(r.type(), r.style()) << " "
-                     << r.destination() << '\n';
-            });
-    }
+//    if (f.type() == common::model::source_file_t::kDirectory) {
+//        util::for_each(f, [this, &ostr](const auto &file) {
+//            generate_relationships(
+//                dynamic_cast<const source_file &>(*file), ostr);
+//        });
+//    }
+//    else {
+//        util::for_each_if(
+//            f.relationships(),
+//            [this](const auto &r) {
+//                return m_model.should_include(r.type()) &&
+//                    util::contains(m_generated_aliases, r.destination());
+//            },
+//            [&f, &ostr](const auto &r) {
+//                ostr << f.alias() << " "
+//                     << plantuml_common::to_plantuml(r.type(), r.style()) << " "
+//                     << r.destination() << '\n';
+//            });
+//    }
 }
 
 void generator::generate(const source_file &f, std::ostream &ostr) const
diff --git a/src/include_diagram/model/diagram.cc b/src/include_diagram/model/diagram.cc
index fe42713f..d217fd89 100644
--- a/src/include_diagram/model/diagram.cc
+++ b/src/include_diagram/model/diagram.cc
@@ -28,7 +28,7 @@ common::model::diagram_t diagram::type() const
     return common::model::diagram_t::kInclude;
 }
 
-type_safe::optional_ref<const common::model::diagram_element> diagram::get(
+common::optional_ref<common::model::diagram_element> diagram::get(
     const std::string &full_name) const
 {
     return get_file(full_name);
@@ -63,7 +63,7 @@ void diagram::add_file(std::unique_ptr<common::model::source_file> &&f)
     add_element(p, std::move(f));
 }
 
-type_safe::optional_ref<const common::model::source_file> diagram::get_file(
+common::optional_ref<common::model::source_file> diagram::get_file(
     const std::string &name) const
 {
     for (const auto &p : files_) {
@@ -72,7 +72,7 @@ type_safe::optional_ref<const common::model::source_file> diagram::get_file(
         }
     }
 
-    return type_safe::nullopt;
+    return {};
 }
 
 std::string diagram::to_alias(const std::string &full_name) const
@@ -94,8 +94,7 @@ std::string diagram::to_alias(const std::string &full_name) const
     return source_file.value().alias();
 }
 
-const std::vector<
-    type_safe::object_ref<const common::model::source_file, false>> &
+const common::reference_vector<const common::model::source_file> &
 diagram::files() const
 {
     return files_;
diff --git a/src/include_diagram/model/diagram.h b/src/include_diagram/model/diagram.h
index a9a1268b..b494f5b7 100644
--- a/src/include_diagram/model/diagram.h
+++ b/src/include_diagram/model/diagram.h
@@ -20,8 +20,7 @@
 #include "common/model/diagram.h"
 #include "common/model/package.h"
 #include "common/model/source_file.h"
-
-#include <type_safe/optional_ref.hpp>
+#include "common/types.h"
 
 #include <string>
 #include <vector>
@@ -42,23 +41,21 @@ public:
 
     common::model::diagram_t type() const override;
 
-    type_safe::optional_ref<const common::model::diagram_element> get(
+    common::optional_ref<common::model::diagram_element> get(
         const std::string &full_name) const override;
 
     void add_file(std::unique_ptr<common::model::source_file> &&f);
 
-    type_safe::optional_ref<const common::model::source_file> get_file(
+    common::optional_ref<common::model::source_file> get_file(
         const std::string &name) const;
 
     std::string to_alias(const std::string &full_name) const;
 
-    const std::vector<
-        type_safe::object_ref<const common::model::source_file, false>> &
+    const common::reference_vector<const common::model::source_file> &
     files() const;
 
 private:
-    std::vector<type_safe::object_ref<const common::model::source_file, false>>
-        files_;
+    common::reference_vector<const common::model::source_file> files_;
 };
 }
 
diff --git a/src/include_diagram/visitor/translation_unit_context.cc.bak b/src/include_diagram/visitor/translation_unit_context.cc.bak
deleted file mode 100644
index 4cff5297..00000000
--- a/src/include_diagram/visitor/translation_unit_context.cc.bak
+++ /dev/null
@@ -1,63 +0,0 @@
-/**
- * src/include_diagram/visitor/translation_unit_context.cc
- *
- * Copyright (c) 2021-2022 Bartek Kryza <bkryza@gmail.com>
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include "translation_unit_context.h"
-
-#include "cx/util.h"
-
-namespace clanguml::include_diagram::visitor {
-
-translation_unit_context::translation_unit_context(
-    cppast::cpp_entity_index &idx,
-    clanguml::include_diagram::model::diagram &diagram,
-    const clanguml::config::include_diagram &config)
-    : entity_index_{idx}
-    , diagram_{diagram}
-    , config_{config}
-{
-}
-
-const cppast::cpp_entity_index &translation_unit_context::entity_index() const
-{
-    return entity_index_;
-}
-
-const clanguml::config::include_diagram &
-translation_unit_context::config() const
-{
-    return config_;
-}
-
-clanguml::include_diagram::model::diagram &translation_unit_context::diagram()
-{
-    return diagram_;
-}
-
-void translation_unit_context::set_current_file(
-    type_safe::optional_ref<common::model::source_file> f)
-{
-    current_file_ = f;
-}
-
-type_safe::optional_ref<common::model::source_file>
-translation_unit_context::get_current_file() const
-{
-    return current_file_;
-}
-
-}
diff --git a/src/package_diagram/generators/plantuml/package_diagram_generator.cc b/src/package_diagram/generators/plantuml/package_diagram_generator.cc
index 38c2bd55..c73d0523 100644
--- a/src/package_diagram/generators/plantuml/package_diagram_generator.cc
+++ b/src/package_diagram/generators/plantuml/package_diagram_generator.cc
@@ -32,7 +32,7 @@ void generator::generate_relationships(
 {
     LOG_DBG("Generating relationships for package {}", p.full_name(true));
 
-    const auto &uns = m_config.using_namespace();
+//    const auto &uns = m_config.using_namespace();
 
     // Generate this packages relationship
     if (m_model.should_include(relationship_t::kDependency)) {
@@ -40,7 +40,7 @@ void generator::generate_relationships(
             std::stringstream relstr;
             try {
                 relstr << p.alias() << " ..> "
-                       << m_model.to_alias(uns.relative(r.destination()))
+                       //<< m_model.get_         to_alias(uns.relative(r.destination()))
                        << '\n';
                 ostr << relstr.str();
             }
diff --git a/src/package_diagram/model/diagram.cc b/src/package_diagram/model/diagram.cc
index 88d31318..18e8d604 100644
--- a/src/package_diagram/model/diagram.cc
+++ b/src/package_diagram/model/diagram.cc
@@ -28,7 +28,8 @@ common::model::diagram_t diagram::type() const
     return common::model::diagram_t::kPackage;
 }
 
-const std::vector<type_safe::object_ref<const common::model::package>> &
+const std::vector<
+    std::reference_wrapper<const clanguml::common::model::package>> &
 diagram::packages() const
 {
     return packages_;
@@ -45,8 +46,8 @@ void diagram::add_package(std::unique_ptr<common::model::package> &&p)
     add_element(ns, std::move(p));
 }
 
-type_safe::optional_ref<const common::model::package> diagram::get_package(
-    const std::string &name) const
+std::optional<std::reference_wrapper<const common::model::package>>
+diagram::get_package(const std::string &name) const
 {
     for (const auto &p : packages_) {
         auto p_full_name = p.get().full_name(false);
@@ -55,11 +56,12 @@ type_safe::optional_ref<const common::model::package> diagram::get_package(
         }
     }
 
-    return type_safe::nullopt;
+    return {};
 }
 
-type_safe::optional_ref<const common::model::diagram_element> diagram::get(
-    const std::string &full_name) const
+std::optional<
+    std::reference_wrapper<const clanguml::common::model::diagram_element>>
+diagram::get(const std::string &full_name) const
 {
     return get_package(full_name);
 }
diff --git a/src/package_diagram/model/diagram.h b/src/package_diagram/model/diagram.h
index 253d8f45..2d39dac5 100644
--- a/src/package_diagram/model/diagram.h
+++ b/src/package_diagram/model/diagram.h
@@ -20,8 +20,6 @@
 #include "common/model/diagram.h"
 #include "common/model/package.h"
 
-#include <type_safe/optional_ref.hpp>
-
 #include <string>
 #include <vector>
 
@@ -41,21 +39,24 @@ public:
 
     common::model::diagram_t type() const override;
 
-    const std::vector<type_safe::object_ref<const common::model::package>> &
+    const std::vector<
+        std::reference_wrapper<const clanguml::common::model::package>> &
     packages() const;
 
-    type_safe::optional_ref<const common::model::diagram_element> get(
-        const std::string &full_name) const override;
+    std::optional<
+        std::reference_wrapper<const clanguml::common::model::diagram_element>>
+    get(const std::string &full_name) const override;
 
     void add_package(std::unique_ptr<common::model::package> &&p);
 
-    type_safe::optional_ref<const common::model::package> get_package(
-        const std::string &name) const;
+    std::optional<
+        std::reference_wrapper<const clanguml::common::model::package>>
+    get_package(const std::string &name) const;
 
     std::string to_alias(const std::string &full_name) const;
 
 private:
-    std::vector<type_safe::object_ref<const common::model::package, false>>
+    std::vector<std::reference_wrapper<const clanguml::common::model::package>>
         packages_;
 };
 }
diff --git a/src/package_diagram/visitor/element_visitor_context.cc.bak b/src/package_diagram/visitor/element_visitor_context.cc.bak
deleted file mode 100644
index 56f3a2ac..00000000
--- a/src/package_diagram/visitor/element_visitor_context.cc.bak
+++ /dev/null
@@ -1,43 +0,0 @@
-/**
- * src/package_diagram/model/visitor/element_visitor_context.cc
- *
- * Copyright (c) 2021-2022 Bartek Kryza <bkryza@gmail.com>
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include "element_visitor_context.h"
-
-#include "translation_unit_context.h"
-
-namespace clanguml::package_diagram::visitor {
-
-template <typename T>
-element_visitor_context<T>::element_visitor_context(
-    clanguml::package_diagram::model::diagram &diagram, T &element)
-    : element_{element}
-    , diagram_{diagram}
-{
-}
-
-template <typename T> T &element_visitor_context<T>::element()
-{
-    return element_;
-}
-
-template <typename T>
-clanguml::package_diagram::model::diagram &element_visitor_context<T>::diagram()
-{
-    return diagram_;
-}
-}
diff --git a/src/package_diagram/visitor/translation_unit_context.cc.bak b/src/package_diagram/visitor/translation_unit_context.cc.bak
deleted file mode 100644
index 38f8c2d9..00000000
--- a/src/package_diagram/visitor/translation_unit_context.cc.bak
+++ /dev/null
@@ -1,194 +0,0 @@
-/**
- * src/package_diagram/visitor/translation_unit_context.cc
- *
- * Copyright (c) 2021-2022 Bartek Kryza <bkryza@gmail.com>
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include "translation_unit_context.h"
-
-#include "cx/util.h"
-
-namespace clanguml::package_diagram::visitor {
-
-translation_unit_context::translation_unit_context(
-    cppast::cpp_entity_index &idx,
-    clanguml::package_diagram::model::diagram &diagram,
-    const clanguml::config::package_diagram &config)
-    : entity_index_{idx}
-    , diagram_{diagram}
-    , config_{config}
-{
-}
-
-bool translation_unit_context::has_namespace_alias(
-    const std::string &full_name) const
-{
-    bool res =
-        namespace_alias_index_.find(full_name) != namespace_alias_index_.end();
-
-    LOG_DBG("Alias {} {} found in index", full_name, res ? "" : "not");
-
-    return res;
-}
-
-void translation_unit_context::add_namespace_alias(const std::string &full_name,
-    type_safe::object_ref<const cppast::cpp_namespace> ref)
-{
-    if (!has_namespace_alias(full_name)) {
-        LOG_DBG("Stored type alias: {} -> {} ", full_name, ref.get().name());
-
-        namespace_alias_index_.emplace(full_name, std::move(ref));
-    }
-}
-
-type_safe::object_ref<const cppast::cpp_namespace>
-translation_unit_context::get_namespace_alias(
-    const std::string &full_name) const
-{
-    assert(has_namespace_alias(full_name));
-
-    return namespace_alias_index_.at(full_name);
-}
-
-type_safe::object_ref<const cppast::cpp_namespace>
-translation_unit_context::get_namespace_alias_final(
-    const cppast::cpp_namespace &ns) const
-{
-    auto ns_full_name = cx::util::full_name({}, ns);
-
-    ns_full_name = cx::util::ns(ns) + "::" + ns_full_name;
-
-    if (has_namespace_alias(ns_full_name)) {
-        return get_namespace_alias_final(
-            namespace_alias_index_.at(ns_full_name).get());
-    }
-
-    return type_safe::ref(ns);
-}
-
-bool translation_unit_context::has_type_alias(
-    const std::string &full_name) const
-{
-    bool res = alias_index_.find(full_name) != alias_index_.end();
-
-    LOG_DBG("Alias {} {} found in index", full_name, res ? "" : "not");
-
-    return res;
-}
-
-void translation_unit_context::add_type_alias(const std::string &full_name,
-    type_safe::object_ref<const cppast::cpp_type> &&ref)
-{
-    if (!has_type_alias(full_name)) {
-        LOG_DBG("Stored type alias: {} -> {} ", full_name,
-            cppast::to_string(ref.get()));
-
-        alias_index_.emplace(full_name, std::move(ref));
-    }
-}
-
-type_safe::object_ref<const cppast::cpp_type>
-translation_unit_context::get_type_alias(const std::string &full_name) const
-{
-    assert(has_type_alias(full_name));
-
-    return alias_index_.at(full_name);
-}
-
-type_safe::object_ref<const cppast::cpp_type>
-translation_unit_context::get_type_alias_final(const cppast::cpp_type &t) const
-{
-    const auto type_full_name =
-        cx::util::full_name(cppast::remove_cv(t), entity_index_, false);
-
-    if (has_type_alias(type_full_name)) {
-        return get_type_alias_final(alias_index_.at(type_full_name).get());
-    }
-
-    return type_safe::ref(t);
-}
-
-bool translation_unit_context::has_type_alias_template(
-    const std::string &full_name) const
-{
-    bool res =
-        alias_template_index_.find(full_name) != alias_template_index_.end();
-
-    LOG_DBG("Alias template {} {} found in index", full_name, res ? "" : "not");
-
-    return res;
-}
-
-void translation_unit_context::add_type_alias_template(
-    const std::string &full_name,
-    type_safe::object_ref<const cppast::cpp_type> &&ref)
-{
-    if (!has_type_alias_template(full_name)) {
-        LOG_DBG("Stored type alias template for: {} ", full_name);
-
-        alias_template_index_.emplace(full_name, std::move(ref));
-    }
-}
-
-type_safe::object_ref<const cppast::cpp_type>
-translation_unit_context::get_type_alias_template(
-    const std::string &full_name) const
-{
-    assert(has_type_alias_template(full_name));
-
-    return alias_template_index_.at(full_name);
-}
-
-void translation_unit_context::push_namespace(const std::string &ns)
-{
-    namespace_ |= ns;
-}
-
-void translation_unit_context::pop_namespace() { namespace_.pop_back(); }
-
-const common::model::namespace_ &translation_unit_context::get_namespace() const
-{
-    return namespace_;
-}
-
-const cppast::cpp_entity_index &translation_unit_context::entity_index() const
-{
-    return entity_index_;
-}
-
-const clanguml::config::package_diagram &
-translation_unit_context::config() const
-{
-    return config_;
-}
-
-clanguml::package_diagram::model::diagram &translation_unit_context::diagram()
-{
-    return diagram_;
-}
-
-void translation_unit_context::set_current_package(
-    type_safe::optional_ref<common::model::package> p)
-{
-    current_package_ = p;
-}
-
-type_safe::optional_ref<common::model::package>
-translation_unit_context::get_current_package() const
-{
-    return current_package_;
-}
-
-}
diff --git a/src/sequence_diagram/model/diagram.cc b/src/sequence_diagram/model/diagram.cc
index 39140f29..26a69a7b 100644
--- a/src/sequence_diagram/model/diagram.cc
+++ b/src/sequence_diagram/model/diagram.cc
@@ -28,7 +28,7 @@ common::model::diagram_t diagram::type() const
     return common::model::diagram_t::kSequence;
 }
 
-type_safe::optional_ref<const common::model::diagram_element> diagram::get(
+common::optional_ref<common::model::diagram_element> diagram::get(
     const std::string & /*full_name*/) const
 {
     return {};
diff --git a/src/sequence_diagram/model/diagram.h b/src/sequence_diagram/model/diagram.h
index 552f0f79..96f297d7 100644
--- a/src/sequence_diagram/model/diagram.h
+++ b/src/sequence_diagram/model/diagram.h
@@ -19,6 +19,7 @@
 
 #include "activity.h"
 #include "common/model/diagram.h"
+#include "common/types.h"
 
 #include <map>
 #include <string>
@@ -36,7 +37,7 @@ public:
 
     common::model::diagram_t type() const override;
 
-    type_safe::optional_ref<const common::model::diagram_element> get(
+    common::optional_ref<common::model::diagram_element> get(
         const std::string &full_name) const override;
 
     std::string to_alias(const std::string &full_name) const;
diff --git a/src/sequence_diagram/visitor/translation_unit_context.cc.bak b/src/sequence_diagram/visitor/translation_unit_context.cc.bak
deleted file mode 100644
index b2b08ba6..00000000
--- a/src/sequence_diagram/visitor/translation_unit_context.cc.bak
+++ /dev/null
@@ -1,65 +0,0 @@
-/**
- * src/sequence_diagram/visitor/translation_unit_context.cc
- *
- * Copyright (c) 2021-2022 Bartek Kryza <bkryza@gmail.com>
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include "translation_unit_context.h"
-
-#include <clang-c/CXCompilationDatabase.h>
-#include <clang-c/Index.h>
-#include <spdlog/spdlog.h>
-
-namespace clanguml::sequence_diagram::visitor {
-
-translation_unit_context::translation_unit_context(
-    cppast::cpp_entity_index &idx,
-    clanguml::sequence_diagram::model::diagram &diagram,
-    const clanguml::config::sequence_diagram &config)
-    : entity_index_{idx}
-    , diagram_{diagram}
-    , config_{config}
-{
-}
-
-void translation_unit_context::push_namespace(const std::string &ns)
-{
-    namespace_.push_back(ns);
-}
-
-void translation_unit_context::pop_namespace() { namespace_.pop_back(); }
-
-const std::vector<std::string> &translation_unit_context::get_namespace() const
-{
-    return namespace_;
-}
-
-const clanguml::config::sequence_diagram &
-translation_unit_context::config() const
-{
-    return config_;
-}
-
-clanguml::sequence_diagram::model::diagram &translation_unit_context::diagram()
-{
-    return diagram_;
-}
-
-const cppast::cpp_entity_index &translation_unit_context::entity_index() const
-{
-    return entity_index_;
-}
-
-}