diff --git a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/AbstractCallGraphEngine.java b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/AbstractCallGraphEngine.java
index 62f786d..75af825 100644
--- a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/AbstractCallGraphEngine.java
+++ b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/AbstractCallGraphEngine.java
@@ -321,32 +321,39 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
             }
 
                 if (declaredType != null) {
-                    List<String> typesToInspect = new ArrayList<>();
-                    if ("inline-instantiation".equals(sourceMethod)) {
-                        typesToInspect.add(declaredType);
-                    } else {
-                        typesToInspect.add(declaredType);
-                        typesToInspect.addAll(context.getImplementations(declaredType));
-                    }
-                    
-                    for (String type : typesToInspect) {
-                        Set<String> visited = new HashSet<>();
-                        TypeDeclaration baseTd = context.getTypeDeclaration(type);
-                        CompilationUnit cuToUse = null;
-                        if (baseTd != null && baseTd.getRoot() instanceof CompilationUnit) {
-                            cuToUse = (CompilationUnit) baseTd.getRoot();
-                        } else {
-                            String firstPathMethod = path.get(0);
-                            String entryClassName = firstPathMethod.substring(0, firstPathMethod.lastIndexOf('.'));
-                            TypeDeclaration entryTd = context.getTypeDeclaration(entryClassName);
-                            if (entryTd != null && entryTd.getRoot() instanceof CompilationUnit) {
-                                cuToUse = (CompilationUnit) entryTd.getRoot();
-                            }
+                    List<String> enumValues = context.getEnumValues(declaredType);
+                    if (enumValues != null && !enumValues.isEmpty()) {
+                        for (String ev : enumValues) {
+                            polymorphicEvents.add(parseEnumSetElement(ev));
                         }
-                        List<String> constants = resolveMethodReturnConstant(type, methodName, 0, visited, cuToUse);
-                        for (String constant : constants) {
-                            if (!polymorphicEvents.contains(constant)) {
-                                polymorphicEvents.add(constant);
+                    } else {
+                        List<String> typesToInspect = new ArrayList<>();
+                        if ("inline-instantiation".equals(sourceMethod)) {
+                            typesToInspect.add(declaredType);
+                        } else {
+                            typesToInspect.add(declaredType);
+                            typesToInspect.addAll(context.getImplementations(declaredType));
+                        }
+                        
+                        for (String type : typesToInspect) {
+                            Set<String> visited = new HashSet<>();
+                            TypeDeclaration baseTd = context.getTypeDeclaration(type);
+                            CompilationUnit cuToUse = null;
+                            if (baseTd != null && baseTd.getRoot() instanceof CompilationUnit) {
+                                cuToUse = (CompilationUnit) baseTd.getRoot();
+                            } else {
+                                String firstPathMethod = path.get(0);
+                                String entryClassName = firstPathMethod.substring(0, firstPathMethod.lastIndexOf('.'));
+                                TypeDeclaration entryTd = context.getTypeDeclaration(entryClassName);
+                                if (entryTd != null && entryTd.getRoot() instanceof CompilationUnit) {
+                                    cuToUse = (CompilationUnit) entryTd.getRoot();
+                                }
+                            }
+                            List<String> constants = resolveMethodReturnConstant(type, methodName, 0, visited, cuToUse);
+                            for (String constant : constants) {
+                                if (!polymorphicEvents.contains(constant)) {
+                                    polymorphicEvents.add(constant);
+                                }
                             }
                         }
                     }