Add ToBuilder to incrementally add additional types to a marshaller (#15

)

src/main/java/org/curioswitch/common/protobuf/json/MarshallerRegistry.java

@@ -62,6 +62,11 @@ TypeSpecificMarshaller<?> findByTypeUrl(String typeUrl) throws InvalidProtocolBu
     return marshaller;
   }
 
+  /** Returns the built parsers in this registry. */
+  Map<Descriptor, TypeSpecificMarshaller<?>> getBuiltParsers() {
+    return descriptorRegistry;
+  }
+
   @SuppressWarnings("StringSplitter")
   private static String getTypeName(String typeUrl) throws InvalidProtocolBufferException {
     String[] parts = typeUrl.split("/");

src/main/java/org/curioswitch/common/protobuf/json/MessageMarshaller.java

@@ -27,6 +27,7 @@
 import java.io.OutputStream;
 import java.lang.reflect.InvocationTargetException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -124,9 +125,29 @@ public SerializableString getEscapeSequence(int ch) {
 
   private final MarshallerRegistry registry;
 
-  private MessageMarshaller(@Nullable PrettyPrinter prettyPrinter, MarshallerRegistry registry) {
-    this.prettyPrinter = prettyPrinter;
+  private final boolean includingDefaultValueFields;
+  private final boolean preservingProtoFieldNames;
+  private final boolean omittingInsignificantWhitespace;
+  private final boolean ignoringUnknownFields;
+  private final boolean printingEnumsAsInts;
+  private final boolean sortingMapKeys;
+
+  private MessageMarshaller(
+      MarshallerRegistry registry,
+      boolean omittingInsignificantWhitespace,
+      boolean includingDefaultValueFields,
+      boolean preservingProtoFieldNames,
+      boolean ignoringUnknownFields,
+      boolean printingEnumsAsInts,
+      boolean sortingMapKeys) {
+    this.prettyPrinter = omittingInsignificantWhitespace ? null : new MessagePrettyPrinter();
     this.registry = registry;
+    this.omittingInsignificantWhitespace = omittingInsignificantWhitespace;
+    this.includingDefaultValueFields = includingDefaultValueFields;
+    this.preservingProtoFieldNames = preservingProtoFieldNames;
+    this.ignoringUnknownFields = ignoringUnknownFields;
+    this.printingEnumsAsInts = printingEnumsAsInts;
+    this.sortingMapKeys = sortingMapKeys;
   }
 
   /**
@@ -256,6 +277,23 @@ public <T extends Message> void writeValue(T message, JsonGenerator gen) throws
     }
   }
 
+  /**
+   * Returns a new {@link Builder} prepopulated with the messages that have been registered in this
+   * {@link MessageMarshaller}. This can be useful to incrementally add more messages to an already
+   * built marshaller. Be careful when doing this on a hot path as compilation can take a relatively
+   * large amount of time.
+   */
+  public Builder toBuilder() {
+    return new Builder(
+        registry.getBuiltParsers(),
+        omittingInsignificantWhitespace,
+        includingDefaultValueFields,
+        preservingProtoFieldNames,
+        ignoringUnknownFields,
+        printingEnumsAsInts,
+        sortingMapKeys);
+  }
+
   /**
    * A {@link Builder} of {@link MessageMarshaller}s, allows registering {@link Message} types to
    * marshall and set options.
@@ -271,6 +309,7 @@ public static final class Builder {
 
     private final List<Message> prototypes = new ArrayList<>();
 
+    private final Map<Descriptor, TypeSpecificMarshaller<?>> preBuiltParsers;
     /**
      * Registers the type of the provided {@link Message} for use with the created {@link
      * MessageMarshaller}. While any instance of the type to register can be used, this will
@@ -383,23 +422,30 @@ public Builder sortingMapKeys(boolean sortingMapKeys) {
      * will not be usable with the returned {@link MessageMarshaller}.
      */
     public MessageMarshaller build() {
-      Map<Descriptor, TypeSpecificMarshaller<?>> builtParsers = new HashMap<>();
-      addStandardParser(BoolValueMarshaller.INSTANCE, builtParsers);
-      addStandardParser(Int32ValueMarshaller.INSTANCE, builtParsers);
-      addStandardParser(UInt32ValueMarshaller.INSTANCE, builtParsers);
-      addStandardParser(Int64ValueMarshaller.INSTANCE, builtParsers);
-      addStandardParser(UInt64ValueMarshaller.INSTANCE, builtParsers);
-      addStandardParser(StringValueMarshaller.INSTANCE, builtParsers);
-      addStandardParser(BytesValueMarshaller.INSTANCE, builtParsers);
-      addStandardParser(FloatValueMarshaller.INSTANCE, builtParsers);
-      addStandardParser(DoubleValueMarshaller.INSTANCE, builtParsers);
-      addStandardParser(TimestampMarshaller.INSTANCE, builtParsers);
-      addStandardParser(DurationMarshaller.INSTANCE, builtParsers);
-      addStandardParser(FieldMaskMarshaller.INSTANCE, builtParsers);
-      addStandardParser(StructMarshaller.INSTANCE, builtParsers);
-      addStandardParser(ValueMarshaller.INSTANCE, builtParsers);
-      addStandardParser(ListValueMarshaller.INSTANCE, builtParsers);
+      Map<Descriptor, TypeSpecificMarshaller<?>> builtParsers;
+      if (!preBuiltParsers.isEmpty()) {
+        builtParsers = new HashMap<>(preBuiltParsers);
+      } else {
+        builtParsers = new HashMap<>();
+        addStandardParser(BoolValueMarshaller.INSTANCE, builtParsers);
+        addStandardParser(Int32ValueMarshaller.INSTANCE, builtParsers);
+        addStandardParser(UInt32ValueMarshaller.INSTANCE, builtParsers);
+        addStandardParser(Int64ValueMarshaller.INSTANCE, builtParsers);
+        addStandardParser(UInt64ValueMarshaller.INSTANCE, builtParsers);
+        addStandardParser(StringValueMarshaller.INSTANCE, builtParsers);
+        addStandardParser(BytesValueMarshaller.INSTANCE, builtParsers);
+        addStandardParser(FloatValueMarshaller.INSTANCE, builtParsers);
+        addStandardParser(DoubleValueMarshaller.INSTANCE, builtParsers);
+        addStandardParser(TimestampMarshaller.INSTANCE, builtParsers);
+        addStandardParser(DurationMarshaller.INSTANCE, builtParsers);
+        addStandardParser(FieldMaskMarshaller.INSTANCE, builtParsers);
+        addStandardParser(StructMarshaller.INSTANCE, builtParsers);
+        addStandardParser(ValueMarshaller.INSTANCE, builtParsers);
+        addStandardParser(ListValueMarshaller.INSTANCE, builtParsers);
+      }
 
+      // AnyMarshaller must be re-registered even if preBuiltParsers are used
+      // to be able to reference the newly registered ones.
       AnyMarshaller anyParser = new AnyMarshaller();
       addStandardParser(anyParser, builtParsers);
 
@@ -418,7 +464,13 @@ public MessageMarshaller build() {
       anyParser.setMarshallerRegistry(registry);
 
       return new MessageMarshaller(
-          omittingInsignificantWhitespace ? null : new MessagePrettyPrinter(), registry);
+          registry,
+          omittingInsignificantWhitespace,
+          includingDefaultValueFields,
+          preservingProtoFieldNames,
+          ignoringUnknownFields,
+          printingEnumsAsInts,
+          sortingMapKeys);
     }
 
     private static <T extends Message> void addStandardParser(
@@ -427,7 +479,26 @@ private static <T extends Message> void addStandardParser(
       marshallers.put(marshaller.getDescriptorForMarshalledType(), marshaller);
     }
 
-    private Builder() {}
+    private Builder() {
+      preBuiltParsers = Collections.emptyMap();
+    }
+
+    Builder(
+        Map<Descriptor, TypeSpecificMarshaller<?>> preBuiltParsers,
+        boolean omittingInsignificantWhitespace,
+        boolean includingDefaultValueFields,
+        boolean preservingProtoFieldNames,
+        boolean ignoringUnknownFields,
+        boolean printingEnumsAsInts,
+        boolean sortingMapKeys) {
+      this.preBuiltParsers = preBuiltParsers;
+      this.omittingInsignificantWhitespace = omittingInsignificantWhitespace;
+      this.includingDefaultValueFields = includingDefaultValueFields;
+      this.preservingProtoFieldNames = preservingProtoFieldNames;
+      this.ignoringUnknownFields = ignoringUnknownFields;
+      this.printingEnumsAsInts = printingEnumsAsInts;
+      this.sortingMapKeys = sortingMapKeys;
+    }
   }
 
   private static class MessagePrettyPrinter extends DefaultPrettyPrinter {

src/test/java/org/curioswitch/common/protobuf/json/MessageMarshallerIncrementalTest.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright (c) Choko (choko@curioswitch.org)
+ * SPDX-License-Identifier: MIT
+ */
+
+package org.curioswitch.common.protobuf.json;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import com.google.protobuf.util.JsonTestProto;
+import org.curioswitch.common.protobuf.json.test.GithubApi;
+import org.junit.jupiter.api.Test;
+
+class MessageMarshallerIncrementalTest {
+
+  @Test
+  void incrementalMarshaller() throws Exception {
+    MessageMarshaller marshaller =
+        MessageMarshaller.builder().register(GithubApi.SearchResponse.getDefaultInstance()).build();
+    GithubApi.SearchResponse response =
+        GithubApi.SearchResponse.newBuilder().setTotalCount(10).build();
+    JsonTestProto.TestAllTypes allTypes = JsonTestUtil.testAllTypesAllFields();
+    assertThat(marshaller.writeValueAsString(response)).isNotEmpty();
+    assertThatThrownBy(() -> marshaller.writeValueAsString(allTypes))
+        .isInstanceOf(IllegalArgumentException.class);
+
+    MessageMarshaller marshaller2 =
+        marshaller.toBuilder().register(JsonTestProto.TestAllTypes.getDefaultInstance()).build();
+    assertThat(marshaller2.writeValueAsString(response)).isNotEmpty();
+    assertThat(marshaller2.writeValueAsString(allTypes)).isNotEmpty();
+  }
+}