From 1da3dcb720b67233536af2f1094b1d3b47ec3ec6 Mon Sep 17 00:00:00 2001
From: Bartek Kryza <bkryza@gmail.com>
Date: Fri, 18 Feb 2022 21:56:15 +0100
Subject: [PATCH] Refactored class diagram model to keep namespace structure

---
 .../plantuml/class_diagram_generator.cc       |  14 +--
 src/class_diagram/model/class.h               |   5 +
 src/class_diagram/model/diagram.cc            |  70 ++++++++----
 src/class_diagram/model/diagram.h             |  26 +++--
 src/class_diagram/model/enum.cc               |   5 +-
 src/class_diagram/model/enum.h                |   6 ++
 .../visitor/translation_unit_context.cc       |  59 +++++++++++
 .../visitor/translation_unit_context.h        |  22 ++++
 .../visitor/translation_unit_visitor.cc       | 100 ++++++++++++++----
 .../visitor/translation_unit_visitor.h        |   3 +-
 10 files changed, 250 insertions(+), 60 deletions(-)

diff --git a/src/class_diagram/generators/plantuml/class_diagram_generator.cc b/src/class_diagram/generators/plantuml/class_diagram_generator.cc
index 3e7b906d..83d07643 100644
--- a/src/class_diagram/generators/plantuml/class_diagram_generator.cc
+++ b/src/class_diagram/generators/plantuml/class_diagram_generator.cc
@@ -277,30 +277,30 @@ void generator::generate(std::ostream &ostr) const
 
     if (m_config.should_include_entities("classes")) {
         for (const auto &c : m_model.classes()) {
-            if (!m_config.should_include(c.name()))
+            if (!m_config.should_include(c.get().name()))
                 continue;
-            generate_alias(c, ostr);
+            generate_alias(*c, ostr);
             ostr << '\n';
         }
 
         for (const auto &e : m_model.enums()) {
-            if (!m_config.should_include(e.name()))
+            if (!m_config.should_include(e.get().name()))
                 continue;
-            generate_alias(e, ostr);
+            generate_alias(*e, ostr);
             ostr << '\n';
         }
 
         for (const auto &c : m_model.classes()) {
-            if (!m_config.should_include(c.name()))
+            if (!m_config.should_include(c.get().name()))
                 continue;
-            generate(c, ostr);
+            generate(*c, ostr);
             ostr << '\n';
         }
     }
 
     if (m_config.should_include_entities("enums"))
         for (const auto &e : m_model.enums()) {
-            generate(e, ostr);
+            generate(*e, ostr);
             ostr << '\n';
         }
 
diff --git a/src/class_diagram/model/class.h b/src/class_diagram/model/class.h
index e84fba7d..0f5dbb11 100644
--- a/src/class_diagram/model/class.h
+++ b/src/class_diagram/model/class.h
@@ -36,6 +36,11 @@ class class_ : public common::model::element,
 public:
     class_(const std::vector<std::string> &using_namespaces);
 
+    class_(const class_ &) = delete;
+    class_(class_ &&) = default;
+    class_ &operator=(const class_ &) = delete;
+    class_ &operator=(class_ &&) = default;
+
     bool is_struct() const;
     void is_struct(bool is_struct);
 
diff --git a/src/class_diagram/model/diagram.cc b/src/class_diagram/model/diagram.cc
index 46ec8185..52a4e834 100644
--- a/src/class_diagram/model/diagram.cc
+++ b/src/class_diagram/model/diagram.cc
@@ -23,40 +23,66 @@
 
 namespace clanguml::class_diagram::model {
 
-const std::vector<class_> diagram::classes() const { return classes_; }
+const std::vector<type_safe::object_ref<class_>> diagram::classes() const
+{
+    return classes_;
+}
 
-const std::vector<enum_> diagram::enums() const { return enums_; }
+const std::vector<type_safe::object_ref<enum_>> diagram::enums() const
+{
+    return enums_;
+}
 
 bool diagram::has_class(const class_ &c) const
 {
     return std::any_of(classes_.cbegin(), classes_.cend(),
-        [&c](const auto &cc) { return cc.full_name() == c.full_name(); });
+        [&c](const auto &cc) { return cc.get().full_name() == c.full_name(); });
 }
 
