package throwing.stream.union;

import java.util.function.Consumer;

import javax.annotation.CheckReturnValue;

import throwing.ExceptionHandler;
import throwing.RethrowChain;
import edu.umd.cs.findbugs.annotations.CleanupObligation;
import edu.umd.cs.findbugs.annotations.CreatesObligation;
import edu.umd.cs.findbugs.annotations.DischargesObligation;

@CleanupObligation
public class UnionThrowable extends Throwable {
  public class ExceptionHandlerBuilder {
    private ExceptionHandler handler = ExceptionHandler.start();

    @CreatesObligation
    public <X> ExceptionHandlerBuilder handle(Class<X> x, Consumer<? super X> action) {
      handler.connect(ExceptionHandler.handle(x, action));
      return this;
    }

    @DischargesObligation
    public void finish() {
      if (!handler.apply(UnionThrowable.this.getCause()).isPresent()) {
        throw UnionThrowable.this.finish();
      }
    }
  }

  private static final long serialVersionUID = -9089451802483526560L;

  public UnionThrowable(Throwable cause) {
    super(cause);
  }

  @CreatesObligation
  public <X extends Throwable> UnionThrowable rethrow(Class<X> x) throws X {
    if (x.isInstance(getCause())) {
      throw x.cast(getCause());
    }

    return this;
  }

  public <X extends Throwable> ExceptionHandlerBuilder handle(Class<X> x, Consumer<? super X> action) {
    return new ExceptionHandlerBuilder().handle(x, action);
  }

  @DischargesObligation
  @CheckReturnValue
  public Error finish() {
    return RethrowChain.END_FUNCTION.apply(getCause());
  }
}

@@ -60,16 +60,52 @@
 onsumer;
+%0Aimport java.util.function.Function;
 %0A%0Aimport
@@ -1176,16 +1176,232 @@
 ows X %7B%0A
+    rethrow(x, Function.identity());%0A%0A    return this;%0A  %7D%0A%0A  @CreatesObligation%0A  public %3CX extends Throwable, Y extends Throwable%3E UnionThrowable rethrow(Class%3CX%3E x,%0A      Function%3C? super X, Y%3E mapper) throws Y %7B%0A
     if (
@@ -1440,16 +1440,29 @@
   throw 
+mapper.apply(
 x.cast(g
@@ -1467,24 +1467,25 @@
 (getCause())
+)
 ;%0A    %7D%0A%0A   

Add a rethrow variant which lets you map the cause