HHH-8354 - New dirty-checking options based on bytecode enhancement

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/spi/CollectionTracker.java

@@ -0,0 +1,51 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.bytecode.enhance.spi;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * @author <a href="mailto:stale.pedersen@jboss.org">Ståle W. Pedersen</a>
+ */
+public class CollectionTracker {
+
+    private Map<String, Integer> tracker;
+
+    public CollectionTracker() {
+        tracker = new HashMap<String, Integer>();
+    }
+
+    public void add(String name, int size) {
+        tracker.put(name, size);
+    }
+
+    public int getSize(String name) {
+        Integer size = tracker.get(name);
+        if(size == null)
+            return -1;
+        else
+            return size;
+    }
+}

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/spi/CompositeOwnerTracker.java

@@ -0,0 +1,85 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.bytecode.enhance.spi;
+
+import org.hibernate.engine.spi.CompositeOwner;
+
+/**
+ * small low memory class to keep references to composite owners
+ *
+ * @author <a href="mailto:stale.pedersen@jboss.org">Ståle W. Pedersen</a>
+ */
+public class CompositeOwnerTracker {
+
+    private String[] names;
+    private CompositeOwner[] owners;
+    private int size = 0;
+
+    public CompositeOwnerTracker() {
+        names = new String[1];
+        owners = new CompositeOwner[1];
+    }
+
+    public void add(String name, CompositeOwner owner) {
+        for(int i=0; i<size;i++) {
+            if(names[i].equals(name)) {
+                owners[i] = owner;
+                return;
+            }
+        }
+        if ( size >= names.length) {
+            String[] tmpNames = new String[size+1];
+            System.arraycopy(names, 0, tmpNames, 0, size);
+            names = tmpNames;
+            CompositeOwner[] tmpOwners = new CompositeOwner[size+1];
+            System.arraycopy(owners, 0, tmpOwners, 0, size);
+            owners = tmpOwners;
+        }
+        names[size] = name;
+        owners[size] = owner;
+        size++;
+    }
+
+    public void callOwner(String fieldName) {
+        for(int i=0; i < size;i++) {
+            owners[i].$$_hibernate_trackChange(names[i]+fieldName);
+        }
+    }
+
+    public void removeOwner(String name) {
+        for(int i=0; i<size;i++) {
+            if(names[i].equals(name)) {
+                if(i < size) {
+                    for(int j=i; j < size-1;j++) {
+                        names[j] = names[j+1];
+                        owners[j] = owners[j+1];
+                    }
+                    names[size-1] = null;
+                    owners[size-1] = null;
+                    size--;
+                }
+            }
+        }
+    }
+}

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/spi/EnhancementContext.java

@@ -121,4 +121,11 @@ public interface EnhancementContext {
 	 * @return {@code true} if the field is lazy loadable; {@code false} otherwise.
 	 */
 	public boolean isLazyLoadable(CtField field);
+
+    /**
+     *
+     * @param field the field to check
+     * @return {@code true} if the field is mapped
+     */
+    public boolean isMappedCollection(CtField field);
 }

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/spi/Enhancer.java

@@ -23,14 +23,21 @@
  */
 package org.hibernate.bytecode.enhance.spi;
 
+import javax.persistence.ElementCollection;
+import javax.persistence.Embedded;
+import javax.persistence.Id;
+import javax.persistence.ManyToMany;
+import javax.persistence.OneToMany;
 import javax.persistence.Transient;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.DataOutputStream;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.IdentityHashMap;
 import java.util.List;
+import java.util.Map;
 
 import javassist.CannotCompileException;
 import javassist.ClassPool;
@@ -49,10 +56,12 @@ import javassist.bytecode.ConstPool;
 import javassist.bytecode.FieldInfo;
 import javassist.bytecode.MethodInfo;
 import javassist.bytecode.Opcode;
+import javassist.bytecode.SignatureAttribute;
 import javassist.bytecode.StackMapTable;
 import javassist.bytecode.annotation.Annotation;
 import javassist.bytecode.stackmap.MapMaker;
 
+import org.hibernate.engine.spi.SelfDirtinessTracker;
 import org.jboss.logging.Logger;
 
 import org.hibernate.HibernateException;
