Merge pull request #364 from komamitsu/complex_type_key

Support complex type key in MessagePackGenerator

Author: komamitsu

msgpack-jackson/src/main/java/org/msgpack/jackson/dataformat/MessagePackFactory.java

@@ -37,6 +37,7 @@ public class MessagePackFactory
     private static final long serialVersionUID = 2578263992015504347L;
 
     private final MessagePack.PackerConfig packerConfig;
+    private boolean reuseResourceInGenerator = true;
 
     public MessagePackFactory()
     {
@@ -48,11 +49,16 @@ public MessagePackFactory(MessagePack.PackerConfig packerConfig)
         this.packerConfig = packerConfig;
     }
 
+    public void setReuseResourceInGenerator(boolean reuseResourceInGenerator)
+    {
+        this.reuseResourceInGenerator = reuseResourceInGenerator;
+    }
+
     @Override
     public JsonGenerator createGenerator(OutputStream out, JsonEncoding enc)
             throws IOException
     {
-        return new MessagePackGenerator(_generatorFeatures, _objectCodec, out, packerConfig);
+        return new MessagePackGenerator(_generatorFeatures, _objectCodec, out, packerConfig, reuseResourceInGenerator);
     }
 
     @Override

msgpack-jackson/src/main/java/org/msgpack/jackson/dataformat/MessagePackGenerator.java

@@ -26,6 +26,7 @@
 import org.msgpack.core.MessagePacker;
 import org.msgpack.core.buffer.OutputStreamBufferOutput;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.math.BigDecimal;
@@ -42,6 +43,8 @@ public class MessagePackGenerator
     private static final Charset DEFAULT_CHARSET = Charset.forName("UTF-8");
     private final MessagePacker messagePacker;
     private static ThreadLocal<OutputStreamBufferOutput> messageBufferOutputHolder = new ThreadLocal<OutputStreamBufferOutput>();
+    private final OutputStream output;
+    private final MessagePack.PackerConfig packerConfig;
     private LinkedList<StackItem> stack;
     private StackItem rootStackItem;
 
@@ -97,20 +100,35 @@ List<Object> getKeys()
         }
     }
 
-    public MessagePackGenerator(int features, ObjectCodec codec, OutputStream out, MessagePack.PackerConfig packerConfig)
+    public MessagePackGenerator(
+            int features,
+            ObjectCodec codec,
+            OutputStream out,
+            MessagePack.PackerConfig packerConfig,
+            boolean reuseResourceInGenerator)
             throws IOException
     {
         super(features, codec);
-        OutputStreamBufferOutput messageBufferOutput = messageBufferOutputHolder.get();
-        if (messageBufferOutput == null) {
-            messageBufferOutput = new OutputStreamBufferOutput(out);
+        this.output = out;
+
+        OutputStreamBufferOutput messageBufferOutput;
+        if (reuseResourceInGenerator) {
+            messageBufferOutput = messageBufferOutputHolder.get();
+            if (messageBufferOutput == null) {
+                messageBufferOutput = new OutputStreamBufferOutput(out);
+                messageBufferOutputHolder.set(messageBufferOutput);
+            }
+            else {
+                messageBufferOutput.reset(out);
+            }
         }
         else {
-            messageBufferOutput.reset(out);
+            messageBufferOutput = new OutputStreamBufferOutput(out);
         }
-        messageBufferOutputHolder.set(messageBufferOutput);
-
         this.messagePacker = packerConfig.newPacker(messageBufferOutput);
+
+        this.packerConfig = packerConfig;
+
         this.stack = new LinkedList<StackItem>();
     }
 
@@ -224,7 +242,11 @@ else if (v instanceof MessagePackExtensionType) {
             messagePacker.writePayload(extData);
         }
         else {
-            throw new IllegalArgumentException(v.toString());
+            messagePacker.flush();
+            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+            MessagePackGenerator messagePackGenerator = new MessagePackGenerator(getFeatureMask(), getCodec(), outputStream, packerConfig, false);
+            getCodec().writeValue(messagePackGenerator, v);
+            output.write(outputStream.toByteArray());
         }
     }
 

msgpack-jackson/src/test/java/org/msgpack/jackson/dataformat/MessagePackGeneratorTest.java

@@ -690,4 +690,46 @@ else if (keyName.equals("bigIntMap")) {
             }
         }
     }
+
+    @Test
+    public void testComplexTypeKey()
+            throws IOException
+    {
+        HashMap<TinyPojo, Integer> map = new HashMap<TinyPojo, Integer>();
+        TinyPojo pojo = new TinyPojo();
+        pojo.t = "foo";
+        map.put(pojo, 42);
+
+        ObjectMapper objectMapper = new ObjectMapper(new MessagePackFactory());
+        objectMapper.setSerializerFactory(new MessagePackSerializerFactory());
+        byte[] bytes = objectMapper.writeValueAsBytes(map);
+
+        MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(bytes);
+        assertThat(unpacker.unpackMapHeader(), is(1));
+        assertThat(unpacker.unpackMapHeader(), is(1));
+        assertThat(unpacker.unpackString(), is("t"));
+        assertThat(unpacker.unpackString(), is("foo"));
+        assertThat(unpacker.unpackInt(), is(42));
+    }
+
+    @Test
+    public void testComplexTypeKeyWithV06Format()
+            throws IOException
+    {
+        HashMap<TinyPojo, Integer> map = new HashMap<TinyPojo, Integer>();
+        TinyPojo pojo = new TinyPojo();
+        pojo.t = "foo";
+        map.put(pojo, 42);
+
+        ObjectMapper objectMapper = new ObjectMapper(new MessagePackFactory());
+        objectMapper.setAnnotationIntrospector(new JsonArrayFormat());
+        objectMapper.setSerializerFactory(new MessagePackSerializerFactory());
+        byte[] bytes = objectMapper.writeValueAsBytes(map);
+
+        MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(bytes);
+        assertThat(unpacker.unpackMapHeader(), is(1));
+        assertThat(unpacker.unpackArrayHeader(), is(1));
+        assertThat(unpacker.unpackString(), is("foo"));
+        assertThat(unpacker.unpackInt(), is(42));
+    }
 }