diff --git a/src/sequence_diagram/visitor/call_expression_context.cc b/src/sequence_diagram/visitor/call_expression_context.cc
index cd75d268..14699a1f 100644
--- a/src/sequence_diagram/visitor/call_expression_context.cc
+++ b/src/sequence_diagram/visitor/call_expression_context.cc
@@ -186,4 +186,119 @@ void call_expression_context::leave_lambda_expression()
     current_lambda_caller_id_.pop();
 }
 
+clang::IfStmt *call_expression_context::current_ifstmt() const
+{
+    if (if_stmt_stack_.empty())
+        return nullptr;
+
+    return if_stmt_stack_.top();
+}
+
+void call_expression_context::enter_ifstmt(clang::IfStmt *stmt)
+{
+    return if_stmt_stack_.push(stmt);
+}
+
+void call_expression_context::leave_ifstmt()
+{
+    if (!if_stmt_stack_.empty()) {
+        if_stmt_stack_.pop();
+        std::stack<clang::IfStmt *>{}.swap(elseif_stmt_stack_);
+    }
+}
+
+void call_expression_context::enter_elseifstmt(clang::IfStmt *stmt)
+{
+    return elseif_stmt_stack_.push(stmt);
+}
+
+void call_expression_context::leave_elseifstmt()
+{
+    if (elseif_stmt_stack_.empty())
+        return elseif_stmt_stack_.pop();
+}
+
+clang::IfStmt *call_expression_context::current_elseifstmt() const
+{
+    if (elseif_stmt_stack_.empty())
+        return nullptr;
+
+    return elseif_stmt_stack_.top();
+}
+
+clang::Stmt *call_expression_context::current_loopstmt() const
+{
+    if (loop_stmt_stack_.empty())
+        return nullptr;
+
+    return loop_stmt_stack_.top();
+}
+
+void call_expression_context::enter_loopstmt(clang::Stmt *stmt)
+{
+    return loop_stmt_stack_.push(stmt);
+}
+
+void call_expression_context::leave_loopstmt()
+{
+    if (loop_stmt_stack_.empty())
+        return loop_stmt_stack_.pop();
+}
+
+bool call_expression_context::is_expr_in_current_control_statement_condition(
+    const clang::Stmt *stmt) const
+{
+    if (current_ifstmt()) {
+        if (common::is_subexpr_of(current_ifstmt()->getCond(), stmt)) {
+            return true;
+        }
+    }
+
+    if (current_elseifstmt()) {
+        if (common::is_subexpr_of(current_elseifstmt()->getCond(), stmt)) {
+            return true;
+        }
+    }
+
+    if (const auto *loop_stmt = current_loopstmt(); loop_stmt != nullptr) {
+        if (const auto *for_stmt = clang::dyn_cast<clang::ForStmt>(loop_stmt);
+            for_stmt != nullptr) {
+            if (common::is_subexpr_of(for_stmt->getCond(), stmt)) {
+                return true;
+            }
+            if (common::is_subexpr_of(for_stmt->getInit(), stmt)) {
+                return true;
+            }
+            if (common::is_subexpr_of(for_stmt->getInc(), stmt)) {
+                return true;
+            }
+        }
+
+        if (const auto *range_for_stmt =
+                clang::dyn_cast<clang::CXXForRangeStmt>(loop_stmt);
+            range_for_stmt != nullptr) {
+            if (common::is_subexpr_of(range_for_stmt->getRangeInit(), stmt)) {
+                return true;
+            }
+        }
+
+        if (const auto *while_stmt =
+                clang::dyn_cast<clang::WhileStmt>(loop_stmt);
+            while_stmt != nullptr) {
+            if (common::is_subexpr_of(while_stmt->getCond(), stmt)) {
+                return true;
+            }
+        }
+
+        if (const auto *do_stmt = clang::dyn_cast<clang::DoStmt>(loop_stmt);
+            do_stmt != nullptr) {
+            if (common::is_subexpr_of(do_stmt->getCond(), stmt)) {
+                return true;
+            }
+        }
+    }
+
+    return false;
+}
+
 }
\ No newline at end of file
diff --git a/src/sequence_diagram/visitor/call_expression_context.h b/src/sequence_diagram/visitor/call_expression_context.h
index cab68e9d..7e4e54b9 100644
--- a/src/sequence_diagram/visitor/call_expression_context.h
+++ b/src/sequence_diagram/visitor/call_expression_context.h
@@ -16,11 +16,8 @@
  * limitations under the License.
  */
 #pragma once
-//
-//#include "common/visitor/translation_unit_visitor.h"
-//#include "config/config.h"
-//#include "sequence_diagram/model/diagram.h"
 
+#include "common/clang_utils.h"
 #include "util/util.h"
 
 #include <clang/AST/Expr.h>
@@ -70,61 +67,22 @@ struct call_expression_context {
 
     void leave_lambda_expression();
 
-    clang::IfStmt *current_ifstmt() const
-    {
-        if (if_stmt_stack_.empty())
-            return nullptr;
+    clang::IfStmt *current_ifstmt() const;
 
-        return if_stmt_stack_.top();
-    }
+    void enter_ifstmt(clang::IfStmt *stmt);
+    void leave_ifstmt();
 
-    void enter_ifstmt(clang::IfStmt *stmt) { return if_stmt_stack_.push(stmt); }
+    void enter_elseifstmt(clang::IfStmt *stmt);
+    void leave_elseifstmt();
 
-    void leave_ifstmt()
-    {
-        if (!if_stmt_stack_.empty()) {
-            if_stmt_stack_.pop();
-            std::stack<clang::IfStmt *>{}.swap(elseif_stmt_stack_);
-        }
-    }
+    clang::IfStmt *current_elseifstmt() const;
+    clang::Stmt *current_loopstmt() const;
 
-    void enter_elseifstmt(clang::IfStmt *stmt)
-    {
-        return elseif_stmt_stack_.push(stmt);
-    }
+    void enter_loopstmt(clang::Stmt *stmt);
+    void leave_loopstmt();
 
-    void leave_elseifstmt()
-    {
-        if (elseif_stmt_stack_.empty())
-            return elseif_stmt_stack_.pop();
-    }
-
-    clang::IfStmt *current_elseifstmt() const
-    {
-        if (elseif_stmt_stack_.empty())
-            return nullptr;
-
-        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();
-    }
+    bool is_expr_in_current_control_statement_condition(
+        const clang::Stmt *stmt) const;
 
     clang::CXXRecordDecl *current_class_decl_;
     clang::ClassTemplateDecl *current_class_template_decl_;
diff --git a/src/sequence_diagram/visitor/translation_unit_visitor.cc b/src/sequence_diagram/visitor/translation_unit_visitor.cc
index 68f92860..061046cb 100644
--- a/src/sequence_diagram/visitor/translation_unit_visitor.cc
+++ b/src/sequence_diagram/visitor/translation_unit_visitor.cc
@@ -775,20 +775,13 @@ bool translation_unit_visitor::VisitCallExpr(clang::CallExpr *expr)
         }
     }
 
-    if (context().current_ifstmt()) {
-        if (common::is_subexpr_of(
-                context().current_ifstmt()->getCond(), expr)) {
-            m.set_message_scope(common::model::message_scope_t::kCondition);
-        }
-    }
-
-    if (context().current_elseifstmt()) {
-        if (common::is_subexpr_of(
-                context().current_elseifstmt()->getCond(), expr)) {
-            m.set_message_scope(common::model::message_scope_t::kCondition);
-        }
+    if (context().is_expr_in_current_control_statement_condition(expr)) {
+        m.set_message_scope(common::model::message_scope_t::kCondition);
     }
 
+    //
+    // Call to an overloaded operator
+    //
     if (const auto *operator_call_expr =
             clang::dyn_cast_or_null<clang::CXXOperatorCallExpr>(expr);
         operator_call_expr != nullptr) {
diff --git a/tests/t20021/t20021.cc b/tests/t20021/t20021.cc
index 5fe5e19f..4f3210ef 100644
--- a/tests/t20021/t20021.cc
+++ b/tests/t20021/t20021.cc
@@ -15,18 +15,31 @@ struct B {
     int b2() const { return 4; }
 };
 
+struct C {
+    int c1() { return 1; }
+    int c2() { return 2; }
+    int c3() { return 3; }
+    int c4() { return c5(); }
+    int c5() { return 5; }
+
+    std::vector<int> &contents() { return contents_; }
+
+    std::vector<int> contents_;
+};
+
 int tmain()
 {
     A a;
     std::vector<B> b;
+    C c;
 
     int i = 10;
-    while (i--) {
+    while (i -= c.c4()) {
         int j = a.a3();
         do {
-            for (int l = a.a2(); l > 0; l--)
+            for (int l = a.a2(); l > c.c1(); l -= c.c2())
                 a.a1();
-        } while (j--);
+        } while (j -= c.c3());
     }
 
     int result = 0;
@@ -34,6 +47,10 @@ int tmain()
         result += bi.b2();
     }
 
+    for (const auto &ci : c.contents()) {
+        result += ci;
+    }
+
     return b.front().b2() + result;
 }
 }
diff --git a/tests/t20021/test_case.h b/tests/t20021/test_case.h
index 5868931e..f6e57132 100644
--- a/tests/t20021/test_case.h
+++ b/tests/t20021/test_case.h
@@ -36,12 +36,27 @@ TEST_CASE("t20021", "[test-case][sequence]")
 
     // 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, HasCallInControlCondition(_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()"));
 
+    REQUIRE_THAT(
+        puml, HasCallInControlCondition(_A("tmain()"), _A("C"), "c1()"));
+    REQUIRE_THAT(
+        puml, HasCallInControlCondition(_A("tmain()"), _A("C"), "c2()"));
+
+    // TODO: Why is this not working?
+    //    REQUIRE_THAT(
+    //        puml, HasCallInControlCondition(_A("tmain()"), _A("C"), "c3()"));
+    REQUIRE_THAT(
+        puml, HasCallInControlCondition(_A("tmain()"), _A("C"), "c4()"));
+    REQUIRE_THAT(puml, HasCall(_A("C"), _A("C"), "c5()"));
+    REQUIRE_THAT(
+        puml, HasCallInControlCondition(_A("tmain()"), _A("C"), "contents()"));
+
     save_puml(
         "./" + config.output_directory() + "/" + diagram->name + ".puml", puml);
 }
\ No newline at end of file