add full spec ext type support to JRuby implementation

.gitignore

@@ -11,6 +11,7 @@ Gemfile*
 pkg
 test/debug.log
 *~
+*.swp
 /rdoc
 tmp
 .classpath

ext/java/org/msgpack/jruby/Decoder.java

@@ -12,6 +12,7 @@
 import org.jruby.RubyClass;
 import org.jruby.RubyBignum;
 import org.jruby.RubyString;
+import org.jruby.RubyArray;
 import org.jruby.RubyHash;
 import org.jruby.exceptions.RaiseException;
 import org.jruby.runtime.builtin.IRubyObject;
@@ -29,33 +30,54 @@ public class Decoder implements Iterator<IRubyObject> {
   private final Encoding utf8Encoding;
   private final RubyClass unpackErrorClass;
   private final RubyClass underflowErrorClass;
+  private final RubyClass malformedFormatErrorClass;
+  private final RubyClass stackErrorClass;
   private final RubyClass unexpectedTypeErrorClass;
+  private final RubyClass unknownExtTypeErrorClass;
 
+  private Unpacker.ExtRegistry registry;
   private ByteBuffer buffer;
   private boolean symbolizeKeys;
+  private boolean allowUnknownExt;
 
   public Decoder(Ruby runtime) {
-    this(runtime, new byte[] {}, 0, 0);
+    this(runtime, null, new byte[] {}, 0, 0);
+  }
+
+  public Decoder(Ruby runtime, Unpacker.ExtRegistry registry) {
+    this(runtime, registry, new byte[] {}, 0, 0);
   }
 
   public Decoder(Ruby runtime, byte[] bytes) {
-    this(runtime, bytes, 0, bytes.length);
+    this(runtime, null, bytes, 0, bytes.length);
+  }
+
+  public Decoder(Ruby runtime, Unpacker.ExtRegistry registry, byte[] bytes) {
+    this(runtime, registry, bytes, 0, bytes.length);
   }
 
-  public Decoder(Ruby runtime, byte[] bytes, int offset, int length) {
+  public Decoder(Ruby runtime, Unpacker.ExtRegistry registry, byte[] bytes, int offset, int length) {
     this.runtime = runtime;
+    this.registry = registry;
     this.binaryEncoding = runtime.getEncodingService().getAscii8bitEncoding();
     this.utf8Encoding = UTF8Encoding.INSTANCE;
     this.unpackErrorClass = runtime.getModule("MessagePack").getClass("UnpackError");
     this.underflowErrorClass = runtime.getModule("MessagePack").getClass("UnderflowError");
+    this.malformedFormatErrorClass = runtime.getModule("MessagePack").getClass("MalformedFormatError");
+    this.stackErrorClass = runtime.getModule("MessagePack").getClass("StackError");
     this.unexpectedTypeErrorClass = runtime.getModule("MessagePack").getClass("UnexpectedTypeError");
+    this.unknownExtTypeErrorClass = runtime.getModule("MessagePack").getClass("UnknownExtTypeError");
     feed(bytes, offset, length);
   }
 
   public void symbolizeKeys(boolean symbolize) {
     this.symbolizeKeys = symbolize;
   }
 
+  public void allowUnknownExt(boolean allow) {
+    this.allowUnknownExt = allow;
+  }
+
   public void feed(byte[] bytes) {
     feed(bytes, 0, bytes.length);
   }
@@ -73,7 +95,7 @@ public void feed(byte[] bytes, int offset, int length) {
   }
 
   public void reset() {
-    buffer.rewind();
+    buffer = null;
   }
 
   public int offset() {
@@ -118,7 +140,20 @@ private IRubyObject consumeHash(int size) {
   private IRubyObject consumeExtension(int size) {
     int type = buffer.get();
     byte[] payload = readBytes(size);
-    return ExtensionValue.newExtensionValue(runtime, type, payload);
+
+    if (this.registry != null) {
+      IRubyObject proc = this.registry.lookup(type);
+      if (proc != null) {
+        ByteList byteList = new ByteList(payload, runtime.getEncodingService().getAscii8bitEncoding());
+        return proc.callMethod(runtime.getCurrentContext(), "call", runtime.newString(byteList));
+      }
+    }
+
+    if (this.allowUnknownExt) {
+      return ExtensionValue.newExtensionValue(runtime, type, payload);
+    }
+
+    throw runtime.newRaiseException(unexpectedTypeErrorClass, "unexpected type");
   }
 
   private byte[] readBytes(int size) {
@@ -142,11 +177,11 @@ public IRubyObject read_array_header() {
     try {
       byte b = buffer.get();
       if ((b & 0xf0) == 0x90) {
-	return runtime.newFixnum(b & 0x0f);
+        return runtime.newFixnum(b & 0x0f);
       } else if (b == ARY16) {
-	return runtime.newFixnum(buffer.getShort() & 0xffff);
+        return runtime.newFixnum(buffer.getShort() & 0xffff);
       } else if (b == ARY32) {
-	return runtime.newFixnum(buffer.getInt());
+        return runtime.newFixnum(buffer.getInt());
       }
       throw runtime.newRaiseException(unexpectedTypeErrorClass, "unexpected type");
     } catch (RaiseException re) {
@@ -163,11 +198,11 @@ public IRubyObject read_map_header() {
     try {
       byte b = buffer.get();
       if ((b & 0xf0) == 0x80) {
-	return runtime.newFixnum(b & 0x0f);
+        return runtime.newFixnum(b & 0x0f);
       } else if (b == MAP16) {
-	return runtime.newFixnum(buffer.getShort() & 0xffff);
+        return runtime.newFixnum(buffer.getShort() & 0xffff);
       } else if (b == MAP32) {
-	return runtime.newFixnum(buffer.getInt());
+        return runtime.newFixnum(buffer.getInt());
       }
       throw runtime.newRaiseException(unexpectedTypeErrorClass, "unexpected type");
     } catch (RaiseException re) {
@@ -233,7 +268,7 @@ public IRubyObject next() {
       default: return runtime.newFixnum(b);
       }
       buffer.position(position);
-      throw runtime.newRaiseException(unpackErrorClass, "Illegal byte sequence");
+      throw runtime.newRaiseException(malformedFormatErrorClass, "Illegal byte sequence");
     } catch (RaiseException re) {
       buffer.position(position);
       throw re;

ext/java/org/msgpack/jruby/Encoder.java

@@ -35,6 +35,7 @@ public class Encoder {
   private final Encoding utf8Encoding;
   private final boolean compatibilityMode;
 
+  private Packer.ExtRegistry registry;
   private ByteBuffer buffer;
 
   public Encoder(Ruby runtime, boolean compatibilityMode) {
@@ -45,6 +46,10 @@ public Encoder(Ruby runtime, boolean compatibilityMode) {
     this.compatibilityMode = compatibilityMode;
   }
 
+  public void setRegistry(Packer.ExtRegistry registry) {
+    this.registry = registry;
+  }
+
   private void ensureRemainingCapacity(int c) {
     if (buffer.remaining() < c) {
       int newLength = Math.max(buffer.capacity() + (buffer.capacity() >> 1), buffer.capacity() + c);
@@ -107,7 +112,7 @@ private void appendObject(IRubyObject object, IRubyObject destination) {
     } else if (object instanceof ExtensionValue) {
       appendExtensionValue((ExtensionValue) object);
     } else {
-      appendCustom(object, destination);
+      appendOther(object, destination);
     }
   }
 
@@ -295,12 +300,7 @@ public void visit(IRubyObject key, IRubyObject value) {
     }
   }
 
-  private void appendExtensionValue(ExtensionValue object) {
-    long type = ((RubyFixnum)object.get_type()).getLongValue();
-    if (type < -128 || type > 127) {
-	    throw object.getRuntime().newRangeError(String.format("integer %d too big to convert to `signed char'", type));
-    }
-    ByteList payloadBytes = ((RubyString)object.payload()).getByteList();
+  private void appendExt(int type, ByteList payloadBytes) {
     int payloadSize = payloadBytes.length();
     int outputSize = 0;
     boolean fixSize = payloadSize == 1 || payloadSize == 2 || payloadSize == 4 || payloadSize == 8 || payloadSize == 16;
@@ -338,6 +338,29 @@ private void appendExtensionValue(ExtensionValue object) {
     buffer.put(payloadBytes.unsafeBytes(), payloadBytes.begin(), payloadSize);
   }
 
+  private void appendExtensionValue(ExtensionValue object) {
+    long type = ((RubyFixnum)object.get_type()).getLongValue();
+    if (type < -128 || type > 127) {
+	    throw object.getRuntime().newRangeError(String.format("integer %d too big to convert to `signed char'", type));
+    }
+    ByteList payloadBytes = ((RubyString)object.payload()).getByteList();
+    appendExt((int) type, payloadBytes);
+  }
+
+  private void appendOther(IRubyObject object, IRubyObject destination) {
+    if (this.registry != null) {
+      IRubyObject[] pair = this.registry.lookup(object.getType());
+      if (pair != null) {
+        RubyString bytes = pair[0].callMethod(runtime.getCurrentContext(), "call", object).asString();
+        int type = (int) ((RubyFixnum) pair[1]).getLongValue();
+        appendExt(type, bytes.getByteList());
+        return;
+      }
+    }
+    // registry is null or type is not registered
+    appendCustom(object, destination);
+  }
+
   private void appendCustom(IRubyObject object, IRubyObject destination) {
     if (destination == null) {
       IRubyObject result = object.callMethod(runtime.getCurrentContext(), "to_msgpack");

ext/java/org/msgpack/jruby/Factory.java

@@ -0,0 +1,133 @@
+package org.msgpack.jruby;
+
+
+import org.jruby.Ruby;
+import org.jruby.RubyClass;
+import org.jruby.RubyObject;
+import org.jruby.RubyArray;
+import org.jruby.RubyHash;
+import org.jruby.RubyIO;
+import org.jruby.RubyInteger;
+import org.jruby.RubyFixnum;
+import org.jruby.RubyString;
+import org.jruby.RubySymbol;
+import org.jruby.runtime.builtin.IRubyObject;
+import org.jruby.anno.JRubyClass;
+import org.jruby.anno.JRubyMethod;
+import org.jruby.runtime.ThreadContext;
+import org.jruby.runtime.ObjectAllocator;
+import org.jruby.util.ByteList;
+
+import static org.jruby.runtime.Visibility.PRIVATE;
+
+@JRubyClass(name="MessagePack::Factory")
+public class Factory extends RubyObject {
+  private Ruby runtime;
+  private Packer.ExtRegistry packerExtRegistry;
+  private Unpacker.ExtRegistry unpackerExtRegistry;
+
+  public Factory(Ruby runtime, RubyClass type) {
+    super(runtime, type);
+    this.runtime = runtime;
+  }
+
+  static class FactoryAllocator implements ObjectAllocator {
+    public IRubyObject allocate(Ruby runtime, RubyClass type) {
+      return new Factory(runtime, type);
+    }
+  }
+
+  @JRubyMethod(name = "initialize")
+  public IRubyObject initialize(ThreadContext ctx) {
+    this.packerExtRegistry = new Packer.ExtRegistry(ctx.getRuntime());
+    this.unpackerExtRegistry = new Unpacker.ExtRegistry(ctx.getRuntime());
+    return this;
+  }
+
+  public Packer.ExtRegistry packerRegistry() {
+    return this.packerExtRegistry.dup();
+  }
+
+  public Unpacker.ExtRegistry unpackerRegistry() {
+    return this.unpackerExtRegistry.dup();
+  }
+
+  @JRubyMethod(name = "packer", optional = 1)
+  public Packer packer(ThreadContext ctx, IRubyObject[] args) {
+    return Packer.newPacker(ctx, packerRegistry(), args);
+  }
+
+  @JRubyMethod(name = "unpacker", optional = 1)
+  public Unpacker unpacker(ThreadContext ctx, IRubyObject[] args) {
+    return Unpacker.newUnpacker(ctx, unpackerRegistry(), args);
+  }
+
+  @JRubyMethod(name = "registered_types_internal", visibility = PRIVATE)
+  public IRubyObject registeredTypesInternal(ThreadContext ctx) {
+    // unpacker
+    Unpacker.ExtRegistry reg = unpackerRegistry();
+    RubyHash mapping = RubyHash.newHash(ctx.getRuntime());
+    for (int i = 0; i < 256; i++) {
+      if (reg.array[i] != null) {
+        mapping.fastASet(RubyFixnum.newFixnum(ctx.getRuntime(), i - 128), reg.array[i]);
+      }
+    }
+
+    IRubyObject[] ary = { packerRegistry().hash, mapping };
+    return RubyArray.newArray(ctx.getRuntime(), ary);
+  }
+
+  @JRubyMethod(name = "register_type", required = 2, optional = 1)
+  public IRubyObject registerType(ThreadContext ctx, IRubyObject[] args) {
+    // register_type(type, Class)
+    // register_type(type, Class, packer: proc-like, unpacker: proc-like)
+    Ruby runtime = ctx.getRuntime();
+    IRubyObject type = args[0];
+    IRubyObject klass = args[1];
+
+    IRubyObject packerArg;
+    IRubyObject unpackerArg;
+    if (args.length == 2) {
+      packerArg = runtime.newSymbol("to_msgpack_ext");
+      unpackerArg = runtime.newSymbol("from_msgpack_ext");
+    } else if (args.length == 3) {
+      if (args[args.length - 1] instanceof RubyHash) {
+        RubyHash options = (RubyHash) args[args.length - 1];
+        packerArg = options.fastARef(runtime.newSymbol("packer"));
+        unpackerArg = options.fastARef(runtime.newSymbol("unpacker"));
+      } else {
+        throw runtime.newArgumentError(String.format("expected Hash but found %s.", args[args.length - 1].getType().getName()));
+      }
+    } else {
+      throw runtime.newArgumentError(String.format("wrong number of arguments (%d for 2..3)", 2 + args.length));
+    }
+
+    long typeId = ((RubyFixnum) type).getLongValue();
+    if (typeId < -128 || typeId > 127) {
+      throw runtime.newRangeError(String.format("integer %d too big to convert to `signed char'", typeId));
+    }
+
+    if (!(klass instanceof RubyClass)) {
+      throw runtime.newArgumentError(String.format("expected Class but found %s.", klass.getType().getName()));
+    }
+    RubyClass extClass = (RubyClass) klass;
+
+    IRubyObject packerProc = runtime.getNil();
+    IRubyObject unpackerProc = runtime.getNil();
+    if (packerArg != runtime.getNil()) {
+      packerProc = packerArg.callMethod(ctx, "to_proc");
+    }
+    if (unpackerArg != runtime.getNil()) {
+      if (unpackerArg instanceof RubyString || unpackerArg instanceof RubySymbol) {
+        unpackerProc = extClass.method(unpackerArg.callMethod(ctx, "to_sym"));
+      } else {
+        unpackerProc = unpackerArg.callMethod(ctx, "method", runtime.newSymbol("call"));
+      }
+    }
+
+    this.packerExtRegistry.put(extClass, (int) typeId, packerProc, packerArg);
+    this.unpackerExtRegistry.put(extClass, (int) typeId, unpackerProc, unpackerArg);
+
+    return runtime.getNil();
+  }
+}