@@ -82,6 +91,7 @@ public class Enhancer  {
 	private final CtClass attributeInterceptableCtClass;
 	private final CtClass entityEntryCtClass;
 	private final CtClass objectCtClass;
+    private boolean isComposite;
 
     /**
 	 * Constructs the Enhancer, using the given context.
@@ -163,12 +173,31 @@ public class Enhancer  {
 			managedCtClass = classPool.makeClassIfNew( new ByteArrayInputStream( originalBytes ) );
 		}
 		catch (IOException e) {
-			log.unableToBuildEnhancementMetamodel( className );
+			log.unableToBuildEnhancementMetamodel(className);
 			return originalBytes;
 		}
 
-		enhance( managedCtClass );
+		enhance(managedCtClass, false);
 
+        return getByteCode(managedCtClass);
+	}
+
+    public byte[] enhanceComposite(String className, byte[] originalBytes) throws EnhancementException {
+		final CtClass managedCtClass;
+		try {
+			managedCtClass = classPool.makeClassIfNew( new ByteArrayInputStream( originalBytes ) );
+		}
+		catch (IOException e) {
+			log.unableToBuildEnhancementMetamodel(className);
+			return originalBytes;
+		}
+
+		enhance(managedCtClass, true);
+
+        return getByteCode(managedCtClass);
+    }
+
+    private byte[] getByteCode(CtClass managedCtClass) {
 		final ByteArrayOutputStream byteStream = new ByteArrayOutputStream();
 		final DataOutputStream out;
 		try {
@@ -192,7 +221,8 @@ public class Enhancer  {
 		}
     }
 
-	private void enhance(CtClass managedCtClass) {
+	private void enhance(CtClass managedCtClass, boolean isComposite) {
+        this.isComposite = isComposite;
 		final String className = managedCtClass.getName();
 		log.debugf( "Enhancing %s", className );
 
@@ -212,10 +242,10 @@ public class Enhancer  {
 			}
 		}
 
-		if ( enhancementContext.isEntityClass( managedCtClass ) ) {
+		if (!isComposite && enhancementContext.isEntityClass( managedCtClass ) ) {
 			enhanceAsEntity( managedCtClass );
 		}
-		else if ( enhancementContext.isCompositeClass( managedCtClass ) ) {
+		else if (isComposite || enhancementContext.isCompositeClass( managedCtClass ) ) {
 			enhanceAsComposite( managedCtClass );
 		}
 		else {
@@ -303,7 +333,7 @@ public class Enhancer  {
 		return annotationsAttribute;
 	}
 
-	private void enhancePersistentAttributes(CtClass managedCtClass) {
+	private void enhancePersistentAttributes(CtClass managedCtClass ) {
 		addInterceptorHandling( managedCtClass );
 		if ( enhancementContext.doDirtyCheckingInline( managedCtClass ) ) {
 			addInLineDirtyHandling( managedCtClass );
@@ -366,6 +396,34 @@ public class Enhancer  {
 		return enhancementContext.order( persistentFieldList.toArray( new CtField[persistentFieldList.size()]) );
 	}
 
+    private List<CtField> collectCollectionFields(CtClass managedCtClass) {
+
+        final List<CtField> collectionList = new ArrayList<CtField>();
+        try {
+            for ( CtField ctField : managedCtClass.getDeclaredFields() ) {
+                // skip static fields
+                if ( Modifier.isStatic( ctField.getModifiers() ) ) {
+                    continue;
+                }
+                // skip fields added by enhancement
+                if ( ctField.getName().startsWith( "$" ) ) {
+                    continue;
+                }
+                if ( enhancementContext.isPersistentField( ctField ) ) {
+                    for(CtClass ctClass : ctField.getType().getInterfaces()) {
+                        if(ctClass.getName().equals("java.util.Collection")) {
+                            collectionList.add(ctField);
+                            break;
+                        }
+                    }
+                }
+            }
+        }
+        catch (NotFoundException ignored) {  }
+
+        return collectionList;
+    }
+
 	private void addInterceptorHandling(CtClass managedCtClass) {
 		// interceptor handling is only needed if either:
 		//		a) in-line dirty checking has *not* been requested
@@ -390,8 +448,265 @@ public class Enhancer  {
 		);
 	}
 
-	private void addInLineDirtyHandling(CtClass managedCtClass) {
+    private boolean isClassAlreadyTrackingDirtyStatus(CtClass managedCtClass) {
-		// todo : implement
+        try {
+            for(CtClass ctInterface : managedCtClass.getInterfaces()) {
+                if(ctInterface.getName().equals(SelfDirtinessTracker.class.getName()))
+                    return true;
+            }
+        }
+        catch (NotFoundException e) {
+            e.printStackTrace();
+        }
+        return false;
+    }
+
+	private void addInLineDirtyHandling(CtClass managedCtClass ) {
+        try {
+
+            //create composite methods
+            if(isComposite) {
+                managedCtClass.addInterface(classPool.get("org.hibernate.engine.spi.CompositeTracker"));
+                CtClass compositeCtType = classPool.get("org.hibernate.bytecode.enhance.spi.CompositeOwnerTracker");
+                addField(managedCtClass, compositeCtType, EnhancerConstants.TRACKER_COMPOSITE_FIELD_NAME, true);
+                createCompositeTrackerMethod(managedCtClass);
+            }
+            // "normal" entity
+            else {
+                managedCtClass.addInterface(classPool.get("org.hibernate.engine.spi.SelfDirtinessTracker"));
+                CtClass trackerCtType = classPool.get("java.util.Set");
+                addField(managedCtClass, trackerCtType, EnhancerConstants.TRACKER_FIELD_NAME, true);
+
+                CtClass collectionTrackerCtType = classPool.get("org.hibernate.bytecode.enhance.spi.CollectionTracker");
+                addField(managedCtClass, collectionTrackerCtType, EnhancerConstants.TRACKER_COLLECTION_NAME, true);
+
+                createDirtyTrackerMethods(managedCtClass);
+            }
+
+
+        }
+        catch (NotFoundException e) {
+            e.printStackTrace();
+        }
+    }
+
+    /**
+     * Create all dirty tracker methods
+     */
+    private void createDirtyTrackerMethods(CtClass managedCtClass) {
+        try {
+            String trackerChangeMethod =
+                    "public void "+EnhancerConstants.TRACKER_CHANGER_NAME+"(String name) {" +
+                            "  if(" +EnhancerConstants.TRACKER_FIELD_NAME+ " == null) {" +
+                            "    "+EnhancerConstants.TRACKER_FIELD_NAME+ " = new java.util.HashSet();" +
+                            "  }" +
+                            "  if(!" +EnhancerConstants.TRACKER_FIELD_NAME+ ".contains(name)) {" +
+                            "    "+EnhancerConstants.TRACKER_FIELD_NAME+".add(name);" +
+                            "  }"+
+                            "}";
+            managedCtClass.addMethod(CtNewMethod.make(trackerChangeMethod, managedCtClass));
+
+            createCollectionDirtyCheckMethod(managedCtClass);
+            createCollectionDirtyCheckGetFieldsMethod(managedCtClass);
+            //createCompositeFieldsDirtyCheckMethod(managedCtClass);
+            //createGetCompositeDirtyFieldsMethod(managedCtClass);
+
+            createHasDirtyAttributesMethod(managedCtClass);
+
+            createClearDirtyCollectionMethod(managedCtClass);
+            createClearDirtyMethod(managedCtClass);
+
+           String trackerGetMethod =
+                    "public java.util.List "+EnhancerConstants.TRACKER_GET_NAME+"() { "+
+                            "if("+ EnhancerConstants.TRACKER_FIELD_NAME+" == null) "+
+                            EnhancerConstants.TRACKER_FIELD_NAME+" = new java.util.HashSet();"+
+                            EnhancerConstants.TRACKER_COLLECTION_CHANGED_FIELD_NAME+"("+
+                            EnhancerConstants.TRACKER_FIELD_NAME+");"+
+                            "return "+EnhancerConstants.TRACKER_FIELD_NAME+"; }";
+            CtMethod getMethod =  CtNewMethod.make(trackerGetMethod, managedCtClass);
+
+            MethodInfo methodInfo = getMethod.getMethodInfo();
+            SignatureAttribute signatureAttribute =
+                    new SignatureAttribute(methodInfo.getConstPool(), "()Ljava/util/Set<Ljava/lang/String;>;");
+            methodInfo.addAttribute(signatureAttribute);
+            managedCtClass.addMethod(getMethod);
+
+        }
+        catch (CannotCompileException e) {
+            e.printStackTrace();
+        }
+        catch (ClassNotFoundException e) {
+            e.printStackTrace();
+        }
+    }
+
+    private void createTrackChangeCompositeMethod(CtClass managedCtClass) {
+        StringBuilder builder = new StringBuilder();
+        builder.append("public void ")
+                .append(EnhancerConstants.TRACKER_CHANGER_NAME)
+                .append("(String name) {")
+                .append("if (")
+                .append(EnhancerConstants.TRACKER_COMPOSITE_FIELD_NAME)
+                .append(" != null) ")
+                .append(EnhancerConstants.TRACKER_COMPOSITE_FIELD_NAME)
+                .append(".callOwner(\".\"+name); }");
+
+        System.out.println("COMPOSITE METHOD: "+builder.toString());
+
+        try {
+            managedCtClass.addMethod(CtNewMethod.make(builder.toString(), managedCtClass));
+        }
+        catch (CannotCompileException e) {
+        }
+    }
+
+    private void createCompositeTrackerMethod(CtClass managedCtClass) {
+        try {
+            StringBuilder builder = new StringBuilder();
+            builder.append("public void ")
+                    .append(EnhancerConstants.TRACKER_COMPOSITE_SET_OWNER)
+                    .append("(String name, org.hibernate.engine.spi.CompositeOwner tracker) {")
+                    .append("if(")
+                    .append(EnhancerConstants.TRACKER_COMPOSITE_FIELD_NAME)
+                    .append(" == null) ")
+                    .append(EnhancerConstants.TRACKER_COMPOSITE_FIELD_NAME)
+                    .append(" = new org.hibernate.bytecode.enhance.spi.CompositeOwnerTracker();")
+                    .append(EnhancerConstants.TRACKER_COMPOSITE_FIELD_NAME)
+                    .append(".add(name, tracker); }");
+
+            managedCtClass.addMethod(CtNewMethod.make(builder.toString(), managedCtClass));
+
+            builder = new StringBuilder();
+            builder.append("public void ")
+                    .append(EnhancerConstants.TRACKER_COMPOSITE_CLEAR_OWNER)
+                    .append("(String name) {")
+                    .append(" if(")
+                    .append(EnhancerConstants.TRACKER_COMPOSITE_FIELD_NAME)
+                    .append(" != null)")
+                    .append(EnhancerConstants.TRACKER_COMPOSITE_FIELD_NAME)
+                    .append(".removeOwner(name);}");
+
+            managedCtClass.addMethod(CtNewMethod.make(builder.toString(), managedCtClass));
+        }
+        catch (CannotCompileException e) {
+            e.printStackTrace();
+        }
+    }
+
+    private void createHasDirtyAttributesMethod(CtClass managedCtClass) throws CannotCompileException {
+        String trackerHasChangedMethod =
+                "public boolean "+EnhancerConstants.TRACKER_HAS_CHANGED_NAME+"() { return ("+
+                        EnhancerConstants.TRACKER_FIELD_NAME+" != null && !" +
+                        EnhancerConstants.TRACKER_FIELD_NAME+".isEmpty()) || "+
+                        EnhancerConstants.TRACKER_COLLECTION_CHANGED_NAME+"(); } ";
+
+        managedCtClass.addMethod(CtNewMethod.make(trackerHasChangedMethod, managedCtClass));
+    }
+
+    /**
+     * Creates _clearDirtyAttributes
+     */
+    private void createClearDirtyMethod(CtClass managedCtClass) throws CannotCompileException, ClassNotFoundException {
+        StringBuilder builder = new StringBuilder();
+        builder.append("public void ")
+                .append( EnhancerConstants.TRACKER_CLEAR_NAME)
+                .append("() {")
+                .append("if (")
+                .append(EnhancerConstants.TRACKER_FIELD_NAME)
+                .append(" != null) ")
+                .append(EnhancerConstants.TRACKER_FIELD_NAME)
+                .append(".clear(); ")
+                .append(EnhancerConstants.TRACKER_COLLECTION_CLEAR_NAME)
+                .append("(); }");
+
+        managedCtClass.addMethod(CtNewMethod.make(builder.toString(), managedCtClass));
+    }
+
+    private void createClearDirtyCollectionMethod(CtClass managedCtClass) throws CannotCompileException {
+        StringBuilder builder = new StringBuilder();
+        builder.append("private void ")
+                .append(EnhancerConstants.TRACKER_COLLECTION_CLEAR_NAME)
+                .append("() { if(")
+                .append(EnhancerConstants.TRACKER_COLLECTION_NAME)
+                .append(" == null)")
+                .append(EnhancerConstants.TRACKER_COLLECTION_NAME)
+                .append(" = new org.hibernate.bytecode.enhance.spi.CollectionTracker();");
+
+        for(CtField ctField : collectCollectionFields(managedCtClass)) {
+            if(!enhancementContext.isMappedCollection(ctField)) {
+                builder.append("if(")
+                        .append(ctField.getName())
+                        .append(" != null) ")
+                        .append(EnhancerConstants.TRACKER_COLLECTION_NAME)
+                        .append(".add(\"")
+                        .append(ctField.getName())
+                        .append("\", ")
+                        .append(ctField.getName())
+                        .append(".size());");
+            }
+        }
+
+        builder.append("}");
+
+        managedCtClass.addMethod(CtNewMethod.make(builder.toString(), managedCtClass));
+    }
+
+    /**
+     * create _areCollectionFieldsDirty
+     */
+    private void createCollectionDirtyCheckMethod(CtClass managedCtClass) throws CannotCompileException {
+        StringBuilder builder = new StringBuilder("private boolean ")
+                .append(EnhancerConstants.TRACKER_COLLECTION_CHANGED_NAME)
+                .append("() { if ($$_hibernate_getInterceptor() == null || ")
+                .append(EnhancerConstants.TRACKER_COLLECTION_NAME)
+                .append(" == null) return false; ");
+
+        for(CtField ctField : collectCollectionFields(managedCtClass)) {
+            if(!enhancementContext.isMappedCollection(ctField)) {
+                builder.append("if(")
+                        .append(EnhancerConstants.TRACKER_COLLECTION_NAME)
+                        .append(".getSize(\"")
+                        .append(ctField.getName())
+                        .append("\") != ")
+                        .append(ctField.getName())
+                        .append(".size()) return true;");
+            }
+        }
+
+        builder.append("return false; }");
+
+        managedCtClass.addMethod(CtNewMethod.make(builder.toString(), managedCtClass));
+    }
+
+    /**
+     * create _getCollectionFieldDirtyNames
+     */
+    private void createCollectionDirtyCheckGetFieldsMethod(CtClass managedCtClass) throws CannotCompileException {
+              StringBuilder collectionFieldDirtyFieldMethod = new StringBuilder("private void ")
+                      .append(EnhancerConstants.TRACKER_COLLECTION_CHANGED_FIELD_NAME)
+                      .append("(java.util.Set trackerSet) { if(")
+                      .append(EnhancerConstants.TRACKER_COLLECTION_NAME)
+                      .append(" == null) return; else {");
+
+        for(CtField ctField : collectCollectionFields(managedCtClass)) {
+            if(!ctField.getName().startsWith("$$_hibernate") &&
+                    !enhancementContext.isMappedCollection(ctField)) {
+                collectionFieldDirtyFieldMethod
+                        .append("if(")
+                        .append(EnhancerConstants.TRACKER_COLLECTION_NAME)
+                        .append(".getSize(\"")
+                        .append(ctField.getName())
+                        .append("\") != ")
+                        .append(ctField.getName())
+                        .append(".size()) trackerSet.add(\"")
+                        .append(ctField.getName())
+                        .append("\");");
+            }
+        }
+
+        collectionFieldDirtyFieldMethod.append("}}");
+
+        managedCtClass.addMethod(CtNewMethod.make(collectionFieldDirtyFieldMethod.toString(), managedCtClass));
     }
 
 	private void addFieldWithGetterAndSetter(
@@ -557,8 +872,37 @@ public class Enhancer  {
 				);
 			}
 
-			if ( enhancementContext.doDirtyCheckingInline( managedCtClass ) ) {
+			if ( enhancementContext.doDirtyCheckingInline( managedCtClass ) && !isComposite ) {
-				writer.insertBefore( typeDescriptor.buildInLineDirtyCheckingBodyFragment( fieldName ) );
+                writer.insertBefore( typeDescriptor.buildInLineDirtyCheckingBodyFragment( persistentField  ));
+			}
+
+            if( isComposite) {
+                StringBuilder builder = new StringBuilder();
+                builder.append(" if(  ")
+                        .append(EnhancerConstants.TRACKER_COMPOSITE_FIELD_NAME)
+                        .append(" != null) ")
+                        .append(EnhancerConstants.TRACKER_COMPOSITE_FIELD_NAME)
+                        .append(".callOwner(\".")
+                        .append(persistentField.getName())
+                        .append("\");");
+
+                writer.insertBefore( builder.toString() );
+            }
+
+            //composite types
+            if(persistentField.getAnnotation(Embedded.class) != null) {
+                //make sure to add the CompositeOwner interface
+                if(!doClassInheritCompositeOwner(managedCtClass)) {
+                    managedCtClass.addInterface(classPool.get("org.hibernate.engine.spi.CompositeOwner"));
+                }
+                //if a composite have a embedded field we need to implement the method as well
+                if(isComposite)
+                    createTrackChangeCompositeMethod(managedCtClass);
+
+
+                writer.insertBefore( cleanupPreviousOwner(persistentField));
+
+                writer.insertAfter( compositeMethodBody(persistentField));
             }
 
 			managedCtClass.addMethod( writer );
@@ -576,6 +920,47 @@ public class Enhancer  {
 		}
 	}
 
+    private boolean doClassInheritCompositeOwner(CtClass managedCtClass) {
+        try {
+            for(CtClass ctClass : managedCtClass.getInterfaces())
+                if(ctClass.getName().equals("org.hibernate.engine.spi.CompositeOwner"))
+                    return true;
+
+            return false;
+        }
+        catch (NotFoundException e) {
+            return false;
+        }
+    }
+
+    private String cleanupPreviousOwner(CtField currentValue) {
+        StringBuilder builder = new StringBuilder();
+        builder.append("if (")
+                .append(currentValue.getName())
+                .append(" != null) ")
+                .append("((org.hibernate.engine.spi.CompositeTracker)")
+                .append(currentValue.getName())
+                .append(").")
+                .append(EnhancerConstants.TRACKER_COMPOSITE_CLEAR_OWNER)
+                .append("(\"")
+                .append(currentValue.getName())
+                .append("\");");
+
+        return  builder.toString();
+    }
+
+    private String compositeMethodBody(CtField currentValue) {
+        StringBuilder builder = new StringBuilder();
+        builder.append("((org.hibernate.engine.spi.CompositeTracker) ")
+                .append(currentValue.getName())
+                .append(").$$_hibernate_setOwner(\"")
+                .append(currentValue.getName())
+                .append("\",(org.hibernate.engine.spi.CompositeOwner) this);")
+                .append(EnhancerConstants.TRACKER_CHANGER_NAME + "(\"" + currentValue.getName() + "\");");
+
+        return builder.toString();
+    }
+
 	private void transformFieldAccessesIntoReadsAndWrites(
 			CtClass managedCtClass,
 			IdentityHashMap<String, PersistentAttributeDescriptor> attributeDescriptorMap) {
@@ -699,7 +1084,7 @@ public class Enhancer  {
 	private static interface AttributeTypeDescriptor {
 		public String buildReadInterceptionBodyFragment(String fieldName);
 		public String buildWriteInterceptionBodyFragment(String fieldName);
-		public String buildInLineDirtyCheckingBodyFragment(String fieldName);
+		public String buildInLineDirtyCheckingBodyFragment(CtField currentField);
 	}
 
 	private AttributeTypeDescriptor resolveAttributeTypeDescriptor(CtField persistentField) throws NotFoundException {
@@ -735,13 +1120,73 @@ public class Enhancer  {
 
 	private abstract static class AbstractAttributeTypeDescriptor implements AttributeTypeDescriptor {
 		@Override
-		public String buildInLineDirtyCheckingBodyFragment(String fieldName) {
+		public String buildInLineDirtyCheckingBodyFragment(CtField currentValue) {
-			// for now...
+            StringBuilder builder = new StringBuilder();
-			// todo : hook-in in-lined dirty checking
+            try {
-			return String.format(
+                //should ignore primary keys
-					"System.out.println( \"DIRTY CHECK (%1$s) : \" + this.%1$s + \" -> \" + $1 + \" (dirty=\" + (this.%1$s != $1) +\")\" );",
+                for(Object o : currentValue.getType().getAnnotations()) {
-					fieldName
+                    if(o instanceof Id)
-			);
+                        return "";
+                }
+
+                builder.append(entityMethodBody(currentValue));
+
+
+            }
+            catch (ClassNotFoundException e) {
+                e.printStackTrace();
+            }
+            catch (NotFoundException e) {
+                e.printStackTrace();
+            }
+            return builder.toString();
+		}
+
+        private String entityMethodBody(CtField currentValue) {
+            StringBuilder inlineBuilder = new StringBuilder();
+            try {
+                inlineBuilder.append("if ( $$_hibernate_getInterceptor() != null ");
+                //primitives || enums
+                if(currentValue.getType().isPrimitive() || currentValue.getType().isEnum()) {
+                    inlineBuilder.append("&& "+currentValue.getName()+" != $1)");
+                }
+                //simple data types
+                else if(currentValue.getType().getName().startsWith("java.lang") ||
+                        currentValue.getType().getName().startsWith("java.math.Big") ||
+                        currentValue.getType().getName().startsWith("java.sql.Time") ||
+                        currentValue.getType().getName().startsWith("java.sql.Date") ||
+                        currentValue.getType().getName().startsWith("java.util.Date") ||
+                        currentValue.getType().getName().startsWith("java.util.Calendar")
+                        ) {
+                    inlineBuilder.append("&& (("+currentValue.getName()+" == null) || (!" +currentValue.getName()+".equals( $1))))");
+                }
+                //all other objects
+                else {
+                    //if the field is a collection we return since we handle that in a separate method
+                    for(CtClass ctClass : currentValue.getType().getInterfaces()) {
+                        if(ctClass.getName().equals("java.util.Collection")) {
+
+                            //if the collection is not managed we should write it to the tracker
+                            //todo: should use EnhancementContext.isMappedCollection here instead
+                            if (currentValue.getAnnotation(OneToMany.class) != null ||
+                                    currentValue.getAnnotation(ManyToMany.class) != null ||
+                                    currentValue.getAnnotation(ElementCollection.class) != null) {
+                                return "";
+                            }
+                        }
+                    }
+
+                    //todo: for now just call equals, should probably do something else here
+                    inlineBuilder.append("&& (("+currentValue.getName()+" == null) || (!" +currentValue.getName()+".equals( $1))))");
+                }
+
+                inlineBuilder.append( EnhancerConstants.TRACKER_CHANGER_NAME+"(\""+currentValue.getName()+"\");");
+            } catch (NotFoundException e) {
+                e.printStackTrace();
+            } catch (ClassNotFoundException e) {
+                e.printStackTrace();
+            }
+            return inlineBuilder.toString();
         }
 	}
 

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/spi/EnhancerConstants.java

@@ -128,6 +128,54 @@ public class EnhancerConstants {
 	 */
 	public static final String INTERCEPTOR_SETTER_NAME = "$$_hibernate_setInterceptor";
 
+    /**
+     * Name of tracker field
+     */
+    public static final String TRACKER_FIELD_NAME = "$$_hibernate_tracker";
+
+    /**
+     * Name of method that add changed fields
+     */
+    public static final String TRACKER_CHANGER_NAME = "$$_hibernate_trackChange";
+
+    /**
+     * Name of method to see if any fields has changed
+     */
+    public static final String TRACKER_HAS_CHANGED_NAME = "$$_hibernate_hasDirtyAttributes";
+
+    /**
+     * Name of method to fetch dirty attributes
+     */
+    public static final String TRACKER_GET_NAME = "$$_hibernate_getDirtyAttributes";
+
+    /**
+     * Name of method to clear stored dirty attributes
+     */
+    public static final String TRACKER_CLEAR_NAME = "$$_hibernate_clearDirtyAttributes";
+
+    /**
+     * Name of method to check if collection fields are dirty
+     */
+    public static final String TRACKER_COLLECTION_CHANGED_NAME = "$$_hibernate_areCollectionFieldsDirty";
+
+    public static final String TRACKER_COLLECTION_NAME = "$$_hibernate_collectionTracker";
+    /**
+     * Name of method to get dirty collection field names
+     */
+    public static final String TRACKER_COLLECTION_CHANGED_FIELD_NAME = "$$_hibernate_getCollectionFieldDirtyNames";
+
+    public static final String TRACKER_COLLECTION_CLEAR_NAME = "$$_hibernate_clearDirtyCollectionNames";
+
+    public static final String TRACKER_COMPOSITE_DIRTY_CHECK = "$$_hibernate_areCompositeFieldsDirty";
+
+    public static final String TRACKER_COMPOSITE_DIRTY_FIELDS_GETTER = "$$_hibernate_getCompositeDirtyFields";
+
+    public static final String TRACKER_COMPOSITE_FIELD_NAME = "$$_hibernate_compositeOwners";
+
+    public static final String TRACKER_COMPOSITE_SET_OWNER = "$$_hibernate_setOwner";
+
+    public static final String TRACKER_COMPOSITE_CLEAR_OWNER = "$$_hibernate_clearOwner";
+
 	private EnhancerConstants() {
 	}
 }

hibernate-core/src/main/java/org/hibernate/engine/spi/CompositeOwner.java

@@ -0,0 +1,35 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.engine.spi;
+
+/**
+ * @author <a href="mailto:stale.pedersen@jboss.org">Ståle W. Pedersen</a>
+ */
+public interface CompositeOwner {
+
+    /**
+     * @param attributeName to be added to the dirty list
+     */
+    void $$_hibernate_trackChange(String attributeName);
+}

hibernate-core/src/main/java/org/hibernate/engine/spi/CompositeTracker.java

@@ -0,0 +1,34 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.engine.spi;
+
+/**
+ * @author <a href="mailto:stale.pedersen@jboss.org">Ståle W. Pedersen</a>
+ */
+public interface CompositeTracker {
+
+    void $$_hibernate_setOwner(String name, CompositeOwner tracker);
+
+    void $$_hibernate_clearOwner(String name);
+}

hibernate-core/src/main/java/org/hibernate/engine/spi/EntityEntry.java

@@ -236,6 +236,9 @@ public final class EntityEntry implements Serializable {
 				interceptor.clearDirty();
 			}
 		}
+        if( entity instanceof SelfDirtinessTracker)
+            ((SelfDirtinessTracker) entity).$$_hibernate_clearDirtyAttributes();
+
 		persistenceContext.getSession()
 				.getFactory()
 				.getCustomEntityDirtinessStrategy()
@@ -299,6 +302,10 @@ public final class EntityEntry implements Serializable {
 
 	@SuppressWarnings( {"SimplifiableIfStatement"})
 	private boolean isUnequivocallyNonDirty(Object entity) {
+
+        if(entity instanceof SelfDirtinessTracker)
+            return ((SelfDirtinessTracker) entity).$$_hibernate_hasDirtyAttributes();
+
 		final CustomEntityDirtinessStrategy customEntityDirtinessStrategy =
 				persistenceContext.getSession().getFactory().getCustomEntityDirtinessStrategy();
 		if ( customEntityDirtinessStrategy.canDirtyCheck( entity, getPersister(), (Session) persistenceContext.getSession() ) ) {

hibernate-core/src/main/java/org/hibernate/engine/spi/SelfDirtinessTracker.java

@@ -0,0 +1,49 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.engine.spi;
+
+import java.util.Set;
+
+/**
+ * Specify if an entity class is instrumented to track field changes
+ *
+ * @author <a href="mailto:stale.pedersen@jboss.org">Ståle W. Pedersen</a>
+ */
+public interface SelfDirtinessTracker {
+
+    /**
+     * Return true if any fields has been changed
+     */
+    boolean $$_hibernate_hasDirtyAttributes();
+
+    /**
+     * Get the field names of all the fields thats been changed
+     */
+    Set<String> $$_hibernate_getDirtyAttributes();
+
+    /**
+     * Clear the stored dirty attributes
+     */
+    void $$_hibernate_clearDirtyAttributes();
+}

hibernate-core/src/main/java/org/hibernate/event/internal/DefaultFlushEntityEventListener.java

@@ -26,6 +26,7 @@ package org.hibernate.event.internal;
 import java.io.Serializable;
 import java.util.Arrays;
 
+import org.hibernate.engine.spi.SelfDirtinessTracker;
 import org.jboss.logging.Logger;
 
 import org.hibernate.AssertionFailure;
@@ -491,6 +492,12 @@ public class DefaultFlushEntityEventListener implements FlushEntityEventListener
 		);
 
 		if ( dirtyProperties == null ) {
+            if(entity instanceof SelfDirtinessTracker) {
+                if(((SelfDirtinessTracker) entity).$$_hibernate_hasDirtyAttributes()) {
+                   dirtyProperties = persister.resolveAttributeIndexes(((SelfDirtinessTracker) entity).$$_hibernate_getDirtyAttributes());
+                }
+            }
+            else {
                 // see if the custom dirtiness strategy can tell us...
                 class DirtyCheckContextImpl implements CustomEntityDirtinessStrategy.DirtyCheckContext {
                     int[] found = null;
@@ -510,6 +517,7 @@ public class DefaultFlushEntityEventListener implements FlushEntityEventListener
                         context
                 );
                 dirtyProperties = context.found;
+            }
 		}
 
 		event.setDatabaseSnapshot(null);

hibernate-core/src/main/java/org/hibernate/persister/entity/AbstractEntityPersister.java

@@ -2148,6 +2148,20 @@ public abstract class AbstractEntityPersister
 		return subclassPropertyNameClosure;
 	}
 
+    @Override
+    public int[] resolveAttributeIndexes(Set<String> properties) {
+        Iterator<String> iter = properties.iterator();
+        int[] fields = new int[properties.size()];
+        int counter = 0;
+        while(iter.hasNext()) {
+            Integer index = entityMetamodel.getPropertyIndexOrNull(iter.next());
+            if(index != null)
+                fields[counter++] = index;
+        }
+
+        return fields;
+    }
+
 	protected String[] getSubclassPropertySubclassNameClosure() {
 		return subclassPropertySubclassNameClosure;
 	}

hibernate-core/src/main/java/org/hibernate/persister/entity/EntityPersister.java

@@ -25,6 +25,7 @@ package org.hibernate.persister.entity;
 
 import java.io.Serializable;
 import java.util.Map;
+import java.util.Set;
 
 import org.hibernate.EntityMode;
 import org.hibernate.HibernateException;
@@ -765,4 +766,7 @@ public interface EntityPersister extends OptimisticCacheSource, EntityDefinition
 	public EntityInstrumentationMetadata getInstrumentationMetadata();
 	
 	public FilterAliasGenerator getFilterAliasGenerator(final String rootAlias);
+
+    public int[] resolveAttributeIndexes(Set<String> properties);
+
 }

hibernate-core/src/main/java/org/hibernate/tool/enhance/EnhancementTask.java

@@ -23,7 +23,11 @@
  */
 package org.hibernate.tool.enhance;
 
+import javax.persistence.ElementCollection;
+import javax.persistence.Embeddable;
 import javax.persistence.Entity;
+import javax.persistence.ManyToMany;
+import javax.persistence.OneToMany;
 import javax.persistence.Transient;
 import java.io.File;
 import java.io.FileInputStream;
@@ -84,39 +88,69 @@ public class EnhancementTask extends Task implements EnhancementContext {
 					continue;
 				}
 
-				processClassFile( javaClassFile );
+				processClassFile( javaClassFile);
-			}
 			}
 		}
 
+	}
+
+    /**
+     * Atm only process files annotated with either @Entity or @Embeddable
+     * @param javaClassFile
+     */
     private void processClassFile(File javaClassFile) {
 		try {
 			final CtClass ctClass = classPool.makeClass( new FileInputStream( javaClassFile ) );
-			if ( ! shouldInclude( ctClass ) ) {
+            if(this.isEntityClass(ctClass))
-				return;
+                processEntityClassFile(javaClassFile, ctClass);
+            else if(this.isCompositeClass(ctClass))
+                processCompositeClassFile(javaClassFile, ctClass);
+
+        }
+        catch (IOException e) {
+            throw new BuildException(
+                    String.format( "Error processing included file [%s]", javaClassFile.getAbsolutePath() ), e );
+        }
     }
 
-			final byte[] enhancedBytecode;
+    private void processEntityClassFile(File javaClassFile, CtClass ctClass ) {
         try {
-				enhancedBytecode = enhancer.enhance( ctClass.getName(), ctClass.toBytecode() );
+            byte[] result = enhancer.enhance( ctClass.getName(), ctClass.toBytecode() );
+            if(result != null)
+                writeEnhancedClass(javaClassFile, result);
         }
         catch (Exception e) {
             log( "Unable to enhance class [" + ctClass.getName() + "]", e, Project.MSG_WARN );
             return;
         }
+    }
 
+    private void processCompositeClassFile(File javaClassFile, CtClass ctClass) {
+        try {
+            byte[] result = enhancer.enhanceComposite(ctClass.getName(), ctClass.toBytecode());
+            if(result != null)
+                writeEnhancedClass(javaClassFile, result);
+        }
+        catch (Exception e) {
+            log( "Unable to enhance class [" + ctClass.getName() + "]", e, Project.MSG_WARN );
+            return;
+        }
+    }
+
+    private void writeEnhancedClass(File javaClassFile, byte[] result) {
+        try {
 			if ( javaClassFile.delete() ) {
                     if ( ! javaClassFile.createNewFile() ) {
-					log( "Unable to recreate class file [" + ctClass.getName() + "]", Project.MSG_INFO );
+                        log( "Unable to recreate class file [" + javaClassFile.getName() + "]", Project.MSG_INFO );
                     }
             }
 			else {
-				log( "Unable to delete class file [" + ctClass.getName() + "]", Project.MSG_INFO );
+				log( "Unable to delete class file [" + javaClassFile.getName() + "]", Project.MSG_INFO );
 			}
 
 			FileOutputStream outputStream = new FileOutputStream( javaClassFile, false );
 			try {
-				outputStream.write( enhancedBytecode );
+				outputStream.write( result);
 				outputStream.flush();
 			}
 			finally {
@@ -132,20 +166,11 @@ public class EnhancementTask extends Task implements EnhancementContext {
         }
         catch (IOException e) {
             throw new BuildException(
-					String.format( "Error processing included file [%s]", javaClassFile.getAbsolutePath() ),
+                    String.format( "Error processing included file [%s]", javaClassFile.getAbsolutePath() ), e );
-					e
-			);
         }
     }
 
-	private boolean shouldInclude(CtClass ctClass) {
-		// we currently only handle entity enhancement
-		return ctClass.hasAnnotation( Entity.class );
-	}
-
-
 	// EnhancementContext impl ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
 	@Override
 	public ClassLoader getLoadingClassLoader() {
 		return getClass().getClassLoader();
@@ -153,18 +178,17 @@ public class EnhancementTask extends Task implements EnhancementContext {
 
 	@Override
 	public boolean isEntityClass(CtClass classDescriptor) {
-		// currently we only call enhance on the classes with @Entity, so here we always return true
+        return classDescriptor.hasAnnotation(Entity.class);
-		return true;
     }
 
 	@Override
 	public boolean isCompositeClass(CtClass classDescriptor) {
-		return false;
+        return classDescriptor.hasAnnotation(Embeddable.class);
 	}
 
 	@Override
 	public boolean doDirtyCheckingInline(CtClass classDescriptor) {
-		return false;
+		return true;
 	}
 
 	@Override
@@ -183,6 +207,18 @@ public class EnhancementTask extends Task implements EnhancementContext {
 		return ! ctField.hasAnnotation( Transient.class );
 	}
 
+    @Override
+    public boolean isMappedCollection(CtField field) {
+        try {
+            return (field.getAnnotation(OneToMany.class) != null ||
+                    field.getAnnotation(ManyToMany.class) != null ||
+                    field.getAnnotation(ElementCollection.class) != null);
+        }
+        catch (ClassNotFoundException e) {
+            return false;
+        }
+    }
+
 	@Override
 	public CtField[] order(CtField[] persistentFields) {
 		// for now...

hibernate-core/src/main/java/org/hibernate/tuple/entity/PojoEntityTuplizer.java

@@ -510,6 +510,11 @@ public class PojoEntityTuplizer extends AbstractEntityTuplizer {
 			//TODO: if we support multiple fetch groups, we would need
 			//      to clone the set of lazy properties!
 			FieldInterceptionHelper.injectFieldInterceptor( entity, getEntityName(), lazyProps, session );
+
+            //also clear the fields that are marked as dirty in the dirtyness tracker
+            if(entity instanceof org.hibernate.engine.spi.SelfDirtinessTracker) {
+                ((org.hibernate.engine.spi.SelfDirtinessTracker) entity).$$_hibernate_clearDirtyAttributes();
+            }
 		}
 	}
 

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/Address.java

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2012 Red Hat, Inc. and/or its affiliates.
+ *
+ * Licensed under the Eclipse Public License version 1.0, available at
+ * http://www.eclipse.org/legal/epl-v10.html
+ */
+package org.hibernate.test.bytecode.enhancement;
+
+import javax.persistence.Embeddable;
+import javax.persistence.Embedded;
+import java.io.Serializable;
+
+/**
+ * @author <a href="mailto:stale.pedersen@jboss.org">Ståle W. Pedersen</a>
+ */
+@Embeddable
+public class Address implements Serializable {
+
+    private String street1;
+    private String street2;
+    private String city;
+    private String state;
+
+    @Embedded
+    private Country country;
+    private String zip;
+    private String phone;
+
+    public Address() {
+    }
+
+    public String getStreet1() {
+        return street1;
+    }
+
+    public void setStreet1(String street1) {
+        this.street1 = street1;
+    }
+
+    public String getStreet2() {
+        return street2;
+    }
+
+    public void setStreet2(String street2) {
+        this.street2 = street2;
+    }
+
+    public String getCity() {
+        return city;
+    }
+
+    public void setCity(String city) {
+        this.city = city;
+    }
+
+    public String getState() {
+        return state;
+    }
+
+    public void setState(String state) {
+        this.state = state;
+    }
+
+    public Country getCountry() {
+        return country;
+    }
+
+    public void setCountry(Country country) {
+        this.country = country;
+    }
+
+    public String getZip() {
+        return zip;
+    }
+
+    public void setZip(String zip) {
+        this.zip = zip;
+    }
+
+    public String getPhone() {
+        return phone;
+    }
+
+    public void setPhone(String phone) {
+        this.phone = phone;
+    }
+
+}

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/CompositeOwnerTrackerTest.java

@@ -0,0 +1,84 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.test.bytecode.enhancement;
+
+import org.hibernate.bytecode.enhance.spi.CompositeOwnerTracker;
+import org.hibernate.engine.spi.CompositeOwner;
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author <a href="mailto:stale.pedersen@jboss.org">Ståle W. Pedersen</a>
+ */
+
+public class CompositeOwnerTrackerTest {
+
+    private int counter = 0;
+
+    @Test
+    public void testCompositeOwnerTracker() {
+
+        CompositeOwnerTracker tracker = new CompositeOwnerTracker();
+        tracker.add("foo", new TestCompositeOwner());
+
+        tracker.callOwner(".street1");
+        assertEquals(1, counter);
+        tracker.add("bar", new TestCompositeOwner());
+        tracker.callOwner(".city");
+        assertEquals(3, counter);
+
+        tracker.removeOwner("foo");
+
+        tracker.callOwner(".country");
+        assertEquals(4, counter);
+        tracker.removeOwner("bar");
+
+        tracker.callOwner(".country");
+
+        tracker.add("moo", new TestCompositeOwner());
+        tracker.callOwner(".country");
+        assertEquals(5, counter);
+    }
+
+    class TestCompositeOwner implements CompositeOwner {
+
+        @Override
+        public void $$_hibernate_trackChange(String attributeName) {
+            if(counter == 0)
+                assertEquals("foo.street1", attributeName);
+            if(counter == 1)
+                assertEquals("foo.city", attributeName);
+            if(counter == 2)
+                assertEquals("bar.city", attributeName);
+            if(counter == 3)
+                assertEquals("bar.country", attributeName);
+            if(counter == 4)
+                assertEquals("moo.country", attributeName);
+            counter++;
+        }
+    }
+}
+
+

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/Country.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2012 Red Hat, Inc. and/or its affiliates.
+ *
+ * Licensed under the Eclipse Public License version 1.0, available at
+ * http://www.eclipse.org/legal/epl-v10.html
+ */
+package org.hibernate.test.bytecode.enhancement;
+
+import javax.persistence.Embeddable;
+
+/**
+ * @author <a href="mailto:stale.pedersen@jboss.org">Ståle W. Pedersen</a>
+ */
+@Embeddable
+public class Country {
+
+    private String name;
+
+    public String getName() {
+        return name;
+    }
+
+    public void setName(String name) {
+        this.name = name;
+    }
+}

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/EnhancerTest.java

@@ -25,7 +25,12 @@ package org.hibernate.test.bytecode.enhancement;
 
 import java.io.ByteArrayInputStream;
 import java.lang.reflect.Method;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Set;
 
 import javassist.ClassPool;
 import javassist.CtClass;
@@ -46,11 +51,17 @@ import org.junit.Test;
 
 import org.hibernate.testing.junit4.BaseUnitTestCase;
 
+import javax.persistence.ElementCollection;
+import javax.persistence.ManyToMany;
+import javax.persistence.OneToMany;
+
 import static org.hibernate.testing.junit4.ExtraAssertions.assertTyping;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
 
 /**
  * @author Steve Ebersole
@@ -87,6 +98,18 @@ public class EnhancerTest extends BaseUnitTestCase {
 			return true;
 		}
 
+        @Override
+        public boolean isMappedCollection(CtField field) {
+            try {
+                return (field.getAnnotation(OneToMany.class) != null ||
+                        field.getAnnotation(ManyToMany.class) != null ||
+                        field.getAnnotation(ElementCollection.class) != null);
+            }
+            catch (ClassNotFoundException e) {
+                return false;
+            }
+        }
+
         @Override
 		public boolean isPersistentField(CtField ctField) {
 			return true;
@@ -108,6 +131,7 @@ public class EnhancerTest extends BaseUnitTestCase {
 		Enhancer enhancer = new Enhancer( enhancementContext );
 		CtClass entityCtClass = generateCtClassForAnEntity( entityClassToEnhance );
 		byte[] original = entityCtClass.toBytecode();
+		//byte[] enhanced = enhancer.enhance( entityCtClass.getName(), original );
         byte[] enhanced = enhancer.enhance( entityCtClass.getName(), original );
 		assertFalse( "entity was not enhanced", Arrays.equals( original, enhanced ) );
 
@@ -115,6 +139,24 @@ public class EnhancerTest extends BaseUnitTestCase {
 		ClassPool cp = new ClassPool( false );
 		cp.appendClassPath( new LoaderClassPath( cl ) );
 		CtClass enhancedCtClass = cp.makeClass( new ByteArrayInputStream( enhanced ) );
+        enhancedCtClass.debugWriteFile("/tmp");
+        //just for debugging
+        Class addressClass = null;
+        Class countryClass = null;
+        if(entityClassToEnhance.getName().endsWith("SimpleEntity")) {
+            CtClass addressCtClass = generateCtClassForAnEntity( Address.class );
+            byte[] enhancedAddress = enhancer.enhanceComposite(Address.class.getName(), addressCtClass.toBytecode());
+            CtClass enhancedCtClassAddress = cp.makeClass( new ByteArrayInputStream( enhancedAddress ) );
+            enhancedCtClassAddress.debugWriteFile("/tmp");
+            addressClass = enhancedCtClassAddress.toClass( cl, this.getClass().getProtectionDomain() );
+
+            CtClass countryCtClass = generateCtClassForAnEntity( Country.class );
+            byte[] enhancedCountry = enhancer.enhanceComposite(Country.class.getName(), countryCtClass.toBytecode());
+            CtClass enhancedCtClassCountry = cp.makeClass( new ByteArrayInputStream( enhancedCountry ) );
+            enhancedCtClassCountry.debugWriteFile("/tmp");
+            countryClass = enhancedCtClassCountry.toClass( cl, this.getClass().getProtectionDomain() );
+
+        }
 		Class entityClass = enhancedCtClass.toClass( cl, this.getClass().getProtectionDomain() );
 		Object entityInstance = entityClass.newInstance();
 
@@ -151,20 +193,96 @@ public class EnhancerTest extends BaseUnitTestCase {
 		entityClass.getMethod( "getId" ).invoke( entityInstance );
 
 		interceptorSetter.invoke( entityInstance, new LocalPersistentAttributeInterceptor() );
-		assertNotNull( interceptorGetter.invoke( entityInstance ) );
+		assertNotNull(interceptorGetter.invoke(entityInstance));
 
 		// dirty checking is unfortunately just printlns for now... just verify the test output
 		entityClass.getMethod( "getId" ).invoke( entityInstance );
 		entityClass.getMethod( "setId", Long.class ).invoke( entityInstance, entityClass.getMethod( "getId" ).invoke( entityInstance ) );
 		entityClass.getMethod( "setId", Long.class ).invoke( entityInstance, 1L );
+        assertTrue((Boolean) entityClass.getMethod("$$_hibernate_hasDirtyAttributes").invoke(entityInstance));
 
 		entityClass.getMethod( "isActive" ).invoke( entityInstance );
 		entityClass.getMethod( "setActive", boolean.class ).invoke( entityInstance, entityClass.getMethod( "isActive" ).invoke( entityInstance ) );
-		entityClass.getMethod( "setActive", boolean.class ).invoke( entityInstance, true );
+		entityClass.getMethod( "setActive", boolean.class ).invoke(entityInstance, true);
 
 		entityClass.getMethod( "getSomeNumber" ).invoke( entityInstance );
 		entityClass.getMethod( "setSomeNumber", long.class ).invoke( entityInstance, entityClass.getMethod( "getSomeNumber" ).invoke( entityInstance ) );
-		entityClass.getMethod( "setSomeNumber", long.class ).invoke( entityInstance, 1L );
+		entityClass.getMethod( "setSomeNumber", long.class ).invoke(entityInstance, 1L);
+        assertEquals(3, ((Set) entityClass.getMethod("$$_hibernate_getDirtyAttributes").invoke(entityInstance)).size());
+
+        entityClass.getMethod("$$_hibernate_clearDirtyAttributes").invoke(entityInstance);
+        //setting the same value should not make it dirty
+        entityClass.getMethod( "setSomeNumber", long.class ).invoke(entityInstance, 1L);
+        //assertEquals(0, ((Set) entityClass.getMethod("$$_hibernate_getDirtyAttributes").invoke(entityInstance)).size());
+        //assertFalse((Boolean) entityClass.getMethod("$$_hibernate_hasDirtyAttributes").invoke(entityInstance));
+
+        if(entityClass.getName().endsWith("SimpleEntity")) {
+
+            List<String> strings = new ArrayList<String>();
+            strings.add("FooBar");
+
+            entityClass.getMethod( "setSomeStrings", java.util.List.class ).invoke(entityInstance, strings);
+
+            assertTrue((Boolean) entityClass.getMethod("$$_hibernate_hasDirtyAttributes").invoke(entityInstance));
+
+            Set tracked = (Set) entityClass.getMethod(EnhancerConstants.TRACKER_GET_NAME).invoke(entityInstance);
+            assertEquals(1, tracked.size());
+            assertEquals("someStrings", tracked.iterator().next());
+
+            entityClass.getMethod("$$_hibernate_clearDirtyAttributes").invoke(entityInstance);
+
+            ((List) entityClass.getMethod( "getSomeStrings").invoke(entityInstance)).add("JADA!");
+            Boolean isDirty = (Boolean) entityClass.getMethod("$$_hibernate_hasDirtyAttributes").invoke(entityInstance);
+            assertTrue(isDirty);
+            assertEquals("someStrings",
+                    ((Set) entityClass.getMethod(EnhancerConstants.TRACKER_GET_NAME).invoke(entityInstance)).iterator().next());
+            entityClass.getMethod("$$_hibernate_clearDirtyAttributes").invoke(entityInstance);
+
+            //this should not set the entity to dirty
+            Set<Integer> ints = new HashSet<Integer>();
+            ints.add(42);
+            entityClass.getMethod( "setSomeInts", java.util.Set.class ).invoke(entityInstance, ints);
+            isDirty = (Boolean) entityClass.getMethod("$$_hibernate_hasDirtyAttributes").invoke(entityInstance);
+            assertFalse(isDirty);
+
+            //testing composite object
+            assert addressClass != null;
+            Object address =  addressClass.newInstance();
+
+            assert countryClass != null;
+            Object country =  countryClass.newInstance();
+
+            //Method adrInterceptorGetter = addressClass.getMethod( EnhancerConstants.INTERCEPTOR_GETTER_NAME );
+            //Method adrInterceptorSetter = addressClass.getMethod( EnhancerConstants.INTERCEPTOR_SETTER_NAME, PersistentAttributeInterceptor.class );
+            //adrInterceptorSetter.invoke( address, new LocalPersistentAttributeInterceptor() );
+
+            entityClass.getMethod("setAddress", addressClass).invoke(entityInstance, address);
+
+            addressClass.getMethod("setCity", String.class).invoke(address, "Arendal");
+
+            tracked = (Set) entityClass.getMethod(EnhancerConstants.TRACKER_GET_NAME).invoke(entityInstance);
+            assertEquals(2, tracked.size());
+            Iterator iter = tracked.iterator();
+            assertEquals("address", iter.next());
+            assertEquals("address.city", iter.next());
+
+            entityClass.getMethod("$$_hibernate_clearDirtyAttributes").invoke(entityInstance);
+
+            //make sure that new composite instances are cleared
+            Object address2 =  addressClass.newInstance();
+
+            entityClass.getMethod("setAddress", addressClass).invoke(entityInstance, address2);
+            addressClass.getMethod("setStreet1", String.class).invoke(address, "Heggedalveien");
+
+            tracked = (Set) entityClass.getMethod(EnhancerConstants.TRACKER_GET_NAME).invoke(entityInstance);
+            assertEquals(1, tracked.size());
+
+            addressClass.getMethod("setCountry", countryClass).invoke(address2, country);
+            countryClass.getMethod("setName", String.class).invoke(country, "Norway");
+
+            tracked = (Set) entityClass.getMethod(EnhancerConstants.TRACKER_GET_NAME).invoke(entityInstance);
+            assertEquals(3, tracked.size());
+        }
 	}
 
 	private CtClass generateCtClassForAnEntity(Class entityClassToEnhance) throws Exception {

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/SimpleEntity.java

@@ -23,8 +23,12 @@
  */
 package org.hibernate.test.bytecode.enhancement;
 
+import javax.persistence.Embedded;
 import javax.persistence.Entity;
 import javax.persistence.Id;
+import javax.persistence.OneToMany;
+import java.util.List;
+import java.util.Set;
 
 /**
  * @author Steve Ebersole
@@ -35,6 +39,13 @@ public class SimpleEntity {
 	private String name;
 	private boolean active;
 	private long someNumber;
+    private List<String> someStrings;
+
+    @OneToMany
+    private Set<Integer> someInts;
+
+    @Embedded
+    private Address address;
 
 	@Id
 	public Long getId() {
@@ -68,4 +79,28 @@ public class SimpleEntity {
 	public void setSomeNumber(long someNumber) {
 		this.someNumber = someNumber;
 	}
+
+    public List<String> getSomeStrings() {
+        return someStrings;
+    }
+
+    public void setSomeStrings(List<String> someStrings) {
+        this.someStrings = someStrings;
+    }
+
+    public Address getAddress() {
+        return address;
+    }
+
+    public void setAddress(Address address) {
+        this.address = address;
+    }
+
+    public Set<Integer> getSomeInts() {
+        return someInts;
+    }
+
+    public void setSomeInts(Set<Integer> someInts) {
+        this.someInts = someInts;
+    }
 }

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/customer/Address.java

@@ -0,0 +1,122 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */package org.hibernate.test.bytecode.enhancement.customer;
+
+import java.io.Serializable;
+
+/**
+ * @author <a href="mailto:stale.pedersen@jboss.org">Ståle W. Pedersen</a>
+ */
+public class Address implements Serializable {
+    private String street1;
+    private String street2;
+    private String city;
+    private String state;
+    private String country;
+    private String zip;
+    private String phone;
+
+    public Address() {
+    }
+    public Address(String street1, String street2, String city, String state,
+                   String country, String zip, String phone) {
+        this.street1 = street1;
+        this.street2 = street2;
+        this.city    = city;
+        this.state   = state;
+        this.country = country;
+        setZip(zip);
+        setPhone(phone);
+    }
+
+    public String toString() {
+        return street1 + "\n" + street2 + "\n" + city + "," + state + " " + zip + "\n" + phone;
+    }
+
+    public String getStreet1() {
+        return street1;
+    }
+
+    public void setStreet1(String street1) {
+        this.street1 = street1;
+    }
+
+    public String getStreet2() {
+        return street2;
+    }
+
+    public void setStreet2(String street2) {
+        this.street2 = street2;
+    }
+
+    public String getCity() {
+        return city;
+    }
+
+    public void setCity(String city) {
+        this.city = city;
+    }
+
+    public String getState() {
+        return state;
+    }
+
+    public void setState(String state) {
+        this.state = state;
+    }
+
+    public String getCountry() {
+        return country;
+    }
+
+    public void setCountry(String country) {
+        this.country = country;
+    }
+
+    public String getZip() {
+        return zip;
+    }
+
+    public void setZip(String zip) {
+        assertNumeric(zip, "Non-numeric zip ");
+        this.zip = zip;
+
+    }
+
+    public String getPhone() {
+        return phone;
+    }
+
+    public void setPhone(String phone) {
+        assertNumeric(zip, "Non-numeric phone ");
+        this.phone = phone;
+    }
+
+    void assertNumeric(String s, String error) {
+        for (int i=0; i<s.length(); i++) {
+            if (!Character.isDigit(s.charAt(i))) {
+                throw new IllegalArgumentException(error + s);
+            }
+        }
+    }
+}

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/customer/Customer.java

@@ -0,0 +1,249 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.test.bytecode.enhancement.customer;
+
+import javax.persistence.AttributeOverride;
+import javax.persistence.AttributeOverrides;
+import javax.persistence.CascadeType;
+import javax.persistence.Column;
+import javax.persistence.Embedded;
+import javax.persistence.Entity;
+import javax.persistence.FetchType;
+import javax.persistence.Id;
+import javax.persistence.OneToMany;
+import javax.persistence.Table;
+import javax.persistence.Temporal;
+import javax.persistence.TemporalType;
+import javax.persistence.Version;
+import java.math.BigDecimal;
+import java.util.ArrayList;
+import java.util.Calendar;
+import java.util.List;
+
+/**
+ * @author <a href="mailto:stale.pedersen@jboss.org">Ståle W. Pedersen</a>
+ */
+@Entity
+@Table(name="O_CUSTOMER")
+public class Customer {
+    public static final String QUERY_ALL = "Customer.selectAll";
+    public static final String QUERY_COUNT = "Customer.count";
+    public static final String QUERY_BY_CREDIT = "Customer.selectByCreditLimit";
+
+    public static final String BAD_CREDIT = "BC";
+
+    @Id
+    @Column(name="C_ID")
+    private int id;
+
+    @Column(name="C_FIRST")
+    private String firstName;
+
+    @Column(name="C_LAST")
+    private String lastName;
+
+    @Column(name="C_CONTACT")
+    private String contact;
+
+    @Column(name="C_CREDIT")
+    private String credit;
+
+    @Column(name="C_CREDIT_LIMIT")
+    private BigDecimal creditLimit;
+
+    @Column(name="C_SINCE")
+    @Temporal(TemporalType.DATE)
+    private Calendar since;
+
+    @Column(name="C_BALANCE")
+    private BigDecimal balance;
+
+    @Column(name="C_YTD_PAYMENT")
+    private BigDecimal ytdPayment;
+
+    @OneToMany(mappedBy="customer", cascade= CascadeType.ALL, fetch= FetchType.EAGER)
+    private List<CustomerInventory> customerInventories;
+
+    @Embedded
+    @AttributeOverrides(
+            {@AttributeOverride(name="street1",column=@Column(name="C_STREET1")),
+                    @AttributeOverride(name="street2",column=@Column(name="C_STREET2")),
+                    @AttributeOverride(name="city",   column=@Column(name="C_CITY")),
+                    @AttributeOverride(name="state",  column=@Column(name="C_STATE")),
+                    @AttributeOverride(name="country",column=@Column(name="C_COUNTRY")),
+                    @AttributeOverride(name="zip",    column=@Column(name="C_ZIP")),
+                    @AttributeOverride(name="phone",  column=@Column(name="C_PHONE"))})
+    private Address       address;
+
+    @Version
+    @Column(name = "C_VERSION")
+    private int version;
+
+    public Customer() {
+    }
+
+    public Customer(String first, String last, Address address,
+                    String contact, String credit, BigDecimal creditLimit,
+                    BigDecimal balance, BigDecimal YtdPayment) {
+
+        this.firstName   = first;
+        this.lastName    = last;
+        this.address     = address;
+        this.contact     = contact;
+        this.since       = Calendar.getInstance();
+        this.credit      = credit;
+        this.creditLimit = creditLimit;
+        this.balance     = balance;
+        this.ytdPayment  = YtdPayment;
+    }
+
+    public Integer getId() {
+        return id;
+    }
+
+    public void setId(Integer customerId) {
+        this.id = customerId;
+    }
+
+    public String getFirstName() {
+        return firstName;
+    }
+
+    public void setFirstName(String firstName) {
+        this.firstName = firstName;
+    }
+
+    public String getLastName() {
+        return lastName;
+    }
+
+    public void setLastName(String lastName) {
+        this.lastName = lastName;
+    }
+
+    public Address getAddress() {
+        return address;
+    }
+
+    public void setAddress(Address address) {
+        this.address = address;
+    }
+
+    public String getContact() {
+        return contact;
+    }
+
+    public void setContact(String contact) {
+        this.contact = contact;
+    }
+
+    public String getCredit() {
+        return credit;
+    }
+
+    public void setCredit(String credit) {
+        this.credit = credit;
+    }
+
+    public BigDecimal getCreditLimit() {
+        return creditLimit;
+    }
+
+    public void setCreditLimit(BigDecimal creditLimit) {
+        this.creditLimit = creditLimit;
+    }
+
+    public Calendar getSince() {
+        return since;
+    }
+
+    public void setSince(Calendar since) {
+        this.since = since;
+    }
+
+    public BigDecimal getBalance() {
+        return balance;
+    }
+
+    public void setBalance(BigDecimal balance) {
+        this.balance = balance;
+    }
+
+    public void changeBalance(BigDecimal change) {
+        setBalance(balance.add(change).setScale(2, BigDecimal.ROUND_DOWN));
+    }
+
+    public BigDecimal getYtdPayment() {
+        return ytdPayment;
+    }
+
+    public void setYtdPayment(BigDecimal ytdPayment) {
+        this.ytdPayment = ytdPayment;
+    }
+
+    public List<CustomerInventory> getInventories() {
+        if (customerInventories == null){
+            customerInventories = new ArrayList<CustomerInventory>();
+        }
+        return customerInventories;
+    }
+
+    public CustomerInventory addInventory(String item, int quantity,
+                                          BigDecimal totalValue) {
+
+        CustomerInventory inventory = new CustomerInventory(this, item,
+                quantity, totalValue);
+        getInventories().add(inventory);
+        return inventory;
+    }
+
+    public int getVersion() {
+        return version;
+    }
+
+    public boolean hasSufficientCredit(BigDecimal amount) {
+        return !BAD_CREDIT.equals(getCredit())
+                && creditLimit != null
+                && creditLimit.compareTo(amount) >= 0;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o)
+            return true;
+        if (o == null || getClass() != o.getClass())
+            return false;
+        return id == ((Customer) o).id;
+    }
+
+    @Override
+    public int hashCode() {
+        return new Integer(id).hashCode();
+    }
+
+    @Override
+    public String toString() {
+        return this.getFirstName() + " " + this.getLastName();
+    }
+}

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/customer/CustomerEnhancerTest.java

@@ -0,0 +1,384 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.test.bytecode.enhancement.customer;
+
+import javassist.ClassPool;
+import javassist.CtClass;
+import javassist.CtField;
+import javassist.LoaderClassPath;
+import org.hibernate.EntityMode;
+import org.hibernate.LockMode;
+import org.hibernate.bytecode.enhance.spi.EnhancementContext;
+import org.hibernate.bytecode.enhance.spi.Enhancer;
+import org.hibernate.bytecode.enhance.spi.EnhancerConstants;
+import org.hibernate.engine.spi.EntityEntry;
+import org.hibernate.engine.spi.ManagedEntity;
+import org.hibernate.engine.spi.PersistentAttributeInterceptor;
+import org.hibernate.engine.spi.Status;
+import org.hibernate.testing.junit4.BaseUnitTestCase;
+import org.junit.Test;
+
+import javax.persistence.ElementCollection;
+import javax.persistence.ManyToMany;
+import javax.persistence.OneToMany;
+import java.io.ByteArrayInputStream;
+import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Set;
+
+import static org.hibernate.testing.junit4.ExtraAssertions.assertTyping;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * @author Steve Ebersole
+ */
+public class CustomerEnhancerTest extends BaseUnitTestCase {
+	private static EnhancementContext enhancementContext = new EnhancementContext() {
+		@Override
+		public ClassLoader getLoadingClassLoader() {
+			return getClass().getClassLoader();
+		}
+
+		@Override
+		public boolean isEntityClass(CtClass classDescriptor) {
+			return true;
+		}
+
+		@Override
+		public boolean isCompositeClass(CtClass classDescriptor) {
+			return false;
+		}
+
+		@Override
+		public boolean doDirtyCheckingInline(CtClass classDescriptor) {
+			return true;
+		}
+
+		@Override
+		public boolean hasLazyLoadableAttributes(CtClass classDescriptor) {
+			return true;
+		}
+
+		@Override
+		public boolean isLazyLoadable(CtField field) {
+			return true;
+		}
+
+        @Override
+        public boolean isMappedCollection(CtField field) {
+            try {
+                return (field.getAnnotation(OneToMany.class) != null ||
+                        field.getAnnotation(ManyToMany.class) != null ||
+                        field.getAnnotation(ElementCollection.class) != null);
+            }
+            catch (ClassNotFoundException e) {
+                return false;
+            }
+        }
+
+        @Override
+		public boolean isPersistentField(CtField ctField) {
+			return true;
+		}
+
+		@Override
+		public CtField[] order(CtField[] persistentFields) {
+			return persistentFields;
+		}
+	};
+
+	@Test
+	public void testEnhancement() throws Exception {
+		testFor( Customer.class );
+	}
+
+	private void testFor(Class entityClassToEnhance) throws Exception {
+		Enhancer enhancer = new Enhancer( enhancementContext );
+		CtClass entityCtClass = generateCtClassForAnEntity( entityClassToEnhance );
+		byte[] original = entityCtClass.toBytecode();
+		//byte[] enhanced = enhancer.enhance( entityCtClass.getName(), original );
+        byte[] enhanced = enhancer.enhance( entityCtClass.getName(), original );
+		assertFalse( "entity was not enhanced", Arrays.equals( original, enhanced ) );
+
+		ClassLoader cl = new ClassLoader() { };
+		ClassPool cp = new ClassPool( false );
+		cp.appendClassPath( new LoaderClassPath( cl ) );
+		CtClass enhancedCtClass = cp.makeClass( new ByteArrayInputStream( enhanced ) );
+        enhancedCtClass.debugWriteFile("/tmp");
+        //just for debugging
+        Class addressClass = null;
+        Class customerInventoryClass = null;
+        Class supplierComponentPKClass = null;
+
+        CtClass addressCtClass = generateCtClassForAnEntity( org.hibernate.test.bytecode.enhancement.customer.Address.class );
+        byte[] enhancedAddress = enhancer.enhanceComposite(Address.class.getName(), addressCtClass.toBytecode());
+        CtClass enhancedCtClassAddress = cp.makeClass( new ByteArrayInputStream( enhancedAddress ) );
+        enhancedCtClassAddress.debugWriteFile("/tmp");
+        addressClass = enhancedCtClassAddress.toClass( cl, this.getClass().getProtectionDomain() );
+
+        CtClass customerInventoryCtClass = generateCtClassForAnEntity( CustomerInventory.class );
+        byte[] enhancedCustomerInventory = enhancer.enhance(CustomerInventory.class.getName(), customerInventoryCtClass.toBytecode());
+        CtClass enhancedCtClassCustomerInventory = cp.makeClass( new ByteArrayInputStream( enhancedCustomerInventory ) );
+        enhancedCtClassCustomerInventory.debugWriteFile("/tmp");
+        customerInventoryClass = enhancedCtClassCustomerInventory.toClass( cl, this.getClass().getProtectionDomain() );
+
+
+        CtClass supplierComponentPKCtClass = generateCtClassForAnEntity( SupplierComponentPK.class );
+        byte[] enhancedSupplierComponentPK = enhancer.enhanceComposite(SupplierComponentPK.class.getName(), supplierComponentPKCtClass.toBytecode());
+        CtClass enhancedCtClassSupplierComponentPK = cp.makeClass( new ByteArrayInputStream( enhancedSupplierComponentPK ) );
+        enhancedCtClassSupplierComponentPK.debugWriteFile("/tmp");
+        supplierComponentPKClass = enhancedCtClassSupplierComponentPK.toClass( cl, this.getClass().getProtectionDomain() );
+
+		Class entityClass = enhancedCtClass.toClass( cl, this.getClass().getProtectionDomain() );
+		Object entityInstance = entityClass.newInstance();
+
+		assertTyping( ManagedEntity.class, entityInstance );
+
+		// call the new methods
+		//
+		Method setter = entityClass.getMethod( EnhancerConstants.ENTITY_ENTRY_SETTER_NAME, EntityEntry.class );
+		Method getter = entityClass.getMethod( EnhancerConstants.ENTITY_ENTRY_GETTER_NAME );
+		assertNull( getter.invoke( entityInstance ) );
+		setter.invoke( entityInstance, makeEntityEntry() );
+		assertNotNull( getter.invoke( entityInstance ) );
+		setter.invoke( entityInstance, new Object[] {null} );
+		assertNull( getter.invoke( entityInstance ) );
+
+		Method entityInstanceGetter = entityClass.getMethod( EnhancerConstants.ENTITY_INSTANCE_GETTER_NAME );
+		assertSame( entityInstance, entityInstanceGetter.invoke( entityInstance ) );
+
+		Method previousGetter = entityClass.getMethod( EnhancerConstants.PREVIOUS_GETTER_NAME );
+		Method previousSetter = entityClass.getMethod( EnhancerConstants.PREVIOUS_SETTER_NAME, ManagedEntity.class );
+		previousSetter.invoke( entityInstance, entityInstance );
+		assertSame( entityInstance, previousGetter.invoke( entityInstance ) );
+
+		Method nextGetter = entityClass.getMethod( EnhancerConstants.PREVIOUS_GETTER_NAME );
+		Method nextSetter = entityClass.getMethod( EnhancerConstants.PREVIOUS_SETTER_NAME, ManagedEntity.class );
+		nextSetter.invoke( entityInstance, entityInstance );
+		assertSame( entityInstance, nextGetter.invoke( entityInstance ) );
+
+		// add an attribute interceptor...
+		Method interceptorGetter = entityClass.getMethod( EnhancerConstants.INTERCEPTOR_GETTER_NAME );
+		Method interceptorSetter = entityClass.getMethod( EnhancerConstants.INTERCEPTOR_SETTER_NAME, PersistentAttributeInterceptor.class );
+
+		assertNull( interceptorGetter.invoke( entityInstance ) );
+		entityClass.getMethod( "getId" ).invoke( entityInstance );
+
+		interceptorSetter.invoke( entityInstance, new LocalPersistentAttributeInterceptor() );
+		assertNotNull(interceptorGetter.invoke(entityInstance));
+
+		// dirty checking is unfortunately just printlns for now... just verify the test output
+		entityClass.getMethod( "getId" ).invoke( entityInstance );
+		entityClass.getMethod( "setId", Integer.class ).invoke( entityInstance, entityClass.getMethod( "getId" ).invoke( entityInstance ) );
+		entityClass.getMethod( "setId", Integer.class ).invoke( entityInstance, 1 );
+        assertTrue((Boolean) entityClass.getMethod("$$_hibernate_hasDirtyAttributes").invoke(entityInstance));
+
+		entityClass.getMethod( "setFirstName", String.class ).invoke(entityInstance, "Erik");
+
+		entityClass.getMethod( "setLastName", String.class ).invoke(entityInstance, "Mykland");
+        assertEquals(3, ((Set) entityClass.getMethod("$$_hibernate_getDirtyAttributes").invoke(entityInstance)).size());
+
+        entityClass.getMethod("$$_hibernate_clearDirtyAttributes").invoke(entityInstance);
+
+        //testing composite object
+        assert addressClass != null;
+        Object address =  addressClass.newInstance();
+
+        assert customerInventoryClass != null;
+        Object customerInventory =  customerInventoryClass.newInstance();
+
+        //Method adrInterceptorGetter = addressClass.getMethod( EnhancerConstants.INTERCEPTOR_GETTER_NAME );
+        //Method adrInterceptorSetter = addressClass.getMethod( EnhancerConstants.INTERCEPTOR_SETTER_NAME, PersistentAttributeInterceptor.class );
+        //adrInterceptorSetter.invoke( address, new LocalPersistentAttributeInterceptor() );
+
+        entityClass.getMethod("setAddress", addressClass).invoke(entityInstance, address);
+
+        addressClass.getMethod("setCity", String.class).invoke(address, "Arendal");
+
+        Set tracked = (Set) entityClass.getMethod(EnhancerConstants.TRACKER_GET_NAME).invoke(entityInstance);
+        assertEquals(2, tracked.size());
+        Iterator iter = tracked.iterator();
+        assertEquals("address", iter.next());
+        assertEquals("address.city", iter.next());
+
+        entityClass.getMethod("$$_hibernate_clearDirtyAttributes").invoke(entityInstance);
+
+        //make sure that new composite instances are cleared
+        Object address2 =  addressClass.newInstance();
+
+        entityClass.getMethod("setAddress", addressClass).invoke(entityInstance, address2);
+        addressClass.getMethod("setStreet1", String.class).invoke(address, "Heggedalveien");
+
+        tracked = (Set) entityClass.getMethod(EnhancerConstants.TRACKER_GET_NAME).invoke(entityInstance);
+        assertEquals(1, tracked.size());
+
+    }
+
+	private CtClass generateCtClassForAnEntity(Class entityClassToEnhance) throws Exception {
+		ClassPool cp = new ClassPool( false );
+		return cp.makeClass(
+				getClass().getClassLoader().getResourceAsStream(
+						entityClassToEnhance.getName().replace( '.', '/' ) + ".class"
+				)
+		);
+	}
+
+	private EntityEntry makeEntityEntry() {
+		return new EntityEntry(
+				Status.MANAGED,
+				null,
+				null,
+				new Long(1),
+				null,
+				LockMode.NONE,
+				false,
+				null,
+				EntityMode.POJO,
+				null,
+				false,
+				false,
+				null
+		);
+	}
+
+
+	private class LocalPersistentAttributeInterceptor implements PersistentAttributeInterceptor {
+		@Override
+		public boolean readBoolean(Object obj, String name, boolean oldValue) {
+			System.out.println( "Reading boolean [" + name + "]" );
+			return oldValue;
+		}
+
+		@Override
+		public boolean writeBoolean(Object obj, String name, boolean oldValue, boolean newValue) {
+			System.out.println( "Writing boolean [" + name + "]" );
+			return newValue;
+		}
+
+		@Override
+		public byte readByte(Object obj, String name, byte oldValue) {
+			System.out.println( "Reading byte [" + name + "]" );
+			return oldValue;
+		}
+
+		@Override
+		public byte writeByte(Object obj, String name, byte oldValue, byte newValue) {
+			System.out.println( "Writing byte [" + name + "]" );
+			return newValue;
+		}
+
+		@Override
+		public char readChar(Object obj, String name, char oldValue) {
+			System.out.println( "Reading char [" + name + "]" );
+			return oldValue;
+		}
+
+		@Override
+		public char writeChar(Object obj, String name, char oldValue, char newValue) {
+			System.out.println( "Writing char [" + name + "]" );
+			return newValue;
+		}
+
+		@Override
+		public short readShort(Object obj, String name, short oldValue) {
+			System.out.println( "Reading short [" + name + "]" );
+			return oldValue;
+		}
+
+		@Override
+		public short writeShort(Object obj, String name, short oldValue, short newValue) {
+			System.out.println( "Writing short [" + name + "]" );
+			return newValue;
+		}
+
+		@Override
+		public int readInt(Object obj, String name, int oldValue) {
+			System.out.println( "Reading int [" + name + "]" );
+			return oldValue;
+		}
+
+		@Override
+		public int writeInt(Object obj, String name, int oldValue, int newValue) {
+			System.out.println( "Writing int [" + name + "]" );
+			return newValue;
+		}
+
+		@Override
+		public float readFloat(Object obj, String name, float oldValue) {
+			System.out.println( "Reading float [" + name + "]" );
+			return oldValue;
+		}
+
+		@Override
+		public float writeFloat(Object obj, String name, float oldValue, float newValue) {
+			System.out.println( "Writing float [" + name + "]" );
+			return newValue;
+		}
+
+		@Override
+		public double readDouble(Object obj, String name, double oldValue) {
+			System.out.println( "Reading double [" + name + "]" );
+			return oldValue;
+		}
+
+		@Override
+		public double writeDouble(Object obj, String name, double oldValue, double newValue) {
+			System.out.println( "Writing double [" + name + "]" );
+			return newValue;
+		}
+
+		@Override
+		public long readLong(Object obj, String name, long oldValue) {
+			System.out.println( "Reading long [" + name + "]" );
+			return oldValue;
+		}
+
+		@Override
+		public long writeLong(Object obj, String name, long oldValue, long newValue) {
+			System.out.println( "Writing long [" + name + "]" );
+			return newValue;
+		}
+
+		@Override
+		public Object readObject(Object obj, String name, Object oldValue) {
+			System.out.println( "Reading Object [" + name + "]" );
+			return oldValue;
+		}
+
+		@Override
+		public Object writeObject(Object obj, String name, Object oldValue, Object newValue) {
+			System.out.println( "Writing Object [" + name + "]" );
+			return newValue;
+		}
+	}
+}

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/customer/CustomerInventory.java

@@ -0,0 +1,152 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */package org.hibernate.test.bytecode.enhancement.customer;
+
+/**
+ * @author <a href="mailto:stale.pedersen@jboss.org">Ståle W. Pedersen</a>
+ */
+
+import javax.persistence.CascadeType;
+import javax.persistence.Column;
+import javax.persistence.Entity;
+import javax.persistence.GeneratedValue;
+import javax.persistence.GenerationType;
+import javax.persistence.Id;
+import javax.persistence.IdClass;
+import javax.persistence.JoinColumn;
+import javax.persistence.ManyToOne;
+import javax.persistence.Table;
+import javax.persistence.TableGenerator;
+import javax.persistence.Version;
+import java.io.Serializable;
+import java.math.BigDecimal;
+import java.util.Comparator;
+
+@SuppressWarnings("serial")
+@Entity
+@Table(name="O_CUSTINVENTORY")
+@IdClass(CustomerInventoryPK.class)
+public class CustomerInventory implements Serializable, Comparator<CustomerInventory> {
+
+        public static final String QUERY_COUNT = "CustomerInventory.count";
+
+        @Id
+        @TableGenerator(name="inventory",
+            table="U_SEQUENCES",
+            pkColumnName="S_ID",
+            valueColumnName="S_NEXTNUM",
+            pkColumnValue="inventory",
+            allocationSize=1000)
+    @GeneratedValue(strategy= GenerationType.TABLE,generator="inventory")
+    @Column(name="CI_ID")
+    private Long         id;
+
+    @Id
+    @Column(name = "CI_CUSTOMERID", insertable = false, updatable = false)
+    private int             custId;
+
+    @ManyToOne(cascade= CascadeType.MERGE)
+    @JoinColumn(name="CI_CUSTOMERID")
+    private Customer        customer;
+
+    @ManyToOne(cascade=CascadeType.MERGE)
+    @JoinColumn(name = "CI_ITEMID")
+    private String            vehicle;
+
+    @Column(name="CI_VALUE")
+    private BigDecimal totalCost;
+
+    @Column(name="CI_QUANTITY")
+    private int             quantity;
+
+    @Version
+    @Column(name = "CI_VERSION")
+    private int             version;
+
+    public CustomerInventory() {
+    }
+
+        CustomerInventory(Customer customer, String vehicle, int quantity,
+                      BigDecimal totalValue) {
+        this.customer = customer;
+        this.vehicle = vehicle;
+        this.quantity = quantity;
+        this.totalCost = totalValue;
+    }
+
+    public String getVehicle() {
+        return vehicle;
+    }
+
+    public BigDecimal getTotalCost() {
+        return totalCost;
+    }
+
+    public int getQuantity() {
+        return quantity;
+    }
+
+    public Long getId() {
+        return id;
+    }
+
+    public Customer getCustomer() {
+        return customer;
+    }
+
+    public int getCustId() {
+        return custId;
+    }
+
+    public int getVersion() {
+        return version;
+    }
+
+    public int compare(CustomerInventory cdb1, CustomerInventory cdb2) {
+        return cdb1.id.compareTo(cdb2.id);
+    }
+
+            @Override
+    public boolean equals(Object obj) {
+        if (obj == this)
+            return true;
+        if (obj == null || !(obj instanceof CustomerInventory))
+            return false;
+        if (this.id == ((CustomerInventory)obj).id)
+            return true;
+        if (this.id != null && ((CustomerInventory)obj).id == null)
+            return false;
+        if (this.id == null && ((CustomerInventory)obj).id != null)
+            return false;
+
+        return this.id.equals(((CustomerInventory)obj).id);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = id.hashCode();
+        result = 31 * result + custId;
+        return result;
+    }
+
+}

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/customer/CustomerInventoryPK.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2012 Red Hat, Inc. and/or its affiliates.
+ *
+ * Licensed under the Eclipse Public License version 1.0, available at
+ * http://www.eclipse.org/legal/epl-v10.html
+ */
+package org.hibernate.test.bytecode.enhancement.customer;
+
+import java.io.Serializable;
+
+/**
+ * @author <a href="mailto:stale.pedersen@jboss.org">Ståle W. Pedersen</a>
+ */
+public class CustomerInventoryPK implements Serializable {
+
+    private Long id;
+    private int custId;
+
+    public CustomerInventoryPK() {
+    }
+
+    public CustomerInventoryPK(Long id, int custId) {
+        this.id = id;
+        this.custId = custId;
+    }
+
+    public boolean equals(Object other) {
+        if (other == this) {
+            return true;
+        }
+        if (other == null || getClass() != other.getClass()) {
+            return false;
+        }
+        CustomerInventoryPK cip = (CustomerInventoryPK) other;
+        return (custId == cip.custId && (id == cip.id ||
+                ( id != null && id.equals(cip.id))));
+    }
+
+    public int hashCode() {
+        return (id == null ? 0 : id.hashCode()) ^ custId;
+    }
+
+    public Long getId() {
+        return id;
+    }
+
+    public int getCustId() {
+        return custId;
+    }
+}

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/customer/SupplierComponentPK.java

@@ -0,0 +1,71 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.test.bytecode.enhancement.customer;
+
+import javax.persistence.Embeddable;
+
+/**
+ * @author <a href="mailto:stale.pedersen@jboss.org">Ståle W. Pedersen</a>
+ */
+@Embeddable
+public class SupplierComponentPK {
+
+    String componentID;
+    int supplierID;
+
+    public SupplierComponentPK() {
+    }
+
+    public SupplierComponentPK(String suppCompID, int suppCompSuppID) {
+        this.componentID = suppCompID;
+        this.supplierID = suppCompSuppID;
+    }
+
+    public String getComponentID() {
+        return componentID;
+    }
+
+    public int getSupplierID() {
+        return supplierID;
+    }
+
+    @Override
+    public int hashCode() {
+        final int PRIME = 31;
+        int result = 1;
+        result = PRIME * result + componentID.hashCode();
+        result = PRIME * result + supplierID;
+        return result;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj)
+            return true;
+        if (obj == null || getClass() != obj.getClass())
+            return false;
+        final SupplierComponentPK other = (SupplierComponentPK) obj;
+        return componentID.equals(other.componentID);
+    }
+}

hibernate-core/src/test/java/org/hibernate/test/cfg/persister/GoofyPersisterClassProvider.java

@@ -25,6 +25,7 @@ import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.util.Comparator;
 import java.util.Map;
+import java.util.Set;
 
 import org.hibernate.EntityMode;
 import org.hibernate.HibernateException;
@@ -603,6 +604,11 @@ public class GoofyPersisterClassProvider implements PersisterClassResolver {
 		public Iterable<AttributeDefinition> getAttributes() {
 			throw new NotYetImplementedException();
 		}
+
+        @Override
+        public int[] resolveAttributeIndexes(Set<String> attributes) {
+            return null;
+        }
 	}
 
 	public static class NoopCollectionPersister implements CollectionPersister {

hibernate-core/src/test/java/org/hibernate/test/legacy/CustomPersister.java

@@ -4,6 +4,7 @@ import java.io.Serializable;
 import java.util.Comparator;
 import java.util.Hashtable;
 import java.util.Map;
+import java.util.Set;
 
 import org.hibernate.EntityMode;
 import org.hibernate.HibernateException;
@@ -693,4 +694,9 @@ public class CustomPersister implements EntityPersister {
 	public Iterable<AttributeDefinition> getAttributes() {
 		throw new NotYetImplementedException();
 	}
+
+    @Override
+    public int[] resolveAttributeIndexes(Set<String> attributes) {
+        return null;
+    }
 }

hibernate-gradle-plugin/src/main/groovy/org/hibernate/tooling/gradle/EnhancerTask.groovy

@@ -32,7 +32,10 @@ import org.gradle.api.tasks.TaskAction
 import org.hibernate.bytecode.enhance.spi.EnhancementContext
 import org.hibernate.bytecode.enhance.spi.Enhancer
 
+import javax.persistence.ElementCollection
 import javax.persistence.Entity
+import javax.persistence.ManyToMany
+import javax.persistence.OneToMany
 import javax.persistence.Transient
 
 /**
@@ -142,7 +145,7 @@ public class EnhancerTask extends DefaultTask implements EnhancementContext {
     }
 
     public boolean doDirtyCheckingInline(CtClass classDescriptor) {
-        return false;
+        return true;
     }
 
     public CtField[] order(CtField[] fields) {
@@ -150,6 +153,17 @@ public class EnhancerTask extends DefaultTask implements EnhancementContext {
         return fields;
     }
 
+    public boolean isMappedCollection(CtField field) {
+        try {
+            return (field.getAnnotation(OneToMany.class) != null ||
+                    field.getAnnotation(ManyToMany.class) != null ||
+                    field.getAnnotation(ElementCollection.class) != null);
+        }
+        catch (ClassNotFoundException e) {
+            return false;
+        }
+    }
+
     public boolean isPersistentField(CtField ctField) {
         return !ctField.hasAnnotation( Transient.class );
     }

hibernate-maven-plugin/src/main/java/org/hibernate/bytecode/enhance/plugins/HibernateEnhancementMojo.java

@@ -37,7 +37,11 @@ import javassist.ClassPool;
 import javassist.CtClass;
 import javassist.CtField;
 
+import javax.persistence.Embeddable;
 import javax.persistence.Entity;
+import javax.persistence.ElementCollection;
+import javax.persistence.ManyToMany;
+import javax.persistence.OneToMany;
 import javax.persistence.Transient;
 
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
@@ -55,13 +59,14 @@ import org.apache.maven.plugins.annotations.Parameter;
  * @author Jeremy Whiting
  */
 @Mojo(name = "enhance")
-public class HibernateEnhancementMojo extends AbstractMojo {
+public class HibernateEnhancementMojo extends AbstractMojo implements EnhancementContext {
 
 	/**
 	 * The contexts to use during enhancement.
 	 */
 	private List<File> classes = new ArrayList<File>();
 	private ClassPool pool = new ClassPool( false );
+    private final Enhancer enhancer = new Enhancer( this);
 
 	private static final String CLASS_EXTENSION = ".class";
 
@@ -74,57 +79,9 @@ public class HibernateEnhancementMojo extends AbstractMojo {
 		File root = new File( this.dir ); 
 		walkDir( root );
 
-		Enhancer enhancer = new Enhancer( new EnhancementContext() {
-
-			private ClassLoader overridden;
-
-			public ClassLoader getLoadingClassLoader() {
-				if ( null == this.overridden ) {
-					return getClass().getClassLoader();
-				}
-				else {
-					return this.overridden;
-				}
-			}
-
-			public void setClassLoader(ClassLoader loader) {
-				this.overridden = loader;
-			}
-
-			public boolean isEntityClass(CtClass classDescriptor) {
-				return true;
-			}
-
-			public boolean hasLazyLoadableAttributes(CtClass classDescriptor) {
-				return true;
-			}
-
-			public boolean isLazyLoadable(CtField field) {
-				return true;
-			}
-
-			public boolean isCompositeClass(CtClass classDescriptor) {
-				return false;
-			}
-
-			public boolean doDirtyCheckingInline(CtClass classDescriptor) {
-				return false;
-			}
-
-			public CtField[] order(CtField[] fields) {
-				// TODO: load ordering from configuration.
-				return fields;
-			}
-
-			public boolean isPersistentField(CtField ctField) {
-				return !ctField.hasAnnotation( Transient.class );
-			}
-
-		} );
-
 		if ( 0 < classes.size() ) {
 			for ( File file : classes ) {
-				enhanceClass( enhancer, file );
+				processClassFile(file);
 			}
 		}
 
@@ -152,6 +109,7 @@ public class HibernateEnhancementMojo extends AbstractMojo {
 	}
 
 	private void walkDir(File dir, FileFilter classesFilter, FileFilter dirFilter) {
+
 		File[] dirs = dir.listFiles( dirFilter );
 		for ( int i = 0; i < dirs.length; i++ ) {
 			walkDir( dirs[i], classesFilter, dirFilter );
@@ -163,67 +121,82 @@ public class HibernateEnhancementMojo extends AbstractMojo {
 		}
 	}
 
-	private void enhanceClass(Enhancer enhancer, File file) {
+
-		byte[] enhancedBytecode = null;
+    /**
-		InputStream is = null;
+     * Atm only process files annotated with either @Entity or @Embeddable
-		CtClass clas = null;
+     * @param javaClassFile
+     */
+    private void processClassFile(File javaClassFile)
+            throws MojoExecutionException {
 		try {
-			is = new FileInputStream( file.toString() );
+			final CtClass ctClass = getClassPool().makeClass( new FileInputStream( javaClassFile ) );
-			clas = getClassPool().makeClass( is );
+            if(this.isEntityClass(ctClass))
-			if ( !clas.hasAnnotation( Entity.class ) ) {
+                processEntityClassFile(javaClassFile, ctClass);
-				getLog().debug( "Class $file not an annotated Entity class. skipping..." );
+            else if(this.isCompositeClass(ctClass))
-			}
+                processCompositeClassFile(javaClassFile, ctClass);
-			else {
+
-				enhancedBytecode = enhancer.enhance( clas.getName(), clas.toBytecode() );
-			}
-		}
-		catch (Exception e) {
-			getLog().error( "Unable to enhance class [${file.toString()}]", e );
-			return;
-		}
-		finally {
-			try {
-				if ( null != is )
-					is.close();
-			}
-			catch (IOException ioe) {}
-		}
-		if ( null != enhancedBytecode ) {
-			if ( file.delete() ) {
-				try {
-					if ( !file.createNewFile() ) {
-						getLog().error( "Unable to recreate class file [" + clas.getName() + "]" );
-					}
-				}
-				catch (IOException ioe) {
-				}
-			}
-			else {
-				getLog().error( "Unable to delete class file [" + clas.getName() + "]" );
-			}
-			FileOutputStream outputStream = null;
-			try {
-				outputStream = new FileOutputStream( file, false );
-				outputStream.write( enhancedBytecode );
-				outputStream.flush();
-			}
-			catch (IOException ioe) {
-			}
-			finally {
-				try {
-					if ( outputStream != null )
-						outputStream.close();
-					clas.detach();// release memory
-				}
-				catch (IOException ignore) {
-				}
         }
+        catch (IOException e) {
+            throw new MojoExecutionException(
+                    String.format( "Error processing included file [%s]", javaClassFile.getAbsolutePath() ), e );
         }
     }
 
-	public void setDir(String dir) {
+    private void processEntityClassFile(File javaClassFile, CtClass ctClass ) {
-		if ( null != dir && !"".equals( dir.trim() ) ) {
+        try {
-			this.dir = dir;
+            byte[] result = enhancer.enhance( ctClass.getName(), ctClass.toBytecode() );
+            if(result != null)
+                writeEnhancedClass(javaClassFile, result);
+        }
+        catch (Exception e) {
+            getLog().error( "Unable to enhance class [" + ctClass.getName() + "]", e);
+            return;
+        }
+    }
+
+    private void processCompositeClassFile(File javaClassFile, CtClass ctClass) {
+        try {
+            byte[] result = enhancer.enhanceComposite(ctClass.getName(), ctClass.toBytecode());
+            if(result != null)
+                writeEnhancedClass(javaClassFile, result);
+        }
+        catch (Exception e) {
+            getLog().error( "Unable to enhance class [" + ctClass.getName() + "]", e);
+            return;
+        }
+    }
+
+    private void writeEnhancedClass(File javaClassFile, byte[] result)
+            throws MojoExecutionException {
+        try {
+			if ( javaClassFile.delete() ) {
+                    if ( ! javaClassFile.createNewFile() ) {
+                        getLog().error( "Unable to recreate class file [" + javaClassFile.getName() + "]");
+                    }
+            }
+			else {
+				getLog().error( "Unable to delete class file [" + javaClassFile.getName() + "]");
+			}
+
+			FileOutputStream outputStream = new FileOutputStream( javaClassFile, false );
+			try {
+				outputStream.write( result);
+				outputStream.flush();
+			}
+			finally {
+				try {
+					outputStream.close();
+				}
+				catch ( IOException ignore) {
+				}
+			}
+        }
+        catch (FileNotFoundException ignore) {
+            // should not ever happen because of explicit checks
+        }
+        catch (IOException e) {
+            throw new MojoExecutionException(
+                    String.format( "Error processing included file [%s]", javaClassFile.getAbsolutePath() ), e );
         }
     }
 
@@ -231,4 +204,63 @@ public class HibernateEnhancementMojo extends AbstractMojo {
 		return this.pool;
 	}
 
+    private boolean shouldInclude(CtClass ctClass) {
+		// we currently only handle entity enhancement
+		return ctClass.hasAnnotation( Entity.class );
+	}
+
+	@Override
+	public ClassLoader getLoadingClassLoader() {
+		return getClass().getClassLoader();
+	}
+
+	@Override
+	public boolean isEntityClass(CtClass classDescriptor) {
+        return classDescriptor.hasAnnotation(Entity.class);
+    }
+
+	@Override
+	public boolean isCompositeClass(CtClass classDescriptor) {
+        return classDescriptor.hasAnnotation(Embeddable.class);
+	}
+
+	@Override
+	public boolean doDirtyCheckingInline(CtClass classDescriptor) {
+		return true;
+	}
+
+	@Override
+	public boolean hasLazyLoadableAttributes(CtClass classDescriptor) {
+		return true;
+	}
+
+	@Override
+	public boolean isLazyLoadable(CtField field) {
+		return true;
+	}
+
+	@Override
+	public boolean isPersistentField(CtField ctField) {
+		// current check is to look for @Transient
+		return ! ctField.hasAnnotation( Transient.class );
+	}
+
+    @Override
+    public boolean isMappedCollection(CtField field) {
+        try {
+            return (field.getAnnotation(OneToMany.class) != null ||
+                    field.getAnnotation(ManyToMany.class) != null ||
+                    field.getAnnotation(ElementCollection.class) != null);
+        }
+        catch (ClassNotFoundException e) {
+            return false;
+        }
+    }
+
+	@Override
+	public CtField[] order(CtField[] persistentFields) {
+		// for now...
+		return persistentFields;
+		// eventually needs to consult the Hibernate metamodel for proper ordering
+	}
 }