Added support for Objective-C categories in class diagrams (#296)

src/class_diagram/generators/json/class_diagram_generator.cc

@@ -127,6 +127,7 @@ void to_json(nlohmann::json &j, const objc_interface &c)
 {
     j = dynamic_cast<const common::model::element &>(c);
     j["is_protocol"] = c.is_protocol();
+    j["is_category"] = c.is_category();
 
     j["members"] = c.members();
     j["methods"] = c.methods();

src/class_diagram/generators/mermaid/class_diagram_generator.cc

@@ -633,6 +633,8 @@ void generator::generate(const objc_interface &c, std::ostream &ostr) const
 
     if (c.is_protocol())
         ostr << indent(2) << "<<ObjC Protocol>>\n";
+    else if (c.is_protocol())
+        ostr << indent(2) << "<<ObjC Category>>\n";
     else
         ostr << indent(2) << "<<ObjC Interface>>\n";
 

src/class_diagram/generators/plantuml/class_diagram_generator.cc

@@ -386,8 +386,13 @@ void generator::generate(const objc_interface &c, std::ostream &ostr) const
 
     ostr << class_type << " " << c.alias();
 
-    ostr << " "
-         << "<<ObjC>>";
+    ostr << " ";
+    if (c.is_protocol())
+        ostr << "<<ObjC Protocol>>";
+    else if (c.is_category())
+        ostr << "<<ObjC Category>>";
+    else
+        ostr << "<<ObjC Interface>>";
 
     if (config().generate_links) {
         common_generator<diagram_config, diagram_model>::generate_link(ostr, c);

src/class_diagram/model/objc_interface.cc

@@ -44,6 +44,10 @@ bool objc_interface::is_protocol() const { return is_protocol_; }
 
 void objc_interface::is_protocol(bool ip) { is_protocol_ = ip; }
 
+bool objc_interface::is_category() const { return is_category_; }
+
+void objc_interface::is_category(bool cat) { is_category_ = cat; }
+
 void objc_interface::add_member(objc_member &&member)
 {
     members_.emplace_back(std::move(member));

src/class_diagram/model/objc_interface.h

@@ -44,6 +44,8 @@ class objc_interface : public common::model::element,
     {
         if (is_protocol())
             return "objc_protocol";
+        else if (is_category())
+            return "objc_category";
         else
             return "objc_interface";
     }
@@ -64,6 +66,10 @@ class objc_interface : public common::model::element,
 
     void is_protocol(bool ip);
 
+    bool is_category() const;
+
+    void is_category(bool cat);
+
     /**
      * @brief Get Doxygen link to documentation page for this element.
      *
@@ -75,6 +81,7 @@ class objc_interface : public common::model::element,
     std::vector<objc_member> members_;
     std::vector<objc_method> methods_;
     bool is_protocol_{false};
+    bool is_category_{false};
 };
 
 } // namespace clanguml::class_diagram::model

src/class_diagram/visitor/translation_unit_visitor.cc

@@ -386,13 +386,53 @@ bool translation_unit_visitor::VisitRecordDecl(clang::RecordDecl *rec)
     return true;
 }
 
+bool translation_unit_visitor::VisitObjCCategoryDecl(
+    clang::ObjCCategoryDecl *decl)
+{
+    if (!should_include(decl))
+        return true;
+
+    LOG_DBG("= Visiting ObjC category declaration {} at {}",
+        decl->getQualifiedNameAsString(),
+        decl->getLocation().printToString(source_manager()));
+
+    auto category_ptr = create_objc_category_declaration(decl);
+
+    if (!category_ptr)
+        return true;
+
+    const auto category_id = category_ptr->id();
+
+    id_mapper().add(decl->getID(), category_id);
+
+    auto &category_model =
+        diagram().find<objc_interface>(category_id).has_value()
+        ? *diagram().find<objc_interface>(category_id).get()
+        : *category_ptr;
+
+    process_objc_category_declaration(*decl, category_model);
+
+    if (diagram().should_include(category_model)) {
+        LOG_DBG("Adding ObjC category {} with id {}",
+            category_model.full_name(false), category_model.id());
+
+        add_objc_interface(std::move(category_ptr));
+    }
+    else {
+        LOG_DBG("Skipping ObjC category {} with id {}",
+            category_model.full_name(), category_model.id());
+    }
+
+    return true;
+}
+
 bool translation_unit_visitor::VisitObjCProtocolDecl(
     clang::ObjCProtocolDecl *decl)
 {
     if (!should_include(decl))
         return true;
 
-    LOG_ERROR("= Visiting ObjC protocol declaration {} at {}",
+    LOG_DBG("= Visiting ObjC protocol declaration {} at {}",
         decl->getQualifiedNameAsString(),
         decl->getLocation().printToString(source_manager()));
 
@@ -432,7 +472,7 @@ bool translation_unit_visitor::VisitObjCInterfaceDecl(
     if (!should_include(decl))
         return true;
 
-    LOG_ERROR("= Visiting ObjC interface declaration {} at {}",
+    LOG_DBG("= Visiting ObjC interface declaration {} at {}",
         decl->getQualifiedNameAsString(),
         decl->getLocation().printToString(source_manager()));
 
@@ -939,6 +979,34 @@ std::unique_ptr<class_> translation_unit_visitor::create_class_declaration(
     return c_ptr;
 }
 
+std::unique_ptr<clanguml::class_diagram::model::objc_interface>
+translation_unit_visitor::create_objc_category_declaration(
+    clang::ObjCCategoryDecl *decl)
+{
+    assert(decl != nullptr);
+
+    if (!should_include(decl))
+        return {};
+
+    auto c_ptr{std::make_unique<class_diagram::model::objc_interface>(
+        config().using_namespace())};
+    auto &c = *c_ptr;
+
+    c.set_name(decl->getNameAsString());
+    c.set_id(common::to_id(*decl));
+    c.is_category(true);
+
+    process_comment(*decl, c);
+    set_source_location(*decl, c);
+
+    if (c.skip())
+        return {};
+
+    c.set_style(c.style_spec());
+
+    return c_ptr;
+}
+
 std::unique_ptr<clanguml::class_diagram::model::objc_interface>
 translation_unit_visitor::create_objc_protocol_declaration(
     clang::ObjCProtocolDecl *decl)
@@ -1090,9 +1158,40 @@ void translation_unit_visitor::process_class_declaration(
     c.complete(true);
 }
 
+void translation_unit_visitor::process_objc_category_declaration(
+    const clang::ObjCCategoryDecl &cls, objc_interface &c)
+{
+    assert(c.is_category());
+
+    // Iterate over class methods (both regular and static)
+    for (const auto *method : cls.methods()) {
+        if (method != nullptr) {
+            process_objc_method(*method, c);
+        }
+    }
+
+    // Add relationship to the ObjC Interface being extended by this
+    // category
+    if (cls.getClassInterface() != nullptr) {
+        eid_t objc_interface_id = common::to_id(*cls.getClassInterface());
+        common::model::relationship r{
+            relationship_t::kInstantiation, objc_interface_id, access_t::kNone};
+
+        LOG_DBG("Found protocol {} [{}] for ObjC interface {}",
+            cls.getClassInterface()->getNameAsString(),
+            objc_interface_id.value(), c.name());
+
+        c.add_relationship(std::move(r));
+    }
+
+    c.complete(true);
+}
+
 void translation_unit_visitor::process_objc_protocol_declaration(
     const clang::ObjCProtocolDecl &cls, objc_interface &c)
 {
+    assert(c.is_protocol());
+
     // Iterate over class methods (both regular and static)
     for (const auto *method : cls.methods()) {
         if (method != nullptr) {
@@ -1106,19 +1205,15 @@ void translation_unit_visitor::process_objc_protocol_declaration(
 void translation_unit_visitor::process_objc_interface_declaration(
     const clang::ObjCInterfaceDecl &cls, objc_interface &c)
 {
+    assert(!c.is_protocol() && !c.is_category());
+
     // Iterate over class methods (both regular and static)
     for (const auto *method : cls.methods()) {
         if (method != nullptr) {
             process_objc_method(*method, c);
         }
     }
 
-    //    for (const auto *member : cls.properties()) {
-    //        if (member != nullptr) {
-    //            process_objc_property(*member, c);
-    //        }
-    //    }
-
     for (const auto *ivar : cls.ivars()) {
         if (ivar != nullptr) {
             process_objc_ivar(*ivar, c);
@@ -1257,34 +1352,6 @@ void translation_unit_visitor::process_objc_ivar(
     c.add_member(std::move(field));
 }
 
-void translation_unit_visitor::process_objc_property(
-    const clang::ObjCPropertyDecl &mf, objc_interface &c)
-{
-    LOG_DBG("== Visiting ObjC property {}", mf.getNameAsString());
-
-    // Default hint for relationship is aggregation
-    //    auto relationship_hint = relationship_t::kAggregation;
-
-    auto field_type = mf.getType();
-    auto field_type_str =
-        common::to_string(field_type, mf.getASTContext(), false);
-
-    auto type_name = common::to_string(field_type, mf.getASTContext());
-
-    const auto field_name = mf.getNameAsString();
-
-    objc_member field{common::access_specifier_to_access_t(mf.getAccess()),
-        field_name, field_type_str};
-
-    process_comment(mf, field);
-    set_source_location(mf, field);
-
-    if (field.skip())
-        return;
-
-    c.add_member(std::move(field));
-}
-
 void translation_unit_visitor::process_class_bases(
     const clang::CXXRecordDecl *cls, class_ &c)
 {