[1.21.2] Adjust how throwables and stack trace are preserved

src/main/java/io/wispforest/owo/mixin/DataResultMixin.java

@@ -1,17 +1,25 @@
 package io.wispforest.owo.mixin;
 
+import com.llamalad7.mixinextras.injector.wrapoperation.Operation;
+import com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation;
 import com.llamalad7.mixinextras.sugar.Local;
 import com.llamalad7.mixinextras.sugar.ref.LocalRef;
+import com.mojang.logging.LogUtils;
 import com.mojang.serialization.DataResult;
+import com.mojang.serialization.Lifecycle;
 import io.wispforest.owo.Owo;
 import io.wispforest.owo.util.StackTraceSupplier;
+import org.slf4j.Logger;
+import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.Unique;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
 
 import java.util.Optional;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
@@ -35,40 +43,85 @@ private static <R> void wrapMessageWithStacktrace(CallbackInfoReturnable<Optiona
         var ogClass = ogSupplier.getClass();
         if (ogSupplier instanceof StackTraceSupplier) return;
 
-        var stackTrace = Thread.currentThread().getStackTrace();
+        StackTraceSupplier stackTraceSupplier = null;
 
         if (ogClass.isSynthetic()) {
             try {
                 for (var field : ogClass.getDeclaredFields()) {
                     if (!Throwable.class.isAssignableFrom(field.getType())) continue;
 
                     field.setAccessible(true);
-                    if (field.get(ogSupplier) instanceof Throwable e) stackTrace = e.getStackTrace().clone();
+                    if (field.get(ogSupplier) instanceof Throwable e) {
+                        stackTraceSupplier = StackTraceSupplier.of(e, ogSupplier);
+                    }
                     break;
                 }
             } catch (IllegalArgumentException | IllegalAccessException ignore) {}
         }
 
-        messageSupplier.set(new StackTraceSupplier(stackTrace, ogSupplier));
+        if (stackTraceSupplier == null) stackTraceSupplier = StackTraceSupplier.of(ogSupplier.get());
+
+        messageSupplier.set(stackTraceSupplier);
     }
 
     @Mixin(value = DataResult.Error.class, remap = false)
     abstract class DataResultErrorMixin<R> {
+
+        @Unique
+        private static final Logger LOGGER = LogUtils.getLogger();
+
         @Shadow(remap = false)
         public abstract Supplier<String> messageSupplier();
 
+        @Shadow @Final
+        private Supplier<String> messageSupplier;
+
         @Inject(method = {"getOrThrow", "getPartialOrThrow"}, at = @At(value = "HEAD"), remap = false)
         private <E extends Throwable> void addStackTraceToException(CallbackInfoReturnable<R> cir, @Local(argsOnly = true) LocalRef<Function<String, E>> exceptionSupplier) {
             final var funcToWrap = exceptionSupplier.get();
 
             exceptionSupplier.set(s -> {
                 var exception = funcToWrap.apply(s);
                 if (this.messageSupplier() instanceof StackTraceSupplier stackTraceSupplier) {
-                    exception.setStackTrace(stackTraceSupplier.stackTrace());
+                    exception.setStackTrace(stackTraceSupplier.getFullStackTrace());
                 }
-
                 return exception;
             });
         }
+
+        @WrapOperation(method ={
+                "resultOrPartial(Ljava/util/function/Consumer;)Ljava/util/Optional;",
+                "promotePartial"
+        }, at = @At(value = "INVOKE", target = "Ljava/util/function/Consumer;accept(Ljava/lang/Object;)V"))
+        private <T> void printStackTrace(Consumer<T> instance, T t, Operation<Void> original) {
+            original.call(instance, t);
+
+            if (Owo.DEBUG && this.messageSupplier instanceof StackTraceSupplier supplier) {
+                LOGGER.error("An error has occurred within DFU: ", supplier.throwable());
+            }
+        }
+
+        @WrapOperation(method = {
+                "ap(Lcom/mojang/serialization/DataResult;)Lcom/mojang/serialization/DataResult$Error;",
+                "flatMap(Ljava/util/function/Function;)Lcom/mojang/serialization/DataResult$Error;"
+        }, at = @At(value = "NEW", target = "(Ljava/util/function/Supplier;Ljava/util/Optional;Lcom/mojang/serialization/Lifecycle;)Lcom/mojang/serialization/DataResult$Error;", ordinal = 1))
+        private DataResult.Error preserveStackTrace1(Supplier<String> messageSupplier, Optional partialValue, Lifecycle lifecycle, Operation<DataResult.Error> original) {
+            if (this.messageSupplier instanceof StackTraceSupplier supplier) {
+                messageSupplier = StackTraceSupplier.of(supplier.throwable(), messageSupplier);
+            }
+
+            return original.call(messageSupplier, partialValue, lifecycle);
+        }
+
+        @WrapOperation(method = {
+                "mapError(Ljava/util/function/UnaryOperator;)Lcom/mojang/serialization/DataResult$Error;"
+        }, at = @At(value = "NEW", target = "(Ljava/util/function/Supplier;Ljava/util/Optional;Lcom/mojang/serialization/Lifecycle;)Lcom/mojang/serialization/DataResult$Error;"))
+        private DataResult.Error preserveStackTrace2(Supplier<String> messageSupplier, Optional partialValue, Lifecycle lifecycle, Operation<DataResult.Error> original) {
+            if (this.messageSupplier instanceof StackTraceSupplier supplier) {
+                messageSupplier = StackTraceSupplier.of(supplier.throwable(), messageSupplier);
+            }
+
+            return original.call(messageSupplier, partialValue, lifecycle);
+        }
     }
 }

