Added option to enable rendering return types in sequence diagrams (fixes #93)

src/sequence_diagram/generators/json/sequence_diagram_generator.cc

@@ -645,16 +645,15 @@ void generator::generate(std::ostream &ostr) const
             generate_activity(
                 m_model.get_activity(start_from), visited_participants);
 
-            json_["sequences"].push_back(std::move(sequence));
-
             block_statements_stack_.pop_back();
 
             if (from.value().type_name() == "method" ||
                 m_config.combine_free_functions_into_file_participants()) {
 
-                // TODO
-                // sequence["return_type"] = from.value().id()
+                sequence["return_type"] = from.value().return_type();
             }
+
+            json_["sequences"].push_back(std::move(sequence));
         }
         else {
             // TODO: Add support for other sequence start location types

src/sequence_diagram/generators/plantuml/sequence_diagram_generator.cc

@@ -128,7 +128,13 @@ void generator::generate_return(const message &m, std::ostream &ostr) const
 
         ostr << to_alias << " "
              << common::generators::plantuml::to_plantuml(message_t::kReturn)
-             << " " << from_alias << '\n';
+             << " " << from_alias;
+
+        if (m_config.generate_return_types()) {
+            ostr << " : //" << m.return_type() << "//";
+        }
+
+        ostr << '\n';
     }
 }
 
@@ -451,7 +457,12 @@ void generator::generate(std::ostream &ostr) const
 
                 if (!from.value().is_void()) {
                     ostr << "[<--"
-                         << " " << from_alias << std::endl;
+                         << " " << from_alias;
+
+                    if (m_config.generate_return_types())
+                        ostr << " : //" << from.value().return_type() << "//";
+
+                    ostr << '\n';
                 }
             }
 

src/sequence_diagram/model/diagram.cc

@@ -312,6 +312,16 @@ void diagram::finalize()
                 assert(block_nest_level >= 0);
             }
             else {
+                if (m.type() == message_t::kCall) {
+                    // Set the message return type based on the callee return
+                    // type
+                    auto to_participant =
+                        get_participant<sequence_diagram::model::function>(
+                            m.to());
+                    if (to_participant.has_value()) {
+                        m.set_return_type(to_participant.value().return_type());
+                    }
+                }
                 block_message_stack.back().push_back(m);
             }
         }

src/sequence_diagram/model/participant.cc

@@ -147,6 +147,10 @@ bool function::is_operator() const { return is_operator_; }
 
 void function::is_operator(bool o) { is_operator_ = o; }
 
+void function::return_type(const std::string &rt) { return_type_ = rt; }
+
+const std::string &function::return_type() const { return return_type_; }
+
 void function::add_parameter(const std::string &a) { parameters_.push_back(a); }
 
 const std::vector<std::string> &function::parameters() const

src/sequence_diagram/model/participant.h

@@ -291,6 +291,20 @@ struct function : public participant {
      */
     void is_operator(bool o);
 
+    /**
+     * @brief Set functions return type
+     *
+     * @param rt Return type
+     */
+    void return_type(const std::string &rt);
+
+    /**
+     * @brief Get function return type
+     *
+     * @return Return type
+     */
+    const std::string &return_type() const;
+
     /**
      * @brief Add a function parameter
      *
@@ -313,6 +327,7 @@ private:
     bool is_void_{false};
     bool is_static_{false};
     bool is_operator_{false};
+    std::string return_type_;
     std::vector<std::string> parameters_;
 };
 

src/sequence_diagram/visitor/translation_unit_visitor.cc

@@ -972,13 +972,6 @@ bool translation_unit_visitor::VisitCallExpr(clang::CallExpr *expr)
         }
     }
 
-    //
-    // This crashes on LLVM <= 12, for now just return empty type
-    //
-    // const auto &return_type =
-    //    function_call_expr->getCallReturnType(current_ast_context);
-    // m.return_type = return_type.getAsString();
-
     if (m.from() > 0 && m.to() > 0) {
         if (diagram().sequences().find(m.from()) ==
             diagram().sequences().end()) {
@@ -1550,6 +1543,10 @@ translation_unit_visitor::build_function_template(
 
     process_template_parameters(declaration, *function_template_model_ptr);
 
+    function_template_model_ptr->return_type(
+        common::to_string(declaration.getAsFunction()->getReturnType(),
+            declaration.getASTContext()));
+
     for (const auto *param : declaration.getTemplatedDecl()->parameters()) {
         function_template_model_ptr->add_parameter(
             simplify_system_template(common::to_string(
@@ -1612,6 +1609,9 @@ std::unique_ptr<model::function> translation_unit_visitor::build_function_model(
     ns.pop_back();
     function_model_ptr->set_namespace(ns);
 
+    function_model_ptr->return_type(common::to_string(
+        declaration.getReturnType(), declaration.getASTContext()));
+
     for (const auto *param : declaration.parameters()) {
         function_model_ptr->add_parameter(
             simplify_system_template(common::to_string(
@@ -2337,6 +2337,9 @@ translation_unit_visitor::create_method_model(clang::CXXMethodDecl *declaration)
         "::" + declaration->getNameAsString());
     method_model_ptr->is_static(declaration->isStatic());
 
+    method_model_ptr->return_type(common::to_string(
+        declaration->getReturnType(), declaration->getASTContext()));
+
     for (const auto *param : declaration->parameters()) {
         method_model_ptr->add_parameter(config().using_namespace().relative(
             simplify_system_template(common::to_string(