diff --git a/src/common/model/enums.cc b/src/common/model/enums.cc
index 3cbf85d2..ea40af08 100644
--- a/src/common/model/enums.cc
+++ b/src/common/model/enums.cc
@@ -74,6 +74,26 @@ std::string to_string(message_t r)
         return "call";
     case message_t::kReturn:
         return "return";
+    case message_t::kIf:
+        return "if";
+    case message_t::kElse:
+        return "else";
+    case message_t::kElseIf:
+        return "else if";
+    case message_t::kIfEnd:
+        return "end if";
+    case message_t::kWhile:
+        return "while";
+    case message_t::kWhileEnd:
+        return "end while";
+    case message_t::kDo:
+        return "do";
+    case message_t::kDoEnd:
+        return "end do";
+    case message_t::kFor:
+        return "for";
+    case message_t::kForEnd:
+        return "end for";
     default:
         assert(false);
         return "";
diff --git a/src/common/model/enums.h b/src/common/model/enums.h
index 78912e37..fbec4a55 100644
--- a/src/common/model/enums.h
+++ b/src/common/model/enums.h
@@ -51,14 +51,15 @@ enum class message_t {
     kDo,
     kDoEnd,
     kFor,
-    kForEnd
+    kForEnd,
+    kNone
 };
 
 std::string to_string(relationship_t r);
 
 std::string to_string(access_t r);
 
-std::string to_string(message_t r);
+std::string to_string(message_t m);
 
 std::string to_string(diagram_t r);
 
diff --git a/src/common/visitor/translation_unit_visitor.h b/src/common/visitor/translation_unit_visitor.h
index eff0e5e9..31a35a63 100644
--- a/src/common/visitor/translation_unit_visitor.h
+++ b/src/common/visitor/translation_unit_visitor.h
@@ -42,8 +42,6 @@ public:
 
     clang::SourceManager &source_manager() const;
 
-    void finalize();
-
 protected:
     void set_source_location(const clang::Decl &decl,
         clanguml::common::model::source_location &element);
diff --git a/src/sequence_diagram/generators/plantuml/sequence_diagram_generator.cc b/src/sequence_diagram/generators/plantuml/sequence_diagram_generator.cc
index 777ff93b..2c6668ef 100644
--- a/src/sequence_diagram/generators/plantuml/sequence_diagram_generator.cc
+++ b/src/sequence_diagram/generators/plantuml/sequence_diagram_generator.cc
@@ -45,16 +45,16 @@ std::string generator::render_name(std::string name) const
 
 void generator::generate_call(const message &m, std::ostream &ostr) const
 {
-    const auto &from = m_model.get_participant<model::participant>(m.from);
-    const auto &to = m_model.get_participant<model::participant>(m.to);
+    const auto &from = m_model.get_participant<model::participant>(m.from());
+    const auto &to = m_model.get_participant<model::participant>(m.to());
 
     if (!from || !to) {
-        LOG_DBG("Skipping empty call from '{}' to '{}'", m.from, m.to);
+        LOG_DBG("Skipping empty call from '{}' to '{}'", m.from(), m.to());
         return;
     }
 
-    generate_participant(ostr, m.from);
-    generate_participant(ostr, m.to);
+    generate_participant(ostr, m.from());
+    generate_participant(ostr, m.to());
 
     std::string message;
 
@@ -89,16 +89,16 @@ void generator::generate_call(const message &m, std::ostream &ostr) const
     ostr << " : " << message << std::endl;
 
     LOG_DBG("Generated call '{}' from {} [{}] to {} [{}]", message, from,
-        m.from, to, m.to);
+        m.from(), to, m.to());
 }
 
 void generator::generate_return(const message &m, std::ostream &ostr) const
 {
     // Add return activity only for messages between different actors and
     // only if the return type is different than void
-    const auto &from = m_model.get_participant<model::participant>(m.from);
-    const auto &to = m_model.get_participant<model::function>(m.to);
-    if ((m.from != m.to) && !to.value().is_void()) {
+    const auto &from = m_model.get_participant<model::participant>(m.from());
+    const auto &to = m_model.get_participant<model::function>(m.to());
+    if ((m.from() != m.to()) && !to.value().is_void()) {
         const std::string from_alias = generate_alias(from.value());
 
         const std::string to_alias = generate_alias(to.value());
@@ -112,15 +112,16 @@ void generator::generate_return(const message &m, std::ostream &ostr) const
 void generator::generate_activity(const activity &a, std::ostream &ostr,
     std::vector<common::model::diagram_element::id_t> &visited) const
 {
-    for (const auto &m : a.messages) {
-        if (m.type == message_t::kCall) {
-            const auto &to = m_model.get_participant<model::participant>(m.to);
+    for (const auto &m : a.messages()) {
+        if (m.type() == message_t::kCall) {
+            const auto &to =
+                m_model.get_participant<model::participant>(m.to());
             if (!to)
                 continue;
 
-            visited.push_back(m.from);
+            visited.push_back(m.from());
 
-            LOG_DBG("Generating message {} --> {}", m.from, m.to);
+            LOG_DBG("Generating message {} --> {}", m.from(), m.to());
 
             generate_call(m, ostr);
 
@@ -128,18 +129,19 @@ void generator::generate_activity(const activity &a, std::ostream &ostr,
 
             ostr << "activate " << to_alias << std::endl;
 
-            if (m_model.sequences.find(m.to) != m_model.sequences.end()) {
-                if (std::find(visited.begin(), visited.end(), m.to) ==
+            if (m_model.sequences().find(m.to()) != m_model.sequences().end()) {
+                if (std::find(visited.begin(), visited.end(), m.to()) ==
                     visited
                         .end()) { // break infinite recursion on recursive calls
                     LOG_DBG("Creating activity {} --> {} - missing sequence {}",
-                        m.from, m.to, m.to);
-                    generate_activity(m_model.sequences[m.to], ostr, visited);
+                        m.from(), m.to(), m.to());
+                    generate_activity(
+                        m_model.get_activity(m.to()), ostr, visited);
                 }
             }
             else
                 LOG_DBG("Skipping activity {} --> {} - missing sequence {}",
-                    m.from, m.to, m.to);
+                    m.from(), m.to(), m.to());
 
             generate_return(m, ostr);
 
@@ -147,34 +149,34 @@ void generator::generate_activity(const activity &a, std::ostream &ostr,
 
             visited.pop_back();
         }
-        else if (m.type == message_t::kIf) {
+        else if (m.type() == message_t::kIf) {
             ostr << "alt\n";
         }
-        else if (m.type == message_t::kElseIf) {
+        else if (m.type() == message_t::kElseIf) {
             ostr << "else\n";
         }
-        else if (m.type == message_t::kElse) {
+        else if (m.type() == message_t::kElse) {
             ostr << "else\n";
         }
-        else if (m.type == message_t::kIfEnd) {
+        else if (m.type() == message_t::kIfEnd) {
             ostr << "end\n";
         }
-        else if (m.type == message_t::kWhile) {
+        else if (m.type() == message_t::kWhile) {
             ostr << "loop\n";
         }
-        else if (m.type == message_t::kWhileEnd) {
+        else if (m.type() == message_t::kWhileEnd) {
             ostr << "end\n";
         }
-        else if (m.type == message_t::kFor) {
+        else if (m.type() == message_t::kFor) {
             ostr << "loop\n";
         }
-        else if (m.type == message_t::kForEnd) {
+        else if (m.type() == message_t::kForEnd) {
             ostr << "end\n";
         }
-        else if (m.type == message_t::kDo) {
+        else if (m.type() == message_t::kDo) {
             ostr << "loop\n";
         }
-        else if (m.type == message_t::kDoEnd) {
+        else if (m.type() == message_t::kDoEnd) {
             ostr << "end\n";
         }
     }
@@ -182,7 +184,7 @@ void generator::generate_activity(const activity &a, std::ostream &ostr,
 
 void generator::generate_participant(std::ostream &ostr, common::id_t id) const
 {
-    for (const auto participant_id : m_model.active_participants_) {
+    for (const auto participant_id : m_model.active_participants()) {
         if (participant_id != id)
             continue;
 
@@ -287,8 +289,8 @@ void generator::generate(std::ostream &ostr) const
     for (const auto &sf : m_config.start_from()) {
         if (sf.location_type == source_location::location_t::function) {
             common::model::diagram_element::id_t start_from;
-            for (const auto &[k, v] : m_model.sequences) {
-                const auto &caller = *m_model.participants.at(v.from);
+            for (const auto &[k, v] : m_model.sequences()) {
+                const auto &caller = *m_model.participants().at(v.from());
                 std::string vfrom = caller.full_name(false);
                 if (vfrom == sf.location) {
                     LOG_DBG("Found sequence diagram start point: {}", k);
@@ -327,7 +329,7 @@ void generator::generate(std::ostream &ostr) const
             ostr << "activate " << from_alias << std::endl;
 
             generate_activity(
-                m_model.sequences[start_from], ostr, visited_participants);
+                m_model.get_activity(start_from), ostr, visited_participants);
 
             if (from.value().type_name() == "method" ||
                 m_config.combine_free_functions_into_file_participants()) {
diff --git a/src/sequence_diagram/model/activity.cc b/src/sequence_diagram/model/activity.cc
index c9e3d226..e2c921ca 100644
--- a/src/sequence_diagram/model/activity.cc
+++ b/src/sequence_diagram/model/activity.cc
@@ -20,4 +20,19 @@
 
 namespace clanguml::sequence_diagram::model {
 
+activity::activity(common::model::diagram_element::id_t id)
+    : from_{id}
+{
+}
+
+void activity::add_message(message m) { messages_.emplace_back(std::move(m)); }
+
+std::vector<message> &activity::messages() { return messages_; }
+
+const std::vector<message> &activity::messages() const { return messages_; }
+
+void activity::set_from(common::model::diagram_element::id_t f) { from_ = f; }
+
+common::model::diagram_element::id_t activity::from() const { return from_; }
+
 }
diff --git a/src/sequence_diagram/model/activity.h b/src/sequence_diagram/model/activity.h
index 41505d9f..2622ad4c 100644
--- a/src/sequence_diagram/model/activity.h
+++ b/src/sequence_diagram/model/activity.h
@@ -25,9 +25,23 @@
 
 namespace clanguml::sequence_diagram::model {
 
-struct activity {
-    common::model::diagram_element::id_t from;
-    std::vector<message> messages;
+class activity {
+public:
+    activity(common::model::diagram_element::id_t id);
+
+    void add_message(message m);
+
+    std::vector<message> &messages();
+
+    const std::vector<message> &messages() const;
+
+    void set_from(common::model::diagram_element::id_t f);
+
+    common::model::diagram_element::id_t from() const;
+
+private:
+    common::model::diagram_element::id_t from_;
+    std::vector<message> messages_;
 };
 
 }
diff --git a/src/sequence_diagram/model/diagram.cc b/src/sequence_diagram/model/diagram.cc
index 61c9af86..21feb99c 100644
--- a/src/sequence_diagram/model/diagram.cc
+++ b/src/sequence_diagram/model/diagram.cc
@@ -31,7 +31,7 @@ common::model::diagram_t diagram::type() const
 common::optional_ref<common::model::diagram_element> diagram::get(
     const std::string &full_name) const
 {
-    for (const auto &[id, participant] : participants) {
+    for (const auto &[id, participant] : participants_) {
         if (participant->full_name(false) == full_name)
             return {*participant};
     }
@@ -42,8 +42,8 @@ common::optional_ref<common::model::diagram_element> diagram::get(
 common::optional_ref<common::model::diagram_element> diagram::get(
     const common::model::diagram_element::id_t id) const
 {
-    if (participants.find(id) != participants.end())
-        return {*participants.at(id)};
+    if (participants_.find(id) != participants_.end())
+        return {*participants_.at(id)};
 
     return {};
 }
@@ -62,7 +62,7 @@ inja::json diagram::context() const
     inja::json::array_t elements{};
 
     // Add classes
-    for (const auto &[id, p] : participants) {
+    for (const auto &[id, p] : participants_) {
         elements.emplace_back(p->context());
     }
 
@@ -71,46 +71,29 @@ inja::json diagram::context() const
     return ctx;
 }
 
-void diagram::print() const
+void diagram::add_participant(std::unique_ptr<participant> p)
 {
-    LOG_DBG(" --- Participants ---");
-    for (const auto &[id, participant] : participants) {
-        LOG_DBG("{} - {}", id, participant->to_string());
+    const auto participant_id = p->id();
+
+    if (participants_.find(participant_id) == participants_.end()) {
+        LOG_DBG("Adding '{}' participant: {}, {} [{}]", p->type_name(),
+            p->full_name(false), p->id(),
+            p->type_name() == "method"
+                ? dynamic_cast<method *>(p.get())->method_name()
+                : "");
+
+        participants_.emplace(participant_id, std::move(p));
     }
+}
 
-    LOG_DBG(" --- Activities ---");
-    for (const auto &[from_id, act] : sequences) {
+void diagram::add_active_participant(common::model::diagram_element::id_t id)
+{
+    active_participants_.emplace(id);
+}
 
-        LOG_DBG("Sequence id={}:", from_id);
-
-        const auto &from_activity = *(participants.at(from_id));
-
-        LOG_DBG("   Activity id={}, from={}:", act.from,
-            from_activity.full_name(false));
-
-        for (const auto &message : act.messages) {
-            if (participants.find(message.from) == participants.end())
-                continue;
-
-            const auto &from_participant = *participants.at(message.from);
-
-            if (participants.find(message.to) == participants.end()) {
-                LOG_DBG("       Message from={}, from_id={}, "
-                        "to={}, to_id={}, name={}",
-                    from_participant.full_name(false), from_participant.id(),
-                    "__UNRESOLVABLE_ID__", message.to, message.message_name);
-            }
-            else {
-                const auto &to_participant = *participants.at(message.to);
-
-                LOG_DBG("       Message from={}, from_id={}, "
-                        "to={}, to_id={}, name={}",
-                    from_participant.full_name(false), from_participant.id(),
-                    to_participant.full_name(false), to_participant.id(),
-                    message.message_name);
-            }
-        }
-    }
+activity &diagram::get_activity(common::model::diagram_element::id_t id)
+{
+    return sequences_.at(id);
 }
 
 void diagram::add_for_stmt(
@@ -158,17 +141,12 @@ void diagram::add_loop_stmt(
     if (current_caller_id == 0)
         return;
 
-    if (sequences.find(current_caller_id) == sequences.end()) {
-        activity a;
-        a.from = current_caller_id;
-        sequences.insert({current_caller_id, std::move(a)});
+    if (sequences_.find(current_caller_id) == sequences_.end()) {
+        activity a{current_caller_id};
+        sequences_.insert({current_caller_id, std::move(a)});
     }
 
-    message m;
-    m.from = current_caller_id;
-    m.type = type;
-
-    sequences[current_caller_id].messages.emplace_back(std::move(m));
+    get_activity(current_caller_id).add_message({type, current_caller_id});
 }
 
 void diagram::end_loop_stmt(
@@ -180,9 +158,7 @@ void diagram::end_loop_stmt(
     if (current_caller_id == 0)
         return;
 
-    message m;
-    m.from = current_caller_id;
-    m.type = type;
+    message m{type, current_caller_id};
 
     message_t loop_type = message_t::kWhile;
 
@@ -191,10 +167,10 @@ void diagram::end_loop_stmt(
     else if (type == message_t::kDoEnd)
         loop_type = message_t::kDo;
 
-    if (sequences.find(current_caller_id) != sequences.end()) {
-        auto &current_messages = sequences[current_caller_id].messages;
+    if (sequences_.find(current_caller_id) != sequences_.end()) {
+        auto &current_messages = get_activity(current_caller_id).messages();
 
-        if (current_messages.back().type == loop_type) {
+        if (current_messages.back().type() == loop_type) {
             current_messages.pop_back();
         }
         else {
@@ -209,16 +185,12 @@ void diagram::add_if_stmt(
 {
     using clanguml::common::model::message_t;
 
-    if (sequences.find(current_caller_id) == sequences.end()) {
-        activity a;
-        a.from = current_caller_id;
-        sequences.insert({current_caller_id, std::move(a)});
+    if (sequences_.find(current_caller_id) == sequences_.end()) {
+        activity a{current_caller_id};
+        sequences_.insert({current_caller_id, std::move(a)});
     }
-    message m;
-    m.from = current_caller_id;
-    m.type = type;
 
-    sequences[current_caller_id].messages.emplace_back(std::move(m));
+    get_activity(current_caller_id).add_message({type, current_caller_id});
 }
 
 void diagram::end_if_stmt(
@@ -227,22 +199,20 @@ void diagram::end_if_stmt(
 {
     using clanguml::common::model::message_t;
 
-    message m;
-    m.from = current_caller_id;
-    m.type = message_t::kIfEnd;
+    message m{message_t::kIfEnd, current_caller_id};
 
-    if (sequences.find(current_caller_id) != sequences.end()) {
+    if (sequences_.find(current_caller_id) != sequences_.end()) {
 
-        auto &current_messages = sequences[current_caller_id].messages;
+        auto &current_messages = get_activity(current_caller_id).messages();
         // Remove the if/else messages if there were no calls
         // added to the diagram between them
         auto last_if_it =
             std::find_if(current_messages.rbegin(), current_messages.rend(),
-                [](const message &m) { return m.type == message_t::kIf; });
+                [](const message &m) { return m.type() == message_t::kIf; });
 
         bool last_if_block_is_empty =
             std::none_of(current_messages.rbegin(), last_if_it,
-                [](const message &m) { return m.type == message_t::kCall; });
+                [](const message &m) { return m.type() == message_t::kCall; });
 
         if (!last_if_block_is_empty) {
             current_messages.emplace_back(std::move(m));
@@ -254,6 +224,88 @@ void diagram::end_if_stmt(
     }
 }
 
+bool diagram::started() const { return started_; }
+
+void diagram::started(bool s) { started_ = s; }
+
+std::map<common::model::diagram_element::id_t, activity> &diagram::sequences()
+{
+    return sequences_;
+}
+
+const std::map<common::model::diagram_element::id_t, activity> &
+diagram::sequences() const
+{
+    return sequences_;
+}
+
+std::map<common::model::diagram_element::id_t, std::unique_ptr<participant>> &
+diagram::participants()
+{
+    return participants_;
+}
+
+const std::map<common::model::diagram_element::id_t,
+    std::unique_ptr<participant>> &
+diagram::participants() const
+{
+    return participants_;
+}
+
+std::set<common::model::diagram_element::id_t> &diagram::active_participants()
+{
+    return active_participants_;
+};
+
+const std::set<common::model::diagram_element::id_t> &
+diagram::active_participants() const
+{
+    return active_participants_;
+};
+
+void diagram::print() const
+{
+    LOG_DBG(" --- Participants ---");
+    for (const auto &[id, participant] : participants_) {
+        LOG_DBG("{} - {}", id, participant->to_string());
+    }
+
+    LOG_DBG(" --- Activities ---");
+    for (const auto &[from_id, act] : sequences_) {
+
+        LOG_DBG("Sequence id={}:", from_id);
+
+        const auto &from_activity = *(participants_.at(from_id));
+
+        LOG_DBG("   Activity id={}, from={}:", act.from(),
+            from_activity.full_name(false));
+
+        for (const auto &message : act.messages()) {
+            if (participants_.find(message.from()) == participants_.end())
+                continue;
+
+            const auto &from_participant = *participants_.at(message.from());
+
+            if (participants_.find(message.to()) == participants_.end()) {
+                LOG_DBG("       Message from={}, from_id={}, "
+                        "to={}, to_id={}, name={}, type={}",
+                    from_participant.full_name(false), from_participant.id(),
+                    "__UNRESOLVABLE_ID__", message.to(), message.message_name(),
+                    to_string(message.type()));
+            }
+            else {
+                const auto &to_participant = *participants_.at(message.to());
+
+                LOG_DBG("       Message from={}, from_id={}, "
+                        "to={}, to_id={}, name={}, type={}",
+                    from_participant.full_name(false), from_participant.id(),
+                    to_participant.full_name(false), to_participant.id(),
+                    message.message_name(), to_string(message.type()));
+            }
+        }
+    }
+}
+
 }
 
 namespace clanguml::common::model {
diff --git a/src/sequence_diagram/model/diagram.h b/src/sequence_diagram/model/diagram.h
index 48f7f792..51d51976 100644
--- a/src/sequence_diagram/model/diagram.h
+++ b/src/sequence_diagram/model/diagram.h
@@ -50,81 +50,84 @@ public:
 
     void print() const;
 
-    bool started{false};
-
     template <typename T>
     common::optional_ref<T> get_participant(
         common::model::diagram_element::id_t id)
     {
-        if (participants.find(id) == participants.end()) {
+        if (participants_.find(id) == participants_.end()) {
             return {};
         }
 
         return common::optional_ref<T>(
-            static_cast<T *>(participants.at(id).get()));
+            static_cast<T *>(participants_.at(id).get()));
     }
 
     template <typename T>
     const common::optional_ref<T> get_participant(
         common::model::diagram_element::id_t id) const
     {
-        if (participants.find(id) == participants.end()) {
+        if (participants_.find(id) == participants_.end()) {
             return {};
         }
 
         return common::optional_ref<T>(
-            static_cast<T *>(participants.at(id).get()));
+            static_cast<T *>(participants_.at(id).get()));
     }
 
-    void add_participant(std::unique_ptr<participant> p)
-    {
-        const auto participant_id = p->id();
+    void add_participant(std::unique_ptr<participant> p);
 
-        if (participants.find(participant_id) == participants.end()) {
-            LOG_DBG("Adding '{}' participant: {}, {} [{}]", p->type_name(),
-                p->full_name(false), p->id(),
-                p->type_name() == "method"
-                    ? dynamic_cast<method *>(p.get())->method_name()
-                    : "");
+    void add_active_participant(common::model::diagram_element::id_t id);
 
-            participants.emplace(participant_id, std::move(p));
-        }
-    }
+    activity& get_activity(common::model::diagram_element::id_t id);
 
-    void add_active_participant(common::model::diagram_element::id_t id)
-    {
-        active_participants_.emplace(id);
-    }
+    void add_if_stmt(common::model::diagram_element::id_t current_caller_id,
+        common::model::message_t type);
+    void end_if_stmt(common::model::diagram_element::id_t current_caller_id,
+        common::model::message_t type);
 
-    std::map<common::model::diagram_element::id_t, activity> sequences;
+    void add_loop_stmt(common::model::diagram_element::id_t current_caller_id,
+        common::model::message_t type);
+    void end_loop_stmt(common::model::diagram_element::id_t current_caller_id,
+        common::model::message_t type);
+
+    void add_while_stmt(common::model::diagram_element::id_t i);
+    void end_while_stmt(common::model::diagram_element::id_t i);
+
+    void add_do_stmt(common::model::diagram_element::id_t i);
+    void end_do_stmt(common::model::diagram_element::id_t i);
+
+    void add_for_stmt(common::model::diagram_element::id_t i);
+    void end_for_stmt(common::model::diagram_element::id_t i);
+
+    bool started() const;
+    void started(bool s);
+
+    std::map<common::model::diagram_element::id_t, activity> &sequences();
+
+    const std::map<common::model::diagram_element::id_t, activity> &
+    sequences() const;
 
     std::map<common::model::diagram_element::id_t, std::unique_ptr<participant>>
-        participants;
+        &participants();
+
+    const std::map<common::model::diagram_element::id_t,
+        std::unique_ptr<participant>> &
+    participants() const;
+
+    std::set<common::model::diagram_element::id_t> &active_participants();
+
+    const std::set<common::model::diagram_element::id_t> &
+    active_participants() const;
+
+private:
+    bool started_{false};
+
+    std::map<common::model::diagram_element::id_t, activity> sequences_;
+
+    std::map<common::model::diagram_element::id_t, std::unique_ptr<participant>>
+        participants_;
 
     std::set<common::model::diagram_element::id_t> active_participants_;
-
-    void add_if_stmt(
-        const common::model::diagram_element::id_t current_caller_id,
-        common::model::message_t type);
-    void end_if_stmt(
-        const common::model::diagram_element::id_t current_caller_id,
-        common::model::message_t type);
-
-    void add_loop_stmt(
-        const common::model::diagram_element::id_t current_caller_id,
-        common::model::message_t type);
-    void end_loop_stmt(
-        const common::model::diagram_element::id_t current_caller_id,
-        common::model::message_t type);
-
-    void add_while_stmt(const common::model::diagram_element::id_t i);
-    void end_while_stmt(const common::model::diagram_element::id_t i);
-
-    void add_do_stmt(const common::model::diagram_element::id_t i);
-    void end_do_stmt(const common::model::diagram_element::id_t i);
-
-    void add_for_stmt(const common::model::diagram_element::id_t i);
-    void end_for_stmt(const common::model::diagram_element::id_t i);
 };
 
 }
diff --git a/src/sequence_diagram/model/message.cc b/src/sequence_diagram/model/message.cc
index ef92071b..93f23937 100644
--- a/src/sequence_diagram/model/message.cc
+++ b/src/sequence_diagram/model/message.cc
@@ -20,4 +20,34 @@
 
 namespace clanguml::sequence_diagram::model {
 
+message::message(
+    common::model::message_t type, common::model::diagram_element::id_t from)
+    : type_{type}
+    , from_{from}
+{
+}
+
+void message::set_type(common::model::message_t t) { type_ = t; }
+
+common::model::message_t message::type() const { return type_; }
+
+void message::set_from(common::model::diagram_element::id_t f) { from_ = f; }
+
+common::model::diagram_element::id_t message::from() const { return from_; }
+
+void message::set_to(common::model::diagram_element::id_t t) { to_ = t; }
+
+common::model::diagram_element::id_t message::to() const { return to_; }
+
+void message::set_message_name(std::string name)
+{
+    message_name_ = std::move(name);
+}
+
+const std::string &message::message_name() const { return message_name_; }
+
+void message::set_return_type(std::string t) { return_type_ = std::move(t); }
+
+const std::string &message::return_type() const { return return_type_; }
+
 }
diff --git a/src/sequence_diagram/model/message.h b/src/sequence_diagram/model/message.h
index b12614c0..48c6ca0a 100644
--- a/src/sequence_diagram/model/message.h
+++ b/src/sequence_diagram/model/message.h
@@ -25,22 +25,40 @@
 
 namespace clanguml::sequence_diagram::model {
 
-struct message : public common::model::diagram_element {
-    message()
-        : from{}
-        , to{}
-        , message_name{}
-        , return_type{}
-    {
-    }
+class message : public common::model::diagram_element {
+public:
+    message() = default;
 
-    common::model::message_t type;
-    common::model::diagram_element::id_t from;
-    common::model::diagram_element::id_t to;
+    message(common::model::message_t type,
+        common::model::diagram_element::id_t from);
 
-    std::string message_name;
+    void set_type(common::model::message_t t);
+    common::model::message_t type() const;
 
-    std::string return_type;
+    void set_from(common::model::diagram_element::id_t f);
+    common::model::diagram_element::id_t from() const;
+
+    void set_to(common::model::diagram_element::id_t t);
+    common::model::diagram_element::id_t to() const;
+
+    void set_message_name(std::string name);
+    const std::string &message_name() const;
+
+    void set_return_type(std::string t);
+    const std::string &return_type() const;
+
+private:
+    common::model::message_t type_{common::model::message_t::kNone};
+
+    common::model::diagram_element::id_t from_{};
+
+    common::model::diagram_element::id_t to_{};
+
+    // This is only for better verbose messages, we cannot rely on this
+    // always
+    std::string message_name_{};
+
+    std::string return_type_{};
 };
 
 }
diff --git a/src/sequence_diagram/visitor/translation_unit_visitor.cc b/src/sequence_diagram/visitor/translation_unit_visitor.cc
index 841ead85..45d48134 100644
--- a/src/sequence_diagram/visitor/translation_unit_visitor.cc
+++ b/src/sequence_diagram/visitor/translation_unit_visitor.cc
@@ -571,12 +571,9 @@ bool translation_unit_visitor::TraverseCompoundStmt(clang::CompoundStmt *stmt)
             const auto current_caller_id = context().caller_id();
 
             if (current_caller_id) {
-                message m;
-                m.from = current_caller_id;
-                m.type = message_t::kElse;
-
-                diagram().sequences[current_caller_id].messages.emplace_back(
-                    std::move(m));
+                diagram()
+                    .get_activity(current_caller_id)
+                    .add_message({message_t::kElse, current_caller_id});
             }
         }
     }
@@ -585,12 +582,9 @@ bool translation_unit_visitor::TraverseCompoundStmt(clang::CompoundStmt *stmt)
             const auto current_caller_id = context().caller_id();
 
             if (current_caller_id) {
-                message m;
-                m.from = current_caller_id;
-                m.type = message_t::kElse;
-
-                diagram().sequences[current_caller_id].messages.emplace_back(
-                    std::move(m));
+                diagram()
+                    .get_activity(current_caller_id)
+                    .add_message({message_t::kElse, current_caller_id});
             }
         }
     }
@@ -763,9 +757,7 @@ bool translation_unit_visitor::VisitCallExpr(clang::CallExpr *expr)
         expr->getBeginLoc().printToString(source_manager()),
         context().caller_id());
 
-    message m;
-    m.type = message_t::kCall;
-    m.from = context().caller_id();
+    message m{message_t::kCall, context().caller_id()};
 
     set_source_location(*expr, m);
 
@@ -774,7 +766,7 @@ bool translation_unit_visitor::VisitCallExpr(clang::CallExpr *expr)
     // Unless the lambda is declared in a function or method call
     if (context().lambda_caller_id() != 0) {
         if (context().current_function_call_expr_ == nullptr) {
-            m.from = context().lambda_caller_id();
+            m.set_from(context().lambda_caller_id());
         }
         else {
             LOG_DBG("Current lambda declaration is passed to a method or "
@@ -847,24 +839,21 @@ bool translation_unit_visitor::VisitCallExpr(clang::CallExpr *expr)
     // const auto &return_type =
     //    function_call_expr->getCallReturnType(current_ast_context);
     // m.return_type = return_type.getAsString();
-    m.return_type = "";
 
-    if (m.from > 0 && m.to > 0) {
-        if (diagram().sequences.find(m.from) == diagram().sequences.end()) {
-            activity a;
-            a.from = m.from;
-            diagram().sequences.insert({m.from, std::move(a)});
+    if (m.from() > 0 && m.to() > 0) {
+        if (diagram().sequences().find(m.from()) ==
+            diagram().sequences().end()) {
+            activity a{m.from()};
+            diagram().sequences().insert({m.from(), std::move(a)});
         }
 
-        diagram().add_active_participant(m.from);
-        diagram().add_active_participant(m.to);
+        diagram().add_active_participant(m.from());
+        diagram().add_active_participant(m.to());
 
-        LOG_DBG("Found call {} from {} [{}] to {} [{}] ", m.message_name,
-            m.from, m.from, m.to, m.to);
+        LOG_DBG("Found call {} from {} [{}] to {} [{}] ", m.message_name(),
+            m.from(), m.from(), m.to(), m.to());
 
-        diagram().sequences[m.from].messages.emplace_back(std::move(m));
-
-        assert(!diagram().sequences.empty());
+        diagram().get_activity(m.from()).add_message(std::move(m));
     }
 
     return true;
@@ -876,6 +865,11 @@ bool translation_unit_visitor::process_operator_call_expression(
     if (operator_call_expr->getCalleeDecl() == nullptr)
         return false;
 
+    // For now we only handle call overloaded operators
+    if (operator_call_expr->getOperator() !=
+        clang::OverloadedOperatorKind::OO_Call)
+        return false;
+
     LOG_DBG("Operator '{}' call expression to {} at {}",
         getOperatorSpelling(operator_call_expr->getOperator()),
         operator_call_expr->getCalleeDecl()->getID(),
@@ -883,14 +877,14 @@ bool translation_unit_visitor::process_operator_call_expression(
 
     auto maybe_id = get_unique_id(operator_call_expr->getCalleeDecl()->getID());
     if (maybe_id.has_value()) {
-        m.to = maybe_id.value();
+        m.set_to(maybe_id.value());
     }
     else {
-        m.to = operator_call_expr->getCalleeDecl()->getID();
+        m.set_to(operator_call_expr->getCalleeDecl()->getID());
     }
 
-    m.message_name = fmt::format(
-        "operator{}", getOperatorSpelling(operator_call_expr->getOperator()));
+    m.set_message_name(fmt::format(
+        "operator{}", getOperatorSpelling(operator_call_expr->getOperator())));
 
     return true;
 }
@@ -912,15 +906,13 @@ bool translation_unit_visitor::process_class_method_call_expression(
             diagram().should_include(callee_decl->getQualifiedNameAsString())))
         return false;
 
-    m.to = method_decl->getID();
-
-    m.message_name = method_decl->getNameAsString();
-
-    m.return_type =
+    m.set_to(method_decl->getID());
+    m.set_message_name(method_decl->getNameAsString());
+    m.set_return_type(
         method_call_expr->getCallReturnType(*context().get_ast_context())
-            .getAsString();
+            .getAsString());
 
-    LOG_DBG("Set callee method id {} for method name {}", m.to,
+    LOG_DBG("Set callee method id {} for method name {}", m.to(),
         method_decl->getQualifiedNameAsString());
 
     diagram().add_active_participant(method_decl->getID());
@@ -951,13 +943,13 @@ bool translation_unit_visitor::process_class_template_method_call_expression(
                 get_participant(template_declaration).value().full_name(false) +
                 "::" + dependent_member_callee->getMember().getAsString();
 
-            for (const auto &[id, p] : diagram().participants) {
+            for (const auto &[id, p] : diagram().participants()) {
                 const auto p_full_name = p->full_name(false);
 
                 if (p_full_name.find(callee_method_full_name + "(") == 0) {
                     // TODO: This selects the first matching template method
                     //       without considering arguments!!!
-                    m.to = id;
+                    m.set_to(id);
                     break;
                 }
             }
@@ -975,12 +967,12 @@ bool translation_unit_visitor::process_class_template_method_call_expression(
                                               .full_name(false) +
                     "::" + dependent_member_callee->getMember().getAsString();
 
-                for (const auto &[id, p] : diagram().participants) {
+                for (const auto &[id, p] : diagram().participants()) {
                     const auto p_full_name = p->full_name(false);
                     if (p_full_name.find(callee_method_full_name + "(") == 0) {
                         // TODO: This selects the first matching template method
                         //       without considering arguments!!!
-                        m.to = id;
+                        m.set_to(id);
                         break;
                     }
                 }
@@ -989,8 +981,7 @@ bool translation_unit_visitor::process_class_template_method_call_expression(
                 return false;
         }
 
-        m.message_name = dependent_member_callee->getMember().getAsString();
-        m.return_type = "";
+        m.set_message_name(dependent_member_callee->getMember().getAsString());
 
         if (get_unique_id(template_declaration->getID()))
             diagram().add_active_participant(
@@ -1027,14 +1018,14 @@ bool translation_unit_visitor::process_function_call_expression(
 
     if (!get_unique_id(callee_function->getID()).has_value()) {
         // This is hopefully not an interesting call...
-        m.to = callee_function->getID();
+        m.set_to(callee_function->getID());
     }
     else {
-        m.to = get_unique_id(callee_function->getID()).value();
+        m.set_to(get_unique_id(callee_function->getID()).value());
     }
 
     auto message_name = callee_name;
-    m.message_name = message_name.substr(0, message_name.size() - 2);
+    m.set_message_name(message_name.substr(0, message_name.size() - 2));
 
     if (f_ptr)
         diagram().add_participant(std::move(f_ptr));
@@ -1057,9 +1048,9 @@ bool translation_unit_visitor::process_unresolved_lookup_call_expression(
                     clang::dyn_cast_or_null<clang::FunctionTemplateDecl>(decl);
 
                 if (!get_unique_id(ftd->getID()).has_value())
-                    m.to = ftd->getID();
+                    m.set_to(ftd->getID());
                 else {
-                    m.to = get_unique_id(ftd->getID()).value();
+                    m.set_to(get_unique_id(ftd->getID()).value());
                 }
 
                 break;
@@ -1878,7 +1869,7 @@ translation_unit_visitor::build_template_instantiation(
     int best_match{};
     common::model::diagram_element::id_t best_match_id{0};
 
-    for (const auto &[id, c] : diagram().participants) {
+    for (const auto &[id, c] : diagram().participants()) {
         const auto *participant_as_class =
             dynamic_cast<model::class_ *>(c.get());
         if ((participant_as_class != nullptr) &&
@@ -1981,9 +1972,9 @@ std::string translation_unit_visitor::make_lambda_name(
 
 void translation_unit_visitor::finalize()
 {
-    decltype(diagram().active_participants_) active_participants_unique;
+    std::set<common::model::diagram_element::id_t> active_participants_unique;
 
-    for (auto id : diagram().active_participants_) {
+    for (auto id : diagram().active_participants()) {
         if (local_ast_id_map_.find(id) != local_ast_id_map_.end()) {
             active_participants_unique.emplace(local_ast_id_map_.at(id));
         }
@@ -1992,12 +1983,12 @@ void translation_unit_visitor::finalize()
         }
     }
 
-    diagram().active_participants_ = std::move(active_participants_unique);
+    diagram().active_participants() = std::move(active_participants_unique);
 
-    for (auto &[id, activity] : diagram().sequences) {
-        for (auto &m : activity.messages) {
-            if (local_ast_id_map_.find(m.to) != local_ast_id_map_.end()) {
-                m.to = local_ast_id_map_.at(m.to);
+    for (auto &[id, activity] : diagram().sequences()) {
+        for (auto &m : activity.messages()) {
+            if (local_ast_id_map_.find(m.to()) != local_ast_id_map_.end()) {
+                m.set_to(local_ast_id_map_.at(m.to()));
             }
         }
     }
diff --git a/src/sequence_diagram/visitor/translation_unit_visitor.h b/src/sequence_diagram/visitor/translation_unit_visitor.h
index 7ff0666f..83eb4dbc 100644
--- a/src/sequence_diagram/visitor/translation_unit_visitor.h
+++ b/src/sequence_diagram/visitor/translation_unit_visitor.h
@@ -116,22 +116,22 @@ public:
     common::optional_ref<T> get_participant(
         const common::model::diagram_element::id_t id)
     {
-        if (diagram().participants.find(id) == diagram().participants.end())
+        if (diagram().participants().find(id) == diagram().participants().end())
             return {};
 
         return common::optional_ref<T>(
-            *(static_cast<T *>(diagram().participants.at(id).get())));
+            *(static_cast<T *>(diagram().participants().at(id).get())));
     }
 
     template <typename T = model::participant>
     const common::optional_ref<T> get_participant(
         const common::model::diagram_element::id_t id) const
     {
-        if (diagram().participants.find(id) == diagram().participants.end())
+        if (diagram().participants().find(id) == diagram().participants().end())
             return {};
 
         return common::optional_ref<T>(
-            *(static_cast<T *>(diagram().participants.at(id).get())));
+            *(static_cast<T *>(diagram().participants().at(id).get())));
     }
 
     /// Store the mapping from local clang entity id (obtained using