src/main/java/io/wispforest/owo/serialization/CodecUtils.java

@@ -29,6 +29,7 @@
 import io.wispforest.owo.serialization.format.nbt.NbtEndec;
 import io.wispforest.owo.serialization.format.nbt.NbtSerializer;
 import io.wispforest.owo.util.Scary;
+import io.wispforest.owo.util.StackTraceSupplier;
 import net.minecraft.nbt.NbtCompound;
 import net.minecraft.nbt.NbtElement;
 import net.minecraft.nbt.NbtOps;
@@ -434,7 +435,7 @@ private static <T> DataResult<T> captureThrows(Supplier<T> action) {
         try {
             return DataResult.success(action.get());
         } catch (Exception e) {
-            return DataResult.error(e::getMessage);
+            return DataResult.error(StackTraceSupplier.of(e));
         }
     }
 

src/main/java/io/wispforest/owo/util/StackTraceSupplier.java

@@ -1,10 +1,51 @@
 package io.wispforest.owo.util;
 
+import org.jetbrains.annotations.Nullable;
+
 import java.util.function.Supplier;
 
-public record StackTraceSupplier(StackTraceElement[] stackTrace, Supplier<String> message) implements Supplier<String> {
+public final class StackTraceSupplier implements Supplier<String> {
+    private final Throwable throwable;
+    private final @Nullable Supplier<String> message;
+
+    private StackTraceSupplier(@Nullable Throwable throwable, @Nullable Supplier<String> message) {
+        this.throwable = throwable;
+        this.message = message;
+    }
+
+    public static StackTraceSupplier of(Throwable throwable) {
+        return new StackTraceSupplier(throwable, null);
+    }
+
+    public static StackTraceSupplier of(Throwable throwable, Supplier<String> supplier) {
+        return new StackTraceSupplier(throwable, supplier);
+    }
+
+    public static StackTraceSupplier of(String message) {
+        return new StackTraceSupplier(new IllegalStateException(message), null);
+    }
+
     @Override
     public String get() {
-        return message.get();
+        return message != null ? message.get() : throwable.getMessage();
+    }
+
+    public StackTraceElement[] getFullStackTrace() {
+        var innerThrowable = throwable();
+        while (innerThrowable.getCause() != null) {
+            innerThrowable = throwable().getCause();
+        }
+        return innerThrowable.getStackTrace();
+    }
+
+    public Throwable throwable() {
+        return throwable;
+    }
+
+    @Override
+    public String toString() {
+        return "StackTraceSupplier[" +
+                "throwable=" + throwable + ", " +
+                "message=" + message + ']';
     }
 }