Serialize/deserialize classes using records as carrier objets

Marshaller.java

package fr.umlv.marshaller;

import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles.Lookup;
import java.lang.reflect.UndeclaredThrowableException;
import java.util.List;
import java.util.Map;

import static java.lang.invoke.MethodType.methodType;
import static java.util.Objects.requireNonNull;
import static java.util.function.Function.identity;
import static java.util.stream.Collectors.toMap;

public interface Marshaller {
  Record deconstruct(Object o);
  <T> T reconstruct(Record record, Class<T> type);

  default void writeObject(ObjectOutputStream out, Object o) throws IOException {
    var type = o.getClass();
    out.writeObject(type);
    out.writeObject(deconstruct(o));
  }

  default Object readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
    var type = (Class<?>) in.readObject();
    var record = (Record) in.readObject();
    return reconstruct(record, type);
  }

  @Target(ElementType.TYPE)
  @Retention(RetentionPolicy.RUNTIME)
  @interface Marshall {
    Class<? extends Record> deconstruct();
    Class<? extends Record>[] reconstructs() default {};
  }

  static Marshaller of(Lookup lookup) {
    requireNonNull(lookup);
    record Info(MethodHandle deconstructor, Map<Class<? extends Record>, MethodHandle> reconstructors) {
      private static MethodHandle deconstructor(Lookup lookup, Class<?> type, Class<?> schema) {
        try {
          return lookup.findVirtual(type, "deconstructor", methodType(schema))
              .asType(methodType(Record.class, Object.class));
        } catch (NoSuchMethodException | IllegalAccessException e) {
          throw new IllegalStateException(e);
        }
      }

      private static MethodHandle reconstructor(Lookup lookup, Class<?> type, Class<?> schema) {
        try {
          return lookup.findStatic(type, "reconstructor", methodType(type, schema))
              .asType(methodType(Object.class, Record.class));
        } catch (NoSuchMethodException | IllegalAccessException e) {
          throw new IllegalStateException(e);
        }
      }

      static Info create(Lookup lookup, Class<?> type, Class<? extends Record> deconstructorSchema, List<Class<? extends Record>> reconstructorSchemas) {
        return new Info(
            deconstructor(lookup, type, deconstructorSchema),
            reconstructorSchemas.stream().collect(toMap(identity(), s -> reconstructor(lookup, type, s)))
        );
      }

      static Info create(Lookup lookup, Class<?> type) {
        var marshall = type.getAnnotation(Marshall.class);
        if (marshall == null) {
          throw new IllegalStateException("no annotation @Marshall declared on " + type);
        }
        var deconstructSchema = marshall.deconstruct();
        var reconstructs = marshall.reconstructs();
        var reconstructSchemas = reconstructs.length == 0?
            List.<Class<? extends Record>>of(deconstructSchema):
            List.of(reconstructs);
        return create(lookup, type, deconstructSchema, reconstructSchemas);
      }
    }

    record Impl(ClassValue<Info> cache) implements Marshaller {
      @Override
      public Record deconstruct(Object o) {
        var deconstructor = cache.get(o.getClass()).deconstructor;
        try {
          return (Record) deconstructor.invokeExact(o);
        } catch (RuntimeException | Error e) {
          throw e;
        } catch(Throwable t) {
          throw new UndeclaredThrowableException(t);
        }
      }

      @Override
      public <T> T reconstruct(Record record, Class<T> type) {
        var reconstructor = cache.get(type).reconstructors.get(record.getClass());
        if (reconstructor == null) {
          throw new IllegalStateException("no reconstructor found for " + type.getName() + " from " + record.getClass());
        }
        try {
          return type.cast(reconstructor.invokeExact(record));
        } catch (RuntimeException | Error e) {
          throw e;
        } catch(Throwable t) {
          throw new UndeclaredThrowableException(t);
        }
      }
    }
    return new Impl(new ClassValue<>() {
      @Override
      protected Info computeValue(Class<?> type) {
        return Info.create(lookup, type);
      }
    });
  }
}

MarshallerTest.java

package fr.umlv.marshaller;

import fr.umlv.marshaller.Marshaller.Marshall;
import org.junit.jupiter.api.Test;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.Serializable;

import static java.lang.invoke.MethodHandles.lookup;
import static org.junit.jupiter.api.Assertions.*;

public class MarshallerTest {
  record Point(int x, int y) implements Serializable { }

  @Marshall(deconstruct = Point.class, reconstructs = Point.class)
  static class MutablePoint {
    int x;
    int y;

    public MutablePoint(int x, int y) {
      this.x = x;
      this.y = y;
    }

    private Point deconstructor() {
      return new Point(x, y);
    }
    private static MutablePoint reconstructor(Point point) {
      return new MutablePoint(point.x, point.y);
    }
  }

  @Test
  public void mutablePoint() {
    var marshaller = Marshaller.of(lookup());

    var mutablePoint = new MutablePoint(1, 2);
    var record = marshaller.deconstruct(mutablePoint);
    assertEquals(new Point(1, 2), record);

    var mutablePoint2 = marshaller.reconstruct(record, MutablePoint.class);
    assertEquals(1, mutablePoint2.x);
    assertEquals(2, mutablePoint2.y);
  }

  @Test
  public void mutablePointSerialization() throws IOException, ClassNotFoundException {
    var marshaller = Marshaller.of(lookup());

    var mutablePoint = new MutablePoint(1, 2);
    var byteArrayOutputStream = new ByteArrayOutputStream();
    marshaller.writeObject(new ObjectOutputStream(byteArrayOutputStream), mutablePoint);

    var byteArrayInputStream = new ByteArrayInputStream(byteArrayOutputStream.toByteArray());
    var mutablePoint2 = (MutablePoint) marshaller.readObject(new ObjectInputStream(byteArrayInputStream));

    assertEquals(1, mutablePoint2.x);
    assertEquals(2, mutablePoint2.y);
  }
}