Additions to ext type support for JRuby

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

@@ -35,28 +35,28 @@ public class Decoder implements Iterator<IRubyObject> {
   private final RubyClass unexpectedTypeErrorClass;
   private final RubyClass unknownExtTypeErrorClass;
 
-  private Unpacker.ExtensionRegistry registry;
+  private ExtensionRegistry registry;
   private ByteBuffer buffer;
   private boolean symbolizeKeys;
   private boolean allowUnknownExt;
 
   public Decoder(Ruby runtime) {
-    this(runtime, null, new byte[] {}, 0, 0);
+    this(runtime, null, new byte[] {}, 0, 0, false, false);
   }
 
-  public Decoder(Ruby runtime, Unpacker.ExtensionRegistry registry) {
-    this(runtime, registry, new byte[] {}, 0, 0);
+  public Decoder(Ruby runtime, ExtensionRegistry registry) {
+    this(runtime, registry, new byte[] {}, 0, 0, false, false);
   }
 
   public Decoder(Ruby runtime, byte[] bytes) {
-    this(runtime, null, bytes, 0, bytes.length);
+    this(runtime, null, bytes, 0, bytes.length, false, false);
   }
 
-  public Decoder(Ruby runtime, Unpacker.ExtensionRegistry registry, byte[] bytes) {
-    this(runtime, registry, bytes, 0, bytes.length);
+  public Decoder(Ruby runtime, ExtensionRegistry registry, byte[] bytes) {
+    this(runtime, registry, bytes, 0, bytes.length, false, false);
   }
 
-  public Decoder(Ruby runtime, Unpacker.ExtensionRegistry registry, byte[] bytes, int offset, int length) {
+  public Decoder(Ruby runtime, ExtensionRegistry registry, byte[] bytes, int offset, int length, boolean symbolizeKeys, boolean allowUnknownExt) {
     this.runtime = runtime;
     this.registry = registry;
     this.binaryEncoding = runtime.getEncodingService().getAscii8bitEncoding();
@@ -67,17 +67,11 @@ public Decoder(Ruby runtime, Unpacker.ExtensionRegistry registry, byte[] bytes,
     this.stackErrorClass = runtime.getModule("MessagePack").getClass("StackError");
     this.unexpectedTypeErrorClass = runtime.getModule("MessagePack").getClass("UnexpectedTypeError");
     this.unknownExtTypeErrorClass = runtime.getModule("MessagePack").getClass("UnknownExtTypeError");
+    this.symbolizeKeys = symbolizeKeys;
+    this.allowUnknownExt = allowUnknownExt;
     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);
   }
@@ -142,7 +136,7 @@ private IRubyObject consumeExtension(int size) {
     byte[] payload = readBytes(size);
 
     if (registry != null) {
-      IRubyObject proc = registry.lookup(type);
+      IRubyObject proc = registry.lookupUnpackerByTypeId(type);
       if (proc != null) {
         ByteList byteList = new ByteList(payload, runtime.getEncodingService().getAscii8bitEncoding());
         return proc.callMethod(runtime.getCurrentContext(), "call", runtime.newString(byteList));

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

@@ -34,19 +34,16 @@ public class Encoder {
   private final Encoding binaryEncoding;
   private final Encoding utf8Encoding;
   private final boolean compatibilityMode;
+  private final ExtensionRegistry registry;
 
-  private Packer.ExtensionRegistry registry;
   private ByteBuffer buffer;
 
-  public Encoder(Ruby runtime, boolean compatibilityMode) {
+  public Encoder(Ruby runtime, boolean compatibilityMode, ExtensionRegistry registry) {
     this.runtime = runtime;
     this.buffer = ByteBuffer.allocate(CACHE_LINE_SIZE - ARRAY_HEADER_SIZE);
     this.binaryEncoding = runtime.getEncodingService().getAscii8bitEncoding();
     this.utf8Encoding = UTF8Encoding.INSTANCE;
     this.compatibilityMode = compatibilityMode;
-  }
-
-  public void setRegistry(Packer.ExtensionRegistry registry) {
     this.registry = registry;
   }
 
@@ -353,7 +350,7 @@ private void appendExtensionValue(ExtensionValue object) {
 
   private void appendOther(IRubyObject object, IRubyObject destination) {
     if (registry != null) {
-      IRubyObject[] pair = registry.lookup(object.getType());
+      IRubyObject[] pair = registry.lookupPackerByClass(object.getType());
       if (pair != null) {
         RubyString bytes = pair[0].callMethod(runtime.getCurrentContext(), "call", object).asString();
         int type = (int) ((RubyFixnum) pair[1]).getLongValue();

ext/java/org/msgpack/jruby/ExtensionRegistry.java

@@ -0,0 +1,158 @@
+package org.msgpack.jruby;
+
+import org.jruby.Ruby;
+import org.jruby.RubyHash;
+import org.jruby.RubyArray;
+import org.jruby.RubyClass;
+import org.jruby.RubyFixnum;
+import org.jruby.runtime.ThreadContext;
+import org.jruby.runtime.builtin.IRubyObject;
+
+import java.util.Map;
+import java.util.HashMap;
+
+public class ExtensionRegistry {
+  private final Map<RubyClass, ExtensionEntry> extensionsByClass;
+  private final Map<RubyClass, ExtensionEntry> extensionsByAncestor;
+  private final ExtensionEntry[] extensionsByTypeId;
+
+  public ExtensionRegistry() {
+    this(new HashMap<RubyClass, ExtensionEntry>());
+  }
+
+  private ExtensionRegistry(Map<RubyClass, ExtensionEntry> extensionsByClass) {
+    this.extensionsByClass = new HashMap<RubyClass, ExtensionEntry>(extensionsByClass);
+    this.extensionsByAncestor = new HashMap<RubyClass, ExtensionEntry>();
+    this.extensionsByTypeId = new ExtensionEntry[256];
+    for (ExtensionEntry entry : extensionsByClass.values()) {
+      if (entry.hasUnpacker()) {
+        extensionsByTypeId[entry.getTypeId() + 128] = entry;
+      }
+    }
+  }
+
+  public ExtensionRegistry dup() {
+    return new ExtensionRegistry(extensionsByClass);
+  }
+
+  public IRubyObject toInternalPackerRegistry(ThreadContext ctx) {
+    RubyHash hash = RubyHash.newHash(ctx.getRuntime());
+    for (RubyClass extensionClass : extensionsByClass.keySet()) {
+      ExtensionEntry entry = extensionsByClass.get(extensionClass);
+      if (entry.hasPacker()) {
+        hash.put(extensionClass, entry.toPackerTuple(ctx));
+      }
+    }
+    return hash;
+  }
+
+  public IRubyObject toInternalUnpackerRegistry(ThreadContext ctx) {
+    RubyHash hash = RubyHash.newHash(ctx.getRuntime());
+    for (ExtensionEntry entry : extensionsByClass.values()) {
+      IRubyObject typeId = RubyFixnum.newFixnum(ctx.getRuntime(), entry.getTypeId());
+      if (entry.hasUnpacker()) {
+        hash.put(typeId, entry.toUnpackerTuple(ctx));
+      }
+    }
+    return hash;
+  }
+
+  public void put(RubyClass cls, int typeId, IRubyObject packerProc, IRubyObject packerArg, IRubyObject unpackerProc, IRubyObject unpackerArg) {
+    ExtensionEntry entry = new ExtensionEntry(cls, typeId, packerProc, packerArg, unpackerProc, unpackerArg);
+    extensionsByClass.put(cls, entry);
+    extensionsByTypeId[typeId + 128] = entry;
+    extensionsByAncestor.clear();
+  }
+
+  public IRubyObject lookupUnpackerByTypeId(int typeId) {
+    ExtensionEntry e = extensionsByTypeId[typeId + 128];
+    if (e != null && e.hasUnpacker()) {
+      return e.getUnpackerProc();
+    } else {
+      return null;
+    }
+  }
+
+  public IRubyObject[] lookupPackerByClass(RubyClass cls) {
+    ExtensionEntry e = extensionsByClass.get(cls);
+    if (e == null) {
+      e = extensionsByAncestor.get(cls);
+    }
+    if (e == null) {
+      e = findEntryByClassOrAncestor(cls);
+      if (e != null) {
+        extensionsByAncestor.put(e.getExtensionClass(), e);
+      }
+    }
+    if (e != null && e.hasPacker()) {
+      return e.toPackerProcTypeIdPair(cls.getRuntime().getCurrentContext());
+    } else {
+      return null;
+    }
+  }
+
+  private ExtensionEntry findEntryByClassOrAncestor(final RubyClass cls) {
+    ThreadContext ctx = cls.getRuntime().getCurrentContext();
+    for (RubyClass extensionClass : extensionsByClass.keySet()) {
+      RubyArray ancestors = (RubyArray) cls.callMethod(ctx, "ancestors");
+      if (ancestors.callMethod(ctx, "include?", extensionClass).isTrue()) {
+        return extensionsByClass.get(extensionClass);
+      }
+    }
+    return null;
+  }
+
+  private static class ExtensionEntry {
+    private final RubyClass cls;
+    private final int typeId;
+    private final IRubyObject packerProc;
+    private final IRubyObject packerArg;
+    private final IRubyObject unpackerProc;
+    private final IRubyObject unpackerArg;
+
+    public ExtensionEntry(RubyClass cls, int typeId, IRubyObject packerProc, IRubyObject packerArg, IRubyObject unpackerProc, IRubyObject unpackerArg) {
+      this.cls = cls;
+      this.typeId = typeId;
+      this.packerProc = packerProc;
+      this.packerArg = packerArg;
+      this.unpackerProc = unpackerProc;
+      this.unpackerArg = unpackerArg;
+    }
+
+    public RubyClass getExtensionClass() {
+      return cls;
+    }
+
+    public int getTypeId() {
+      return typeId;
+    }
+
+    public boolean hasPacker() {
+      return packerProc != null;
+    }
+
+    public boolean hasUnpacker() {
+      return unpackerProc != null;
+    }
+
+    public IRubyObject getPackerProc() {
+      return packerProc;
+    }
+
+    public IRubyObject getUnpackerProc() {
+      return unpackerProc;
+    }
+
+    public RubyArray toPackerTuple(ThreadContext ctx) {
+      return RubyArray.newArray(ctx.getRuntime(), new IRubyObject[] {cls, packerProc, packerArg});
+    }
+
+    public RubyArray toUnpackerTuple(ThreadContext ctx) {
+      return RubyArray.newArray(ctx.getRuntime(), new IRubyObject[] {RubyFixnum.newFixnum(ctx.getRuntime(), typeId), unpackerProc, unpackerArg});
+    }
+
+    public IRubyObject[] toPackerProcTypeIdPair(ThreadContext ctx) {
+      return new IRubyObject[] {packerProc, RubyFixnum.newFixnum(ctx.getRuntime(), typeId)};
+    }
+  }
+}

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

@@ -22,13 +22,13 @@
 
 @JRubyClass(name="MessagePack::Factory")
 public class Factory extends RubyObject {
-  private Ruby runtime;
-  private Packer.ExtensionRegistry packerExtensionRegistry;
-  private Unpacker.ExtensionRegistry unpackerExtensionRegistry;
+  private final Ruby runtime;
+  private final ExtensionRegistry extensionRegistry;
 
   public Factory(Ruby runtime, RubyClass type) {
     super(runtime, type);
     this.runtime = runtime;
+    this.extensionRegistry = new ExtensionRegistry();
   }
 
   static class FactoryAllocator implements ObjectAllocator {
@@ -37,43 +37,31 @@ public IRubyObject allocate(Ruby runtime, RubyClass type) {
     }
   }
 
+  public ExtensionRegistry extensionRegistry() {
+    return extensionRegistry.dup();
+  }
+
   @JRubyMethod(name = "initialize")
   public IRubyObject initialize(ThreadContext ctx) {
-    this.packerExtensionRegistry = new Packer.ExtensionRegistry(ctx.getRuntime());
-    this.unpackerExtensionRegistry = new Unpacker.ExtensionRegistry(ctx.getRuntime());
     return this;
   }
 
-  public Packer.ExtensionRegistry packerRegistry() {
-    return this.packerExtensionRegistry.dup();
-  }
-
-  public Unpacker.ExtensionRegistry unpackerRegistry() {
-    return this.unpackerExtensionRegistry.dup();
-  }
-
   @JRubyMethod(name = "packer", optional = 1)
   public Packer packer(ThreadContext ctx, IRubyObject[] args) {
-    return Packer.newPacker(ctx, packerRegistry(), args);
+    return Packer.newPacker(ctx, extensionRegistry(), args);
   }
 
   @JRubyMethod(name = "unpacker", optional = 1)
   public Unpacker unpacker(ThreadContext ctx, IRubyObject[] args) {
-    return Unpacker.newUnpacker(ctx, unpackerRegistry(), args);
+    return Unpacker.newUnpacker(ctx, extensionRegistry(), args);
   }
 
   @JRubyMethod(name = "registered_types_internal", visibility = PRIVATE)
   public IRubyObject registeredTypesInternal(ThreadContext ctx) {
-    Unpacker.ExtensionRegistry 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);
+    return RubyArray.newArray(ctx.getRuntime(), new IRubyObject[] {
+      extensionRegistry.toInternalPackerRegistry(ctx),
+      extensionRegistry.toInternalUnpackerRegistry(ctx)
+    });
   }
 
   @JRubyMethod(name = "register_type", required = 2, optional = 1)
@@ -122,8 +110,7 @@ public IRubyObject registerType(ThreadContext ctx, IRubyObject[] args) {
       }
     }
 
-    this.packerExtensionRegistry.put(extClass, (int) typeId, packerProc, packerArg);
-    this.unpackerExtensionRegistry.put(extClass, (int) typeId, unpackerProc, unpackerArg);
+    extensionRegistry.put(extClass, (int) typeId, packerProc, packerArg, unpackerProc, unpackerArg);
 
     return runtime.getNil();
   }

ext/java/org/msgpack/jruby/MessagePackLibrary.java

@@ -117,13 +117,15 @@ public static IRubyObject pack(ThreadContext ctx, IRubyObject recv, IRubyObject[
 
     @JRubyMethod(module = true, required = 1, optional = 1, alias = {"load"})
     public static IRubyObject unpack(ThreadContext ctx, IRubyObject recv, IRubyObject[] args) {
-      Unpacker.ExtensionRegistry registry = MessagePackLibrary.defaultFactory.unpackerRegistry();
-      Decoder decoder = new Decoder(ctx.getRuntime(), registry, args[0].asString().getBytes());
+      ExtensionRegistry registry = MessagePackLibrary.defaultFactory.extensionRegistry();
+      boolean symbolizeKeys = false;
       if (args.length > 1 && !args[args.length - 1].isNil()) {
         RubyHash hash = args[args.length - 1].convertToHash();
-        IRubyObject symbolizeKeys = hash.fastARef(ctx.getRuntime().newSymbol("symbolize_keys"));
-        decoder.symbolizeKeys(symbolizeKeys != null && symbolizeKeys.isTrue());
+        IRubyObject symbolizeKeysVal = hash.fastARef(ctx.getRuntime().newSymbol("symbolize_keys"));
+        symbolizeKeys = symbolizeKeysVal != null && symbolizeKeysVal.isTrue();
       }
+      byte[] bytes = args[0].asString().getBytes();
+      Decoder decoder = new Decoder(ctx.getRuntime(), registry, bytes, 0, bytes.length, symbolizeKeys, false);
       return decoder.next();
     }
   }