-void diagram::add_type_alias(type_alias &&ta)
+bool diagram::has_enum(const enum_ &e) const
 {
-    LOG_DBG("Adding global alias: {} -> {}", ta.alias(), ta.underlying_type());
-
-    type_aliases_[ta.alias()] = std::move(ta);
+    return std::any_of(enums_.cbegin(), enums_.cend(),
+        [&e](const auto &ee) { return ee.get().full_name() == e.full_name(); });
 }
 
-void diagram::add_class(class_ &&c)
+void diagram::add_type_alias(std::unique_ptr<type_alias> &&ta)
 {
-    LOG_DBG("Adding class: {}, {}", c.name(), c.full_name());
-    if (!has_class(c))
-        classes_.emplace_back(std::move(c));
+    LOG_DBG(
+        "Adding global alias: {} -> {}", ta->alias(), ta->underlying_type());
+
+    type_aliases_[ta->alias()] = std::move(ta);
+}
+
+void diagram::add_package(std::unique_ptr<common::model::package> &&p)
+{
+    LOG_DBG("Adding namespace package: {}, {}", p->name(), p->full_name(true));
+
+    add_element(p->get_namespace(), std::move(p));
+}
+
+void diagram::add_class(std::unique_ptr<class_> &&c)
+{
+    LOG_DBG("Adding class: {}, {}", c->name(), c->full_name());
+
+    if (!has_class(*c)) {
+        classes_.emplace_back(*c);
+        add_element(c->get_namespace(), std::move(c));
+    }
     else
-        LOG_DBG("Class {} ({}) already in the model", c.name(), c.full_name());
+        LOG_DBG(
+            "Class {} ({}) already in the model", c->name(), c->full_name());
 }
 
-void diagram::add_enum(enum_ &&e)
+void diagram::add_enum(std::unique_ptr<enum_> &&e)
 {
-    LOG_DBG("Adding enum: {}", e.name());
-    auto it = std::find(enums_.begin(), enums_.end(), e);
-    if (it == enums_.end())
-        enums_.emplace_back(std::move(e));
+    LOG_DBG("Adding enum: {}", e->name());
+
+    if (!has_enum(*e)) {
+        enums_.emplace_back(*e);
+        add_element(e->get_namespace(), std::move(e));
+    }
     else
-        LOG_DBG("Enum {} already in the model", e.name());
+        LOG_DBG("Enum {} already in the model", e->name());
 }
 
 std::string diagram::to_alias(const std::string &full_name) const
@@ -64,14 +90,14 @@ std::string diagram::to_alias(const std::string &full_name) const
     LOG_DBG("Looking for alias for {}", full_name);
 
     for (const auto &c : classes_) {
-        if (c.full_name() == full_name) {
-            return c.alias();
+        if (c->full_name() == full_name) {
+            return c->alias();
         }
     }
 
     for (const auto &e : enums_) {
-        if (e.full_name() == full_name) {
-            return e.alias();
+        if (e->full_name() == full_name) {
+            return e->alias();
         }
     }
 
diff --git a/src/class_diagram/model/diagram.h b/src/class_diagram/model/diagram.h
index 5ceecfb9..42aa479b 100644
--- a/src/class_diagram/model/diagram.h
+++ b/src/class_diagram/model/diagram.h
@@ -19,6 +19,8 @@
 
 #include "class.h"
 #include "common/model/diagram.h"
+#include "common/model/nested_trait.h"
+#include "common/model/package.h"
 #include "enum.h"
 #include "type_alias.h"
 
@@ -27,7 +29,9 @@
 
 namespace clanguml::class_diagram::model {
 
-class diagram : public clanguml::common::model::diagram {
+class diagram : public clanguml::common::model::diagram,
+                public clanguml::common::model::nested_trait<
+                    clanguml::common::model::element> {
 public:
     diagram() = default;
 
@@ -36,23 +40,27 @@ public:
     diagram &operator=(const diagram &) = delete;
     diagram &operator=(diagram &&) = default;
 
-    const std::vector<class_> classes() const;
+    const std::vector<type_safe::object_ref<class_>> classes() const;
 
-    const std::vector<enum_> enums() const;
+    const std::vector<type_safe::object_ref<enum_>> enums() const;
 
     bool has_class(const class_ &c) const;
 
-    void add_type_alias(type_alias &&ta);
+    bool has_enum(const enum_ &e) const;
 
-    void add_class(class_ &&c);
+    void add_type_alias(std::unique_ptr<type_alias> &&ta);
 
-    void add_enum(enum_ &&e);
+    void add_class(std::unique_ptr<class_> &&c);
+
+    void add_enum(std::unique_ptr<enum_> &&e);
+
+    void add_package(std::unique_ptr<common::model::package> &&p);
 
     std::string to_alias(const std::string &full_name) const;
 
 private:
-    std::vector<class_> classes_;
-    std::vector<enum_> enums_;
-    std::map<std::string, type_alias> type_aliases_;
+    std::vector<type_safe::object_ref<class_>> classes_;
+    std::vector<type_safe::object_ref<enum_>> enums_;
+    std::map<std::string, std::unique_ptr<type_alias>> type_aliases_;
 };
 }
diff --git a/src/class_diagram/model/enum.cc b/src/class_diagram/model/enum.cc
index 4b8fb341..49b75c77 100644
--- a/src/class_diagram/model/enum.cc
+++ b/src/class_diagram/model/enum.cc
@@ -29,7 +29,10 @@ enum_::enum_(const std::vector<std::string> &using_namespaces)
 {
 }
 
-bool operator==(const enum_ &l, const enum_ &r) { return l.name() == r.name(); }
+bool operator==(const enum_ &l, const enum_ &r)
+{
+    return (l.get_namespace() == r.get_namespace()) && (l.name() == r.name());
+}
 
 std::string enum_::full_name(bool relative) const
 {
diff --git a/src/class_diagram/model/enum.h b/src/class_diagram/model/enum.h
index cd2070fa..f54da03e 100644
--- a/src/class_diagram/model/enum.h
+++ b/src/class_diagram/model/enum.h
@@ -29,6 +29,12 @@ class enum_ : public common::model::element,
 public:
     enum_(const std::vector<std::string> &using_namespaces);
 
+    enum_(const enum_ &) = delete;
+    enum_(enum_ &&) = default;
+    enum_ &operator=(const enum_ &) = delete;
+    enum_ &operator=(enum_ &&) = default;
+
+    // TODO: Do we need this?
     friend bool operator==(const enum_ &l, const enum_ &r);
 
     std::string full_name(bool relative = true) const override;
diff --git a/src/class_diagram/visitor/translation_unit_context.cc b/src/class_diagram/visitor/translation_unit_context.cc
index 77e9c50e..7d9f612d 100644
--- a/src/class_diagram/visitor/translation_unit_context.cc
+++ b/src/class_diagram/visitor/translation_unit_context.cc
@@ -32,6 +32,53 @@ translation_unit_context::translation_unit_context(
 {
 }
 
+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 namespace 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
 {
@@ -132,4 +179,16 @@ clanguml::class_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/class_diagram/visitor/translation_unit_context.h b/src/class_diagram/visitor/translation_unit_context.h
index a5dff73c..453d2805 100644
--- a/src/class_diagram/visitor/translation_unit_context.h
+++ b/src/class_diagram/visitor/translation_unit_context.h
@@ -20,6 +20,7 @@
 #include "config/config.h"
 
 #include <cppast/cpp_entity_index.hpp>
+#include <cppast/cpp_namespace.hpp>
 #include <cppast/cpp_type.hpp>
 #include <type_safe/reference.hpp>
 
@@ -31,6 +32,17 @@ public:
         clanguml::class_diagram::model::diagram &diagram,
         const clanguml::config::class_diagram &config);
 
+    bool has_namespace_alias(const std::string &full_name) const;
+
+    void add_namespace_alias(const std::string &full_name,
+        type_safe::object_ref<const cppast::cpp_namespace> ref);
+
+    type_safe::object_ref<const cppast::cpp_namespace> get_namespace_alias(
+        const std::string &full_name) const;
+
+    type_safe::object_ref<const cppast::cpp_namespace>
+    get_namespace_alias_final(const cppast::cpp_namespace &t) const;
+
     bool has_type_alias(const std::string &full_name) const;
 
     void add_type_alias(const std::string &full_name,
@@ -62,6 +74,10 @@ public:
 
     clanguml::class_diagram::model::diagram &diagram();
 
+    void set_current_package(type_safe::optional_ref<common::model::package> p);
+
+    type_safe::optional_ref<common::model::package> get_current_package() const;
+
 private:
     // Current visitor namespace
     std::vector<std::string> namespace_;
@@ -75,6 +91,10 @@ private:
     // Reference to class diagram config
     const clanguml::config::class_diagram &config_;
 
+    // Map of discovered aliases (declared with 'namespace' keyword)
+    std::map<std::string, type_safe::object_ref<const cppast::cpp_namespace>>
+        namespace_alias_index_;
+
     // Map of discovered aliases (declared with 'using' keyword)
     std::map<std::string, type_safe::object_ref<const cppast::cpp_type>>
         alias_index_;
@@ -82,6 +102,8 @@ private:
     // Map of discovered template aliases (declared with 'using' keyword)
     std::map<std::string, type_safe::object_ref<const cppast::cpp_type>>
         alias_template_index_;
+
+    type_safe::optional_ref<common::model::package> current_package_;
 };
 
 }
diff --git a/src/class_diagram/visitor/translation_unit_visitor.cc b/src/class_diagram/visitor/translation_unit_visitor.cc
index 2ce17a31..ac4abf65 100644
--- a/src/class_diagram/visitor/translation_unit_visitor.cc
+++ b/src/class_diagram/visitor/translation_unit_visitor.cc
@@ -94,8 +94,51 @@ void translation_unit_visitor::operator()(const cppast::cpp_entity &file)
                     const auto &ns_declaration =
                         static_cast<const cppast::cpp_namespace &>(e);
                     if (!ns_declaration.is_anonymous() &&
-                        !ns_declaration.is_inline())
+                        !ns_declaration.is_inline()) {
+
+                        std::vector<std::string> package_parent =
+                            ctx.get_namespace();
+                        auto package_path = package_parent;
+                        package_path.push_back(e.name());
+
+                        auto usn = util::split(
+                            ctx.config().using_namespace()[0], "::");
+
+                        if (!util::starts_with(usn, package_path)) {
+                            auto p = std::make_unique<common::model::package>(
+                                ctx.config().using_namespace());
+                            util::remove_prefix(package_path, usn);
+                            util::remove_prefix(package_parent, usn);
+
+                            p->set_name(e.name());
+                            p->set_namespace(package_parent);
+
+                            if (ns_declaration.comment().has_value())
+                                p->add_decorators(decorators::parse(
+                                    ns_declaration.comment().value()));
+
+                            p->set_style(p->style_spec());
+
+                            for (const auto &attr :
+                                ns_declaration.attributes()) {
+                                if (attr.kind() ==
+                                    cppast::cpp_attribute_kind::deprecated) {
+                                    p->set_deprecated(true);
+                                    break;
+                                }
+                            }
+
+                            if (!p->skip()) {
+                                ctx.diagram().add_element(
+                                    package_parent, std::move(p));
+                                ctx.set_current_package(
+                                    ctx.diagram()
+                                        .get_element<common::model::package>(
+                                            package_path));
+                            }
+                        }
                         ctx.push_namespace(e.name());
+                    }
                 }
                 else {
                     LOG_DBG("========== Leaving '{}' - {}", e.name(),
@@ -108,6 +151,15 @@ void translation_unit_visitor::operator()(const cppast::cpp_entity &file)
                         ctx.pop_namespace();
                 }
             }
+            else if (e.kind() == cppast::cpp_entity_kind::namespace_alias_t) {
+                auto &na = static_cast<const cppast::cpp_namespace_alias &>(e);
+
+                for (const auto &alias_target :
+                    na.target().get(ctx.entity_index())) {
+                    auto full_ns = cx::util::full_name(ctx.get_namespace(), na);
+                    ctx.add_namespace_alias(full_ns, alias_target);
+                }
+            }
             else if (e.kind() ==
                 cppast::cpp_entity_kind::class_template_specialization_t) {
                 LOG_DBG("========== Visiting '{}' - {}",
@@ -158,9 +210,9 @@ void translation_unit_visitor::operator()(const cppast::cpp_entity &file)
                     cppast::to_string(e.kind()));
 
                 auto &ta = static_cast<const cppast::cpp_type_alias &>(e);
-                type_alias t;
-                t.set_alias(cx::util::full_name(ctx.get_namespace(), ta));
-                t.set_underlying_type(cx::util::full_name(ta.underlying_type(),
+                std::unique_ptr<type_alias> t;
+                t->set_alias(cx::util::full_name(ctx.get_namespace(), ta));
+                t->set_underlying_type(cx::util::full_name(ta.underlying_type(),
                     ctx.entity_index(), cx::util::is_inside_class(e)));
 
                 ctx.add_type_alias(cx::util::full_name(ctx.get_namespace(), ta),
@@ -183,7 +235,7 @@ void translation_unit_visitor::operator()(const cppast::cpp_entity &file)
                             .name());
                 }
                 else {
-                    class_ tinst = build_template_instantiation(static_cast<
+                    auto tinst = build_template_instantiation(static_cast<
                         const cppast::cpp_template_instantiation_type &>(
                         at.type_alias().underlying_type()));
 
@@ -202,8 +254,10 @@ void translation_unit_visitor::process_enum_declaration(
         return;
     }
 
-    enum_ e{ctx.config().using_namespace()};
+    auto e_ptr = std::make_unique<enum_>(ctx.config().using_namespace());
+    auto &e = *e_ptr;
     e.set_name(cx::util::full_name(ctx.get_namespace(), enm));
+    e.set_namespace(ctx.get_namespace());
 
     if (enm.comment().has_value())
         e.add_decorators(decorators::parse(enm.comment().value()));
@@ -235,16 +289,18 @@ void translation_unit_visitor::process_enum_declaration(
         }
     }
 
-    ctx.diagram().add_enum(std::move(e));
+    ctx.diagram().add_enum(std::move(e_ptr));
 }
 
 void translation_unit_visitor::process_class_declaration(
     const cppast::cpp_class &cls,
     type_safe::optional_ref<const cppast::cpp_template_specialization> tspec)
 {
-    class_ c{ctx.config().using_namespace()};
+    auto c_ptr = std::make_unique<class_>(ctx.config().using_namespace());
+    auto &c = *c_ptr;
     c.is_struct(cls.class_kind() == cppast::cpp_class_kind::struct_t);
     c.set_name(cx::util::full_name(ctx.get_namespace(), cls));
+    c.set_namespace(ctx.get_namespace());
 
     if (cls.comment().has_value())
         c.add_decorators(decorators::parse(cls.comment().value()));
@@ -520,7 +576,7 @@ void translation_unit_visitor::process_class_declaration(
         static_cast<const char *>(cls.user_data()),
         fmt::ptr(reinterpret_cast<const void *>(&cls)));
 
-    ctx.diagram().add_class(std::move(c));
+    ctx.diagram().add_class(std::move(c_ptr));
 }
 
 bool translation_unit_visitor::process_field_with_template_instantiation(
@@ -539,7 +595,8 @@ bool translation_unit_visitor::process_field_with_template_instantiation(
         static_cast<const cppast::cpp_template_instantiation_type &>(
             resolve_alias(template_instantiation_type));
 
-    class_ tinst = build_template_instantiation(unaliased);
+    auto tinst_ptr = build_template_instantiation(unaliased);
+    auto &tinst = *tinst_ptr;
 
     // Infer the relationship of this field to the template
     // instantiation
@@ -577,7 +634,7 @@ bool translation_unit_visitor::process_field_with_template_instantiation(
 
         LOG_DBG("Created template instantiation: {}", tinst.full_name());
 
-        ctx.diagram().add_class(std::move(tinst));
+        ctx.diagram().add_class(std::move(tinst_ptr));
     }
 
     return res;
@@ -953,9 +1010,9 @@ void translation_unit_visitor::process_function_parameter(
                     }
                     else {
                         // First check if tinst already exists
-                        class_ tinst = build_template_instantiation(
+                        auto tinst_ptr = build_template_instantiation(
                             template_instantiation_type);
-
+                        auto &tinst = *tinst_ptr;
                         relationship rr{
                             relationship_t::kDependency, tinst.full_name()};
 
@@ -967,7 +1024,7 @@ void translation_unit_visitor::process_function_parameter(
 
                         c.add_relationship(std::move(rr));
 
-                        ctx.diagram().add_class(std::move(tinst));
+                        ctx.diagram().add_class(std::move(tinst_ptr));
                     }
                 }
             }
@@ -1201,15 +1258,18 @@ bool translation_unit_visitor::find_relationships(const cppast::cpp_type &t_,
     return found;
 }
 
-class_ translation_unit_visitor::build_template_instantiation(
+std::unique_ptr<class_> translation_unit_visitor::build_template_instantiation(
     const cppast::cpp_template_instantiation_type &t,
     std::optional<clanguml::class_diagram::model::class_ *> parent)
 {
-    class_ tinst{ctx.config().using_namespace()};
+    auto tinst_ptr = std::make_unique<class_>(ctx.config().using_namespace());
+    auto &tinst = *tinst_ptr;
     std::string full_template_name;
 
     std::deque<std::tuple<std::string, int, bool>> template_base_params{};
 
+    tinst.set_namespace(ctx.get_namespace());
+
     if (t.primary_template().get(ctx.entity_index()).size()) {
         const auto &primary_template_ref =
             static_cast<const cppast::cpp_class_template &>(
@@ -1349,16 +1409,16 @@ class_ translation_unit_visitor::build_template_instantiation(
                     if (parent)
                         nnn = (*parent)->name();
 
-                    class_ nested_tinst =
+                    auto nested_tinst =
                         build_template_instantiation(nested_template_parameter,
                             ctx.config().should_include(tinst.full_name(false))
                                 ? std::make_optional(&tinst)
                                 : parent);
 
                     relationship tinst_dependency{
-                        relationship_t::kDependency, nested_tinst.full_name()};
+                        relationship_t::kDependency, nested_tinst->full_name()};
 
-                    auto nested_tinst_full_name = nested_tinst.full_name();
+                    auto nested_tinst_full_name = nested_tinst->full_name();
 
                     if (ctx.config().should_include(fn)) {
                         ctx.diagram().add_class(std::move(nested_tinst));
@@ -1480,7 +1540,7 @@ class_ translation_unit_visitor::build_template_instantiation(
     relationship r{relationship_t::kInstantiation, destination};
     tinst.add_relationship(std::move(r));
 
-    return tinst;
+    return tinst_ptr;
 }
 
 const cppast::cpp_type &translation_unit_visitor::resolve_alias(
diff --git a/src/class_diagram/visitor/translation_unit_visitor.h b/src/class_diagram/visitor/translation_unit_visitor.h
index 508866f8..5f158325 100644
--- a/src/class_diagram/visitor/translation_unit_visitor.h
+++ b/src/class_diagram/visitor/translation_unit_visitor.h
@@ -120,7 +120,8 @@ public:
         clanguml::class_diagram::model::class_ &parent);
 
 private:
-    clanguml::class_diagram::model::class_ build_template_instantiation(
+    std::unique_ptr<clanguml::class_diagram::model::class_>
+    build_template_instantiation(
         const cppast::cpp_template_instantiation_type &t,
         std::optional<clanguml::class_diagram::model::class_ *> parent = {});