Fix double-counting of StructLayout fields

Author: vp-of-awesome

src/org/jruby/ext/ffi/StructLayout.java

@@ -59,10 +59,13 @@ public final class StructLayout extends Type {
     static final String CLASS_NAME = "StructLayout";
 
     /** The name:offset map for this struct */
-    private final Map<IRubyObject, Member> fields;
+    private final Map<IRubyObject, Member> fieldMap;
 
     /** The ordered list of field names (as symbols) */
     private final List<RubySymbol> fieldNames;
+
+    /** The ordered list of field names (as symbols) */
+    private final List<Member> fields;
 
     private final int cacheableFieldCount;
     private final int[] cacheIndexMap;
@@ -95,17 +98,21 @@ public static RubyClass createStructLayoutClass(Ruby runtime, RubyModule module)
      * @param size the total size of the struct.
      * @param minAlign The minimum alignment required when allocating memory.
      */
-    StructLayout(Ruby runtime, Map<IRubyObject, Member> fields, int size, int minAlign) {
+    StructLayout(Ruby runtime, Collection<RubySymbol> fieldNames, Map<IRubyObject, Member> fields, int size, int minAlign) {
         super(runtime, runtime.fastGetModule("FFI").fastGetClass(CLASS_NAME), NativeType.STRUCT, size, minAlign);
         //
         // fields should really be an immutable map as it is never modified after construction
         //
-        this.fields = immutableMap(fields);
-        this.cacheIndexMap = new int[fields.size()];
-        this.referenceIndexMap = new int[fields.size()];
+        this.fieldMap = immutableMap(fields);
+        this.cacheIndexMap = new int[fieldNames.size()];
+        this.referenceIndexMap = new int[fieldNames.size()];
 
         int cfCount = 0, refCount = 0;
-        for (Member m : fields.values()) {
+        List<Member> fieldList = new ArrayList<Member>(fieldNames.size());
+
+        for (RubySymbol fieldName : fieldNames) {
+            Member m = fields.get(fieldName);
+            fieldList.add(m);
             if (m.isCacheable()) {
                 cacheIndexMap[m.index] = cfCount++;
             } else {
@@ -121,13 +128,8 @@ public static RubyClass createStructLayoutClass(Ruby runtime, RubyModule module)
         this.referenceFieldCount = refCount;
 
         // Create the ordered list of field names from the map
-        List<RubySymbol> names = new ArrayList<RubySymbol>(fields.size());
-        for (Map.Entry<IRubyObject, Member> e : fields.entrySet()) {
-            if (e.getKey() instanceof RubySymbol) {
-                names.add((RubySymbol) e.getKey());
-            }
-        }
-        this.fieldNames = Collections.unmodifiableList(names);
+        this.fieldNames = Collections.unmodifiableList(new ArrayList<RubySymbol>(fieldNames));
+        this.fields = Collections.unmodifiableList(fieldList);
     }
 
     /**
@@ -192,15 +194,15 @@ public IRubyObject members(ThreadContext context) {
     public IRubyObject offsets(ThreadContext context) {
         Ruby runtime = context.getRuntime();
         RubyArray offsets = RubyArray.newArray(runtime);
-        for (Map.Entry<IRubyObject, Member> e : fields.entrySet()) {
-            if (e.getKey() instanceof RubySymbol) {
-                RubyArray offset = RubyArray.newArray(runtime);
-                // Assemble a [ :name, offset ] array
-                offset.append(e.getKey());
-                offset.append(runtime.newFixnum(e.getValue().offset));
-                offsets.append(offset);
-            }
+
+        for (RubySymbol name : fieldNames) {
+            RubyArray offset = RubyArray.newArray(runtime);
+            // Assemble a [ :name, offset ] array
+            offset.append(name);
+            offset.append(runtime.newFixnum(fieldMap.get(name).offset));
+            offsets.append(offset);
         }
+
         return offsets;
     }
 
@@ -242,7 +244,7 @@ public IRubyObject offset_of(ThreadContext context, IRubyObject fieldName) {
      * @return A <tt>Member</tt> descriptor.
      */
     final Member getMember(Ruby runtime, IRubyObject name) {
-        Member f = fields.get(name);
+        Member f = fieldMap.get(name);
         if (f != null) {
             return f;
         }
@@ -278,7 +280,7 @@ public final int getFieldCount() {
     }
 
     public final java.util.Collection<Member> getFields() {
-        return Collections.unmodifiableCollection(fields.values());
+        return fields;
     }
 
     /**

src/org/jruby/ext/ffi/StructLayoutBuilder.java

@@ -31,6 +31,8 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.LinkedHashMap;
+import java.util.LinkedList;
+import java.util.List;
 import java.util.Map;
 import org.jruby.Ruby;
 import org.jruby.RubyClass;
@@ -55,6 +57,8 @@
 public final class StructLayoutBuilder extends RubyObject {
     public static final String CLASS_NAME = "StructLayoutBuilder";
 
+    private final List<RubySymbol> fieldNames = new LinkedList<RubySymbol>();
+
     private final Map<IRubyObject, StructLayout.Member> fields = new LinkedHashMap<IRubyObject, StructLayout.Member>();
     /** The current size of the layout in bytes */
     private int size = 0;
@@ -98,7 +102,7 @@ public static StructLayoutBuilder newInstance(ThreadContext context, IRubyObject
 
     @JRubyMethod(name = "build")
     public StructLayout build(ThreadContext context) {
-        return new StructLayout(context.getRuntime(), fields, minAlign + ((size - 1) & ~(minAlign - 1)), minAlign);
+        return new StructLayout(context.getRuntime(), fieldNames, fields, minAlign + ((size - 1) & ~(minAlign - 1)), minAlign);
     }
 
     @JRubyMethod(name = "size")
@@ -119,9 +123,9 @@ private static final int alignMember(int offset, int align) {
         return align + ((offset - 1) & ~(align - 1));
     }
 
-    private static final IRubyObject createSymbolKey(Ruby runtime, IRubyObject key) {
+    private static final RubySymbol createSymbolKey(Ruby runtime, IRubyObject key) {
         if (key instanceof RubySymbol) {
-            return key;
+            return (RubySymbol) key;
         }
         return runtime.getSymbolTable().getSymbol(key.asJavaString());
     }
@@ -131,8 +135,10 @@ private static IRubyObject createStringKey(Ruby runtime, IRubyObject key) {
     }
 
     private final IRubyObject storeField(Ruby runtime, IRubyObject name, StructLayout.Member field, int align, int size) {
+        
         fields.put(createStringKey(runtime, name), field);
         fields.put(createSymbolKey(runtime, name), field);
+        fieldNames.add(createSymbolKey(runtime, name));
         this.size = Math.max(this.size, (int) field.offset + size);
         this.minAlign = Math.max(this.minAlign, align);
         return this;