Added loop statement sequence diagram support

src/common/model/enums.h

@@ -39,7 +39,20 @@ enum class relationship_t {
     kDependency
 };
 
-enum class message_t { kCall, kReturn, kIf, kElse, kElseIf, kIfEnd };
+enum class message_t {
+    kCall,
+    kReturn,
+    kIf,
+    kElse,
+    kElseIf,
+    kIfEnd,
+    kWhile,
+    kWhileEnd,
+    kDo,
+    kDoEnd,
+    kFor,
+    kForEnd
+};
 
 std::string to_string(relationship_t r);
 

src/sequence_diagram/generators/plantuml/sequence_diagram_generator.cc

@@ -159,6 +159,24 @@ void generator::generate_activity(const activity &a, std::ostream &ostr,
         else if (m.type == message_t::kIfEnd) {
             ostr << "end\n";
         }
+        else if (m.type == message_t::kWhile) {
+            ostr << "loop\n";
+        }
+        else if (m.type == message_t::kWhileEnd) {
+            ostr << "end\n";
+        }
+        else if (m.type == message_t::kFor) {
+            ostr << "loop\n";
+        }
+        else if (m.type == message_t::kForEnd) {
+            ostr << "end\n";
+        }
+        else if (m.type == message_t::kDo) {
+            ostr << "loop\n";
+        }
+        else if (m.type == message_t::kDoEnd) {
+            ostr << "end\n";
+        }
     }
 }
 

src/sequence_diagram/model/diagram.cc

@@ -102,6 +102,96 @@ void diagram::print() const
     }
 }
 
+void diagram::add_for_stmt(
+    const common::model::diagram_element::id_t current_caller_id)
+{
+    add_loop_stmt(current_caller_id, common::model::message_t::kFor);
+}
+
+void diagram::end_for_stmt(
+    const common::model::diagram_element::id_t current_caller_id)
+{
+    end_loop_stmt(current_caller_id, common::model::message_t::kForEnd);
+}
+
+void diagram::add_while_stmt(
+    const common::model::diagram_element::id_t current_caller_id)
+{
+    add_loop_stmt(current_caller_id, common::model::message_t::kWhile);
+}
+
+void diagram::end_while_stmt(
+    const common::model::diagram_element::id_t current_caller_id)
+{
+    end_loop_stmt(current_caller_id, common::model::message_t::kWhileEnd);
+}
+
+void diagram::add_do_stmt(
+    const common::model::diagram_element::id_t current_caller_id)
+{
+    add_loop_stmt(current_caller_id, common::model::message_t::kDo);
+}
+
+void diagram::end_do_stmt(
+    const common::model::diagram_element::id_t current_caller_id)
+{
+    end_loop_stmt(current_caller_id, common::model::message_t::kDoEnd);
+}
+
+void diagram::add_loop_stmt(
+    const common::model::diagram_element::id_t current_caller_id,
+    common::model::message_t type)
+{
+    using clanguml::common::model::message_t;
+
+    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)});
+    }
+
+    message m;
+    m.from = current_caller_id;
+    m.type = type;
+
+    sequences[current_caller_id].messages.emplace_back(std::move(m));
+}
+
+void diagram::end_loop_stmt(
+    const common::model::diagram_element::id_t current_caller_id,
+    common::model::message_t type)
+{
+    using clanguml::common::model::message_t;
+
+    if (current_caller_id == 0)
+        return;
+
+    message m;
+    m.from = current_caller_id;
+    m.type = type;
+
+    message_t loop_type = message_t::kWhile;
+
+    if (type == message_t::kForEnd)
+        loop_type = message_t::kFor;
+    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 (current_messages.back().type == loop_type) {
+            current_messages.pop_back();
+        }
+        else {
+            current_messages.emplace_back(std::move(m));
+        }
+    }
+}
+
 }
 
 namespace clanguml::common::model {

src/sequence_diagram/model/diagram.h

@@ -102,6 +102,19 @@ public:
         participants;
 
     std::set<common::model::diagram_element::id_t> active_participants_;
+
+    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);
 };
 
 }

src/sequence_diagram/visitor/call_expression_context.h

@@ -107,6 +107,25 @@ struct call_expression_context {
         return elseif_stmt_stack_.top();
     }
 
+    clang::Stmt *current_loopstmt() const
+    {
+        if (loop_stmt_stack_.empty())
+            return nullptr;
+
+        return loop_stmt_stack_.top();
+    }
+
+    void enter_loopstmt(clang::Stmt *stmt)
+    {
+        return loop_stmt_stack_.push(stmt);
+    }
+
+    void leave_loopstmt()
+    {
+        if (loop_stmt_stack_.empty())
+            return loop_stmt_stack_.pop();
+    }
+
     clang::CXXRecordDecl *current_class_decl_;
     clang::ClassTemplateDecl *current_class_template_decl_;
     clang::ClassTemplateSpecializationDecl
@@ -122,6 +141,8 @@ private:
     std::stack<std::int64_t> current_lambda_caller_id_;
     std::stack<clang::IfStmt *> if_stmt_stack_;
     std::stack<clang::IfStmt *> elseif_stmt_stack_;
+
+    std::stack<clang::Stmt *> loop_stmt_stack_;
 };
 
 }

src/sequence_diagram/visitor/translation_unit_visitor.cc

@@ -566,7 +566,7 @@ bool translation_unit_visitor::TraverseCompoundStmt(clang::CompoundStmt *stmt)
             }
         }
     }
-    else if(current_ifstmt != nullptr) {
+    else if (current_ifstmt != nullptr) {
         if (current_ifstmt->getElse() == stmt) {
             const auto current_caller_id = context().caller_id();
 
@@ -671,6 +671,92 @@ bool translation_unit_visitor::TraverseIfStmt(clang::IfStmt *stmt)
     return true;
 }
 
+bool translation_unit_visitor::TraverseWhileStmt(clang::WhileStmt *stmt)
+{
+    using clanguml::common::model::message_t;
+    using clanguml::sequence_diagram::model::activity;
+    using clanguml::sequence_diagram::model::message;
+
+    context().enter_loopstmt(stmt);
+
+    const auto current_caller_id = context().caller_id();
+
+    diagram().add_while_stmt(current_caller_id);
+
+    RecursiveASTVisitor<translation_unit_visitor>::TraverseWhileStmt(stmt);
+
+    diagram().end_while_stmt(current_caller_id);
+
+    context().leave_loopstmt();
+
+    return true;
+}
+
+bool translation_unit_visitor::TraverseDoStmt(clang::DoStmt *stmt)
+{
+    using clanguml::common::model::message_t;
+    using clanguml::sequence_diagram::model::activity;
+    using clanguml::sequence_diagram::model::message;
+
+    context().enter_loopstmt(stmt);
+
+    const auto current_caller_id = context().caller_id();
+
+    diagram().add_do_stmt(current_caller_id);
+
+    RecursiveASTVisitor<translation_unit_visitor>::TraverseDoStmt(stmt);
+
+    context().leave_loopstmt();
+
+    diagram().end_do_stmt(current_caller_id);
+
+    return true;
+}
+
+bool translation_unit_visitor::TraverseForStmt(clang::ForStmt *stmt)
+{
+    using clanguml::common::model::message_t;
+    using clanguml::sequence_diagram::model::activity;
+    using clanguml::sequence_diagram::model::message;
+
+    context().enter_loopstmt(stmt);
+
+    const auto current_caller_id = context().caller_id();
+
+    diagram().add_for_stmt(current_caller_id);
+
+    RecursiveASTVisitor<translation_unit_visitor>::TraverseForStmt(stmt);
+
+    context().leave_loopstmt();
+
+    diagram().end_for_stmt(current_caller_id);
+
+    return true;
+}
+
+bool translation_unit_visitor::TraverseCXXForRangeStmt(
+    clang::CXXForRangeStmt *stmt)
+{
+    using clanguml::common::model::message_t;
+    using clanguml::sequence_diagram::model::activity;
+    using clanguml::sequence_diagram::model::message;
+
+    context().enter_loopstmt(stmt);
+
+    const auto current_caller_id = context().caller_id();
+
+    diagram().add_for_stmt(current_caller_id);
+
+    RecursiveASTVisitor<translation_unit_visitor>::TraverseCXXForRangeStmt(
+        stmt);
+
+    context().leave_loopstmt();
+
+    diagram().end_for_stmt(current_caller_id);
+
+    return true;
+}
+
 bool translation_unit_visitor::VisitCallExpr(clang::CallExpr *expr)
 {
     using clanguml::common::model::message_t;

src/sequence_diagram/visitor/translation_unit_visitor.h

@@ -68,6 +68,14 @@ public:
 
     bool TraverseIfStmt(clang::IfStmt *stmt);
 
+    bool TraverseWhileStmt(clang::WhileStmt *stmt);
+
+    bool TraverseDoStmt(clang::DoStmt *stmt);
+
+    bool TraverseForStmt(clang::ForStmt *stmt);
+
+    bool TraverseCXXForRangeStmt(clang::CXXForRangeStmt *stmt);
+
     clanguml::sequence_diagram::model::diagram &diagram();
 
     const clanguml::sequence_diagram::model::diagram &diagram() const;

tests/t20021/.clang-uml

@@ -0,0 +1,14 @@
+compilation_database_dir: ..
+output_directory: puml
+diagrams:
+  t20021_sequence:
+    type: sequence
+    glob:
+      - ../../tests/t20021/t20021.cc
+    include:
+      namespaces:
+        - clanguml::t20021
+    using_namespace:
+      - clanguml::t20021
+    start_from:
+      - function: "clanguml::t20021::tmain()"

tests/t20021/t20021.cc

@@ -0,0 +1,40 @@
+#include <vector>
+
+namespace clanguml {
+namespace t20021 {
+struct A {
+    int a1() { return 0; }
+    int a2() { return 1; }
+    int a3() { return 2; }
+};
+
+struct B {
+    void log() { }
+
+    int b1() const { return 3; }
+    int b2() const { return 4; }
+};
+
+int tmain()
+{
+    A a;
+    std::vector<B> b;
+
+    int i = 10;
+    while (i--) {
+        int j = a.a3();
+        do {
+            for (int l = a.a2(); l > 0; l--)
+                a.a1();
+        } while (j--);
+    }
+
+    int result = 0;
+    for (const auto &bi : b) {
+        result += bi.b2();
+    }
+
+    return b.front().b2() + result;
+}
+}
+}

tests/t20021/test_case.h

@@ -0,0 +1,48 @@
+/**
+ * tests/t20021/test_case.h
+ *
+ * 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.
+ */
+
+TEST_CASE("t20021", "[test-case][sequence]")
+{
+    auto [config, db] = load_config("t20021");
+
+    auto diagram = config.diagrams["t20021_sequence"];
+
+    REQUIRE(diagram->name == "t20021_sequence");
+
+    auto model = generate_sequence_diagram(*db, diagram);
+
+    REQUIRE(model->name() == "t20021_sequence");
+
+    auto puml = generate_sequence_puml(diagram, *model);
+    AliasMatcher _A(puml);
+
+    REQUIRE_THAT(puml, StartsWith("@startuml"));
+    REQUIRE_THAT(puml, EndsWith("@enduml\n"));
+
+    
+    // Check if all calls exist
+    REQUIRE_THAT(puml, HasCall(_A("tmain()"), _A("A"), "a1()"));
+    REQUIRE_THAT(puml, HasCall(_A("tmain()"), _A("A"), "a2()"));
+    REQUIRE_THAT(puml, HasCall(_A("tmain()"), _A("A"), "a3()"));
+
+    REQUIRE_THAT(puml, HasCall(_A("tmain()"), _A("B"), "b1()"));
+    REQUIRE_THAT(puml, HasCall(_A("tmain()"), _A("B"), "b2()"));
+
+    save_puml(
+        "./" + config.output_directory() + "/" + diagram->name + ".puml", puml);
+}

tests/test_cases.cc

@@ -267,6 +267,7 @@ using namespace clanguml::test::matchers;
 #include "t20018/test_case.h"
 #include "t20019/test_case.h"
 #include "t20020/test_case.h"
+#include "t20021/test_case.h"
 
 ///
 /// Package diagram tests

tests/test_cases.yaml

@@ -208,6 +208,9 @@ test_cases:
     - name: t20020
       title: If statement sequence diagram test case
       description:
+    - name: t20021
+      title: Loop statements sequence diagram test case
+      description:
   Package diagrams:
     - name: t30001
       title: Basic package diagram test case