1 files changed, 47 insertions, 5 deletions
diff --git a/src/main/java/io/trygvis/rules/acme/AcmeIo.java b/src/main/java/io/trygvis/rules/acme/AcmeIo.java
index 67dd4cb..9235992 100644
--- a/src/main/java/io/trygvis/rules/acme/AcmeIo.java
+++ b/src/main/java/io/trygvis/rules/acme/AcmeIo.java
@@ -1,9 +1,11 @@
 package io.trygvis.rules.acme;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.type.TypeFactory;
 import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
 import com.fasterxml.jackson.dataformat.yaml.YAMLGenerator;
 import org.drools.core.common.DefaultFactHandle;
+import org.kie.api.KieBase;
 import org.kie.api.runtime.rule.FactHandle;
 
 import java.io.File;
@@ -12,19 +14,21 @@ import java.io.IOException;
 import java.util.AbstractMap;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
-import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 
 public class AcmeIo {
     private final ObjectMapper mapper;
 
-    public AcmeIo() {
+    public AcmeIo(KieBase kieBase) {
         var factory = new YAMLFactory();
         factory.enable(YAMLGenerator.Feature.USE_NATIVE_TYPE_ID);
         factory.enable(YAMLGenerator.Feature.USE_NATIVE_OBJECT_ID);
         mapper = new ObjectMapper(factory);
+        var typeFactory = TypeFactory.defaultInstance()
+                .withClassLoader(new AcmeClassLoader(kieBase));
+        mapper.setTypeFactory(typeFactory);
         mapper.findAndRegisterModules();
     }
 
@@ -36,7 +40,7 @@ public class AcmeIo {
         List<Object> items = new ArrayList<>(objects.size());
         for (AcmeObject object : objects) {
             try {
-                var type = Class.forName(object.type);
+                var type = mapper.getTypeFactory().findClass(object.type);
                 var x = mapper.treeToValue(object.data, type);
                 items.add(x);
             } catch (ClassNotFoundException e) {
@@ -48,6 +52,10 @@ public class AcmeIo {
     }
 
     public void dump(String s, Collection<FactHandle> factHandles) throws IOException {
+        dump(s, factHandles, (o) -> true);
+    }
+
+    public void dump(String s, Collection<FactHandle> factHandles, Function<Object, Boolean> filter) throws IOException {
         var out = new File("out");
 
         if (!out.isDirectory()) {
@@ -59,9 +67,14 @@ public class AcmeIo {
         List<Map.Entry<String, Object>> facts = new ArrayList<>(factHandles.size());
         for (var handle : factHandles) {
             if (handle instanceof DefaultFactHandle h) {
+                var obj = h.getObject();
+                if (!filter.apply(obj)) {
+                    continue;
+                }
+
                 facts.add(new AbstractMap.SimpleImmutableEntry<>(
                         h.getObjectClassName(),
-                        h.getObject()));
+                        obj));
             }
         }
 
@@ -78,4 +91,33 @@ public class AcmeIo {
             }
         }
     }
+
+    private static class AcmeClassLoader extends ClassLoader {
+        private final KieBase kieBase;
+
+        public AcmeClassLoader(KieBase kieBase) {this.kieBase = kieBase;}
+
+        @Override
+        public Class<?> loadClass(String name) throws ClassNotFoundException {
+            try {
+                return super.loadClass(name);
+            } catch (ClassNotFoundException e) {
+                var i = name.lastIndexOf('.');
+                String pkg, klass;
+                if (i == -1) {
+                    pkg = null;
+                    klass = name;
+                } else {
+                    pkg = name.substring(0, i);
+                    klass = name.substring(i + 1);
+                }
+                var clazz = kieBase.getFactType(pkg, klass);
+                if (clazz == null) {
+                    throw e;
+                }
+
+                return clazz.getFactClass();
+            }
+        }
+    }
 }