support ref tracking for map chunk encoding

java/fury-core/src/main/java/org/apache/fury/builder/BaseObjectCodecBuilder.java

@@ -25,7 +25,6 @@
 import static org.apache.fury.codegen.Expression.Reference.fieldRef;
 import static org.apache.fury.codegen.ExpressionOptimizer.invokeGenerated;
 import static org.apache.fury.codegen.ExpressionUtils.add;
-import static org.apache.fury.codegen.ExpressionUtils.and;
 import static org.apache.fury.codegen.ExpressionUtils.bitand;
 import static org.apache.fury.codegen.ExpressionUtils.bitor;
 import static org.apache.fury.codegen.ExpressionUtils.cast;
@@ -506,7 +505,7 @@ protected Expression writeClassInfo(
     Expression classInfo = classInfoRef.f0;
     writeClassAction.add(
         new If(
-            neq(new Invoke(classInfo, "getCls", CLASS_TYPE), clsExpr),
+            neq(inlineInvoke(classInfo, "getCls", CLASS_TYPE), clsExpr),
             new Assign(
                 classInfo,
                 inlineInvoke(classResolverRef, "getClassInfo", classInfoTypeRef, clsExpr))));
@@ -1149,10 +1148,18 @@ protected Expression writeChunk(
     Expression chunkSizeOffset =
         subtract(
             inlineInvoke(buffer, "writerIndex", PRIMITIVE_INT_TYPE), ofInt(1), "chunkSizeOffset");
+    expressions.add(
+        key,
+        value,
+        keyTypeExpr,
+        valueTypeExpr,
+        writePlaceHolder,
+        chunkSizeOffset,
+        writePlaceHolder,
+        chunkSizeOffset);
 
     Expression chunkHeader;
     Expression keySerializer, valueSerializer;
-
     boolean trackingKeyRef =
         visitFury(fury -> fury.getClassResolver().needToWriteRef(keyTypeRawType));
     boolean trackingValueRef =
@@ -1171,60 +1178,75 @@ protected Expression writeChunk(
         header |= TRACKING_VALUE_REF;
       }
       chunkHeader = ofInt(header);
+      expressions.add(chunkHeader);
     } else if (keyMonomorphic) {
       int header = KEY_DECL_TYPE;
       if (trackingKeyRef) {
         header |= TRACKING_KEY_REF;
       }
       keySerializer = getOrCreateSerializer(keyTypeRawType);
+      walkPath.add("value:" + valueType);
       valueSerializer = writeClassInfo(buffer, valueTypeExpr, valueTypeRawType, true);
-      chunkHeader = ofInt(header);
+      walkPath.removeLast();
+      chunkHeader = ExpressionUtils.ofInt("chunkHeader", header);
+      expressions.add(chunkHeader);
       if (trackingValueRef) {
         // value type may be subclass and not track ref.
         valueWriteRef =
-            new Invoke(valueSerializer, "needToWriteRef", PRIMITIVE_BOOLEAN_TYPE, valueTypeExpr);
-        chunkHeader = and(chunkHeader, valueWriteRef, "chunkHeader");
+            new Invoke(valueSerializer, "needToWriteRef", "valueWriteRef", PRIMITIVE_BOOLEAN_TYPE);
+        expressions.add(
+            new If(
+                valueWriteRef,
+                new Assign(chunkHeader, bitor(chunkHeader, ofInt(TRACKING_VALUE_REF)))));
       }
     } else if (valueMonomorphic) {
+      walkPath.add("key:" + keyType);
       keySerializer = writeClassInfo(buffer, keyTypeExpr, keyTypeRawType, true);
+      walkPath.removeLast();
       valueSerializer = getOrCreateSerializer(valueTypeRawType);
       int header = VALUE_DECL_TYPE;
       if (trackingValueRef) {
         header |= TRACKING_VALUE_REF;
       }
-      chunkHeader = ofInt(header);
+      chunkHeader = ExpressionUtils.ofInt("chunkHeader", header);
+      expressions.add(chunkHeader);
       if (trackingKeyRef) {
         // key type may be subclass and not track ref.
         keyWriteRef =
-            new Invoke(keySerializer, "needToWriteRef", PRIMITIVE_BOOLEAN_TYPE, keyTypeExpr);
-        chunkHeader = and(chunkHeader, keyWriteRef, "chunkHeader");
+            new Invoke(keySerializer, "needToWriteRef", "keyWriteRef", PRIMITIVE_BOOLEAN_TYPE);
+        expressions.add(
+            new If(
+                keyWriteRef, new Assign(chunkHeader, bitor(chunkHeader, ofInt(TRACKING_KEY_REF)))));
       }
     } else {
+      walkPath.add("key:" + keyType);
       keySerializer = writeClassInfo(buffer, keyTypeExpr, keyTypeRawType, true);
+      walkPath.removeLast();
+      walkPath.add("value:" + valueType);
       valueSerializer = writeClassInfo(buffer, valueTypeExpr, valueTypeRawType, true);
-      chunkHeader = ofInt(0);
+      walkPath.removeLast();
+      chunkHeader = ExpressionUtils.ofInt("chunkHeader", 0);
+      expressions.add(chunkHeader);
       if (trackingKeyRef) {
         // key type may be subclass and not track ref.
-        valueWriteRef =
-            new Invoke(valueSerializer, "needToWriteRef", PRIMITIVE_BOOLEAN_TYPE, valueTypeExpr);
-        chunkHeader = and(chunkHeader, valueWriteRef, "chunkHeader");
+        keyWriteRef =
+            new Invoke(keySerializer, "needToWriteRef", "keyWriteRef", PRIMITIVE_BOOLEAN_TYPE);
+        expressions.add(
+            new If(
+                keyWriteRef, new Assign(chunkHeader, bitor(chunkHeader, ofInt(TRACKING_KEY_REF)))));
       }
       if (trackingValueRef) {
         // key type may be subclass and not track ref.
-        keyWriteRef =
-            new Invoke(keySerializer, "needToWriteRef", PRIMITIVE_BOOLEAN_TYPE, keyTypeExpr);
-        chunkHeader = and(chunkHeader, keyWriteRef, "chunkHeader");
+        valueWriteRef =
+            new Invoke(valueSerializer, "needToWriteRef", "valueWriteRef", PRIMITIVE_BOOLEAN_TYPE);
+        expressions.add(
+            new If(
+                valueWriteRef,
+                new Assign(chunkHeader, bitor(chunkHeader, ofInt(TRACKING_VALUE_REF)))));
       }
     }
     Expression chunkSize = ExpressionUtils.ofInt("chunkSize", 0);
     expressions.add(
-        key,
-        value,
-        keyTypeExpr,
-        valueTypeExpr,
-        writePlaceHolder,
-        chunkSizeOffset,
-        chunkHeader,
         keySerializer,
         valueSerializer,
         keyWriteRef,
@@ -1447,8 +1469,10 @@ protected Expression deserializeForNotNull(
       } else if (useMapSerialization(typeRef)) {
         obj = deserializeForMap(buffer, typeRef, serializer, invokeHint);
       } else {
+        if (serializer != null) {
+          return new Invoke(serializer, "read", OBJECT_TYPE, buffer);
+        }
         if (isMonomorphic(cls)) {
-          Preconditions.checkState(serializer == null);
           serializer = getOrCreateSerializer(cls);
           Class<?> returnType =
               ReflectionUtils.getReturnType(getRawType(serializer.type()), "read");
@@ -1812,43 +1836,29 @@ private Expression readChunk(
           new If(
               keyIsDeclaredType,
               keySerializer,
-              new Invoke(
-                  readClassInfo(keyTypeRawType, buffer),
-                  "getSerializer",
-                  "keySerializer",
-                  SERIALIZER_TYPE),
+              inlineInvoke(readClassInfo(keyTypeRawType, buffer), "getSerializer", SERIALIZER_TYPE),
               false);
       valueSerializer =
           new If(
               valueIsDeclaredType,
               valueSerializer,
-              new Invoke(
-                  readClassInfo(valueTypeRawType, buffer),
-                  "getSerializer",
-                  "valueSerializer",
-                  SERIALIZER_TYPE),
+              inlineInvoke(
+                  readClassInfo(valueTypeRawType, buffer), "getSerializer", SERIALIZER_TYPE),
               false);
     } else if (!keyMonomorphic) {
       keySerializer =
           new If(
               keyIsDeclaredType,
               keySerializer,
-              new Invoke(
-                  readClassInfo(keyTypeRawType, buffer),
-                  "getSerializer",
-                  "keySerializer",
-                  SERIALIZER_TYPE),
+              inlineInvoke(readClassInfo(keyTypeRawType, buffer), "getSerializer", SERIALIZER_TYPE),
               false);
     } else if (!valueMonomorphic) {
       valueSerializer =
           new If(
               valueIsDeclaredType,
               valueSerializer,
-              new Invoke(
-                  readClassInfo(valueTypeRawType, buffer),
-                  "getSerializer",
-                  "valueSerializer",
-                  SERIALIZER_TYPE),
+              inlineInvoke(
+                  readClassInfo(valueTypeRawType, buffer), "getSerializer", SERIALIZER_TYPE),
               false);
     }
     Expression keySerializerExpr = uninline(keySerializer);
@@ -1877,7 +1887,8 @@ private Expression readChunk(
                             expr -> expr,
                             () ->
                                 deserializeForNotNull(buffer, keyType, keySerializerExpr, keyHint)),
-                        deserializeForNotNull(buffer, keyType, keySerializerExpr, keyHint));
+                        deserializeForNotNull(buffer, keyType, keySerializerExpr, keyHint),
+                        false);
               } else {
                 keyAction = deserializeForNotNull(buffer, keyType, keySerializerExpr, keyHint);
               }
@@ -1893,7 +1904,8 @@ private Expression readChunk(
                             () ->
                                 deserializeForNotNull(
                                     buffer, valueType, valueSerializerExpr, valueHint)),
-                        deserializeForNotNull(buffer, valueType, valueSerializerExpr, valueHint));
+                        deserializeForNotNull(buffer, valueType, valueSerializerExpr, valueHint),
+                        false);
               } else {
                 valueAction =
                     deserializeForNotNull(buffer, valueType, valueSerializerExpr, valueHint);
@@ -1913,14 +1925,15 @@ private Expression readChunk(
                   chunkHeader, inlineInvoke(buffer, "readUnsignedByte", PRIMITIVE_INT_TYPE))));
       return expressions;
     } else {
-      Expression sizeAndHeader =
+      Expression returnSizeAndHeader =
           new If(
               gt(size, ofInt(0)),
-              (bitor(
-                  shift("<<", size, 8),
-                  inlineInvoke(buffer, "readUnsignedByte", PRIMITIVE_INT_TYPE))),
-              ofInt(0));
-      expressions.add(new Return(sizeAndHeader));
+              new Return(
+                  (bitor(
+                      shift("<<", size, 8),
+                      inlineInvoke(buffer, "readUnsignedByte", PRIMITIVE_INT_TYPE)))),
+              new Return(ofInt(0)));
+      expressions.add(returnSizeAndHeader);
       // method too big, spilt it into a new method.
       // Generate similar signature as `AbstractMapSerializer.writeJavaChunk`(
       //   MemoryBuffer buffer,

java/fury-core/src/main/java/org/apache/fury/codegen/Expression.java

@@ -2438,7 +2438,6 @@ public ForLoop(
       Preconditions.checkArgument(maxType.isPrimitive());
       iref = new Reference(String.valueOf(System.identityHashCode(this)), TypeRef.of(maxType));
       this.loopAction = action.apply(iref);
-      ;
     }
 
     @Override

java/fury-core/src/main/java/org/apache/fury/codegen/ExpressionUtils.java

@@ -103,6 +103,12 @@ public static BitOr bitor(Expression left, Expression right) {
     return new BitOr(left, right);
   }
 
+  public static BitAnd bitand(Expression left, Expression right, String name) {
+    BitAnd bitAnd = new BitAnd(left, right);
+    bitAnd.inline(false);
+    return bitAnd;
+  }
+
   public static BitAnd bitand(Expression left, Expression right) {
     return new BitAnd(left, right);
   }

java/fury-core/src/test/java/org/apache/fury/serializer/collection/MapSerializersTest.java

@@ -53,6 +53,7 @@
 import org.apache.fury.reflect.TypeRef;
 import org.apache.fury.serializer.Serializer;
 import org.apache.fury.serializer.collection.CollectionSerializersTest.TestEnum;
+import org.apache.fury.test.bean.BeanB;
 import org.apache.fury.test.bean.Cyclic;
 import org.apache.fury.test.bean.MapFields;
 import org.apache.fury.type.GenericType;
@@ -787,4 +788,27 @@ public void testMapFieldStructCodegen2(boolean referenceTrackingConfig) {
     struct1.map3 = ofHashMap(1, "v1", 2, "v2");
     serDeCheck(fury, struct1);
   }
+
+  @Data
+  public static class MapFieldStruct3 {
+    public Map<Object, Object> map1;
+    public Map<BeanB, Object> map2;
+    public Map<Object, BeanB> map3;
+  }
+
+  @Test(dataProvider = "referenceTrackingConfig")
+  public void testMapFieldStructCodegen3(boolean referenceTrackingConfig) {
+    Fury fury =
+        Fury.builder()
+            .withRefTracking(referenceTrackingConfig)
+            .withCodegen(true)
+            .requireClassRegistration(false)
+            .build();
+    MapFieldStruct3 struct1 = new MapFieldStruct3();
+    BeanB beanB = BeanB.createBeanB(2);
+    struct1.map1 = ofHashMap(BeanB.createBeanB(2), BeanB.createBeanB(2));
+    struct1.map2 = ofHashMap(BeanB.createBeanB(2), 1, beanB, beanB);
+    struct1.map3 = ofHashMap(1, beanB, 2, beanB, 3, BeanB.createBeanB(2));
+    serDeCheck(fury, struct1);
+  }
 }