3 files changed, 25 insertions, 4 deletions

diff --git a/subprojects/language-semantics/src/main/java/tools/refinery/language/semantics/ModelInitializer.java b/subprojects/language-semantics/src/main/java/tools/refinery/language/semantics/ModelInitializer.java
index c2243f08..8bfc8acd 100644
--- a/subprojects/language-semantics/src/main/java/tools/refinery/language/semantics/ModelInitializer.java
+++ b/subprojects/language-semantics/src/main/java/tools/refinery/language/semantics/ModelInitializer.java
@@ -967,6 +967,7 @@ public class ModelInitializer {
                                         .orElseGet(() -> "::rule" + ruleCount);
                         ruleCount++;
                         var rule = toRule(name, ruleDefinition);
+                        problemTrace.putRuleDefinition(ruleDefinition, rule);
                         switch (ruleDefinition.getKind()) {
                         case DECISION -> storeBuilder.tryGetAdapter(DesignSpaceExplorationBuilder.class)
                                         .ifPresent(dseBuilder -> dseBuilder.transformation(rule));
diff --git a/subprojects/language-semantics/src/main/java/tools/refinery/language/semantics/ProblemTrace.java b/subprojects/language-semantics/src/main/java/tools/refinery/language/semantics/ProblemTrace.java
index b8d0b804..9122cf1f 100644
--- a/subprojects/language-semantics/src/main/java/tools/refinery/language/semantics/ProblemTrace.java
+++ b/subprojects/language-semantics/src/main/java/tools/refinery/language/semantics/ProblemTrace.java
@@ -10,6 +10,8 @@ import org.eclipse.xtext.naming.QualifiedName;
 import tools.refinery.language.model.problem.Node;
 import tools.refinery.language.model.problem.Problem;
 import tools.refinery.language.model.problem.Relation;
+import tools.refinery.language.model.problem.RuleDefinition;
+import tools.refinery.store.dse.transition.Rule;
 import tools.refinery.store.reasoning.representation.AnyPartialSymbol;
 import tools.refinery.store.reasoning.representation.PartialRelation;
 import tools.refinery.store.reasoning.translator.TranslationException;
@@ -36,6 +38,8 @@ public interface ProblemTrace {
 
         Relation getRelation(AnyPartialSymbol partialSymbol);
 
+        RuleDefinition getRuleDefinition(Rule rule);
+
         RuntimeException wrapException(TranslationException translationException);
 
         PartialRelation getPartialRelation(Relation relation);
diff --git a/subprojects/language-semantics/src/main/java/tools/refinery/language/semantics/ProblemTraceImpl.java b/subprojects/language-semantics/src/main/java/tools/refinery/language/semantics/ProblemTraceImpl.java
index 4b44339c..5584af10 100644
--- a/subprojects/language-semantics/src/main/java/tools/refinery/language/semantics/ProblemTraceImpl.java
+++ b/subprojects/language-semantics/src/main/java/tools/refinery/language/semantics/ProblemTraceImpl.java
@@ -13,10 +13,8 @@ import org.eclipse.xtext.naming.IQualifiedNameConverter;
 import org.eclipse.xtext.naming.QualifiedName;
 import org.eclipse.xtext.scoping.IScope;
 import org.eclipse.xtext.scoping.IScopeProvider;
-import tools.refinery.language.model.problem.Node;
-import tools.refinery.language.model.problem.Problem;
-import tools.refinery.language.model.problem.ProblemPackage;
-import tools.refinery.language.model.problem.Relation;
+import tools.refinery.language.model.problem.*;
+import tools.refinery.store.dse.transition.Rule;
 import tools.refinery.store.reasoning.representation.AnyPartialSymbol;
 import tools.refinery.store.reasoning.representation.PartialRelation;
 import tools.refinery.store.reasoning.translator.TranslationException;
@@ -46,6 +44,7 @@ class ProblemTraceImpl implements ProblemTrace {
                         Collections.unmodifiableMap(mutableRelationTrace);
         private final Map<AnyPartialSymbol, Relation> mutableInverseTrace = new HashMap<>();
         private final Map<AnyPartialSymbol, Relation> inverseTrace = Collections.unmodifiableMap(mutableInverseTrace);
+        private final Map<Rule, RuleDefinition> mutableInverseRuleDefinitionTrace = new LinkedHashMap<>();
 
         @Override
         public Problem getProblem() {
@@ -129,6 +128,23 @@ class ProblemTraceImpl implements ProblemTrace {
                 return relation;
         }
 
+        void putRuleDefinition(RuleDefinition ruleDefinition, Rule rule) {
+                var oldRuleDefinition = mutableInverseRuleDefinitionTrace.put(rule, ruleDefinition);
+                if (oldRuleDefinition != null) {
+                        throw new TracedException(oldRuleDefinition, "Rule definition %s was already mapped to rule"
+                                        .formatted(rule.getName()));
+                }
+        }
+
+        @Override
+        public RuleDefinition getRuleDefinition(Rule rule) {
+                var ruleDefinition = mutableInverseRuleDefinitionTrace.get(rule);
+                if (ruleDefinition == null) {
+                        throw new IllegalArgumentException("No definition for rule: " + rule.getName());
+                }
+                return ruleDefinition;
+        }
+
         @Override
         public RuntimeException wrapException(TranslationException translationException) {
                 var partialSymbol = translationException.getPartialSymbol();