finish chunk writing codegen

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

@@ -25,11 +25,14 @@
 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.cast;
 import static org.apache.fury.codegen.ExpressionUtils.eq;
+import static org.apache.fury.codegen.ExpressionUtils.eqNull;
 import static org.apache.fury.codegen.ExpressionUtils.gt;
 import static org.apache.fury.codegen.ExpressionUtils.inline;
-import static org.apache.fury.codegen.ExpressionUtils.isNull;
+import static org.apache.fury.codegen.ExpressionUtils.list;
 import static org.apache.fury.codegen.ExpressionUtils.neq;
+import static org.apache.fury.codegen.ExpressionUtils.neqNull;
 import static org.apache.fury.codegen.ExpressionUtils.not;
 import static org.apache.fury.codegen.ExpressionUtils.notNull;
 import static org.apache.fury.codegen.ExpressionUtils.nullValue;
@@ -349,7 +352,7 @@ protected Expression serializeFor(
                   buffer, "writeByte", new Literal(Fury.REF_VALUE_FLAG, PRIMITIVE_BYTE_TYPE)),
               serializeForNotNull(inputObject, buffer, typeRef, serializer, generateNewMethod));
       return new If(
-          ExpressionUtils.eqNull(inputObject),
+          eqNull(inputObject),
           new Invoke(buffer, "writeByte", new Literal(Fury.NULL_FLAG, PRIMITIVE_BYTE_TYPE)),
           action);
     }
@@ -1050,9 +1053,16 @@ private Expression jitWriteMap(
     map = new Invoke(serializer, "onMapWrite", TypeUtils.mapOf(keyType, valueType), buffer, map);
     Expression iterator =
         new Invoke(inlineInvoke(map, "entrySet", SET_TYPE), "iterator", ITERATOR_TYPE);
-    Expression entry = new Invoke(iterator, "next", MAP_ENTRY_TYPE);
+    Expression entry =
+        new Cast(inlineInvoke(iterator, "next", OBJECT_TYPE), MAP_ENTRY_TYPE, "entry");
     boolean keyMonomorphic = isMonomorphic(keyType);
     boolean valueMonomorphic = isMonomorphic(valueType);
+    Class<?> keyTypeRawType = keyType.getRawType();
+    Class<?> valueTypeRawType = valueType.getRawType();
+    boolean trackingKeyRef =
+        visitFury(fury -> fury.getClassResolver().needToWriteRef(keyTypeRawType));
+    boolean trackingValueRef =
+        visitFury(fury -> fury.getClassResolver().needToWriteRef(valueTypeRawType));
     Expression keySerializer, valueSerializer;
     if (keyMonomorphic && valueMonomorphic) {
       keySerializer = getOrCreateSerializer(keyType.getRawType());
@@ -1069,22 +1079,28 @@ private Expression jitWriteMap(
     }
     Expression.While whileAction =
         new Expression.While(
-            notNull(entry),
-            () ->
-                new ListExpression(
-                    new Assign(
-                        inlineInvoke(
-                            serializer,
-                            "writeJavaNullChunk",
-                            MAP_ENTRY_TYPE,
-                            buffer,
-                            entry,
-                            iterator,
-                            keySerializer,
-                            valueSerializer),
-                        entry),
-                    new If(
-                        notNull(entry), writeChunk(buffer, entry, iterator, keyType, valueType))));
+            neqNull(entry),
+            () -> {
+              String method = "writeJavaNullChunk";
+              if (keyMonomorphic && valueMonomorphic) {
+                if (!trackingKeyRef && !trackingValueRef) {
+                  method = "writeNullChunkKVNoRef";
+                }
+              }
+              return new ListExpression(
+                  new Assign(
+                      entry,
+                      inlineInvoke(
+                          serializer,
+                          method,
+                          MAP_ENTRY_TYPE,
+                          buffer,
+                          entry,
+                          iterator,
+                          keySerializer,
+                          valueSerializer)),
+                  new If(neqNull(entry), writeChunk(buffer, entry, iterator, keyType, valueType)));
+            });
 
     return new If(not(inlineInvoke(map, "isEmpty", PRIMITIVE_BOOLEAN_TYPE)), whileAction);
   }
@@ -1096,27 +1112,38 @@ protected Expression writeChunk(
       TypeRef<?> keyType,
       TypeRef<?> valueType) {
     ListExpression expressions = new ListExpression();
-    Expression key =
-        tryCastIfPublic(new Invoke(entry, "getKey", "keyObj", OBJECT_TYPE), keyType, "key");
+    Expression key = tryCastIfPublic(inlineInvoke(entry, "getKey", OBJECT_TYPE), keyType, "key");
     Expression value =
-        tryCastIfPublic(new Invoke(entry, "getValue", "valueObj", OBJECT_TYPE), valueType, "value");
-    Expression keyTypeExpr = new Invoke(key, "getClass", "keyType", CLASS_TYPE);
-    Expression valueTypeExpr = new Invoke(value, "getClass", "valueType", CLASS_TYPE);
-    Expression writePlaceHolder = new Invoke(buffer, "writeInt16", Literal.ofShort((short) -1));
-    Expression chunkSizeOffset =
-        subtract(inlineInvoke(buffer, "writerIndex", PRIMITIVE_INT_TYPE), Literal.ofInt(-1));
+        tryCastIfPublic(inlineInvoke(entry, "getValue", OBJECT_TYPE), valueType, "value");
+    boolean keyMonomorphic = isMonomorphic(keyType);
+    boolean valueMonomorphic = isMonomorphic(valueType);
     Class<?> keyTypeRawType = keyType.getRawType();
     Class<?> valueTypeRawType = valueType.getRawType();
+    Expression keyTypeExpr =
+        keyMonomorphic
+            ? getClassExpr(keyTypeRawType)
+            : new Invoke(key, "getClass", "keyType", CLASS_TYPE);
+    Expression valueTypeExpr =
+        valueMonomorphic
+            ? getClassExpr(valueTypeRawType)
+            : new Invoke(value, "getClass", "valueType", CLASS_TYPE);
+    Expression writePlaceHolder = new Invoke(buffer, "writeInt16", Literal.ofShort((short) -1));
+    Expression chunkSizeOffset =
+        subtract(
+            inlineInvoke(buffer, "writerIndex", PRIMITIVE_INT_TYPE),
+            Literal.ofInt(1),
+            "chunkSizeOffset");
+
     Expression chunkHeader;
     Expression keySerializer, valueSerializer;
-    boolean keyMonomorphic = isMonomorphic(keyType);
-    boolean valueMonomorphic = isMonomorphic(valueType);
+
     boolean trackingKeyRef =
         visitFury(fury -> fury.getClassResolver().needToWriteRef(keyTypeRawType));
     boolean trackingValueRef =
         visitFury(fury -> fury.getClassResolver().needToWriteRef(valueTypeRawType));
     Expression keyWriteRef = Literal.ofBoolean(trackingKeyRef);
     Expression valueWriteRef = Literal.ofBoolean(trackingValueRef);
+    boolean inline = keyMonomorphic && valueMonomorphic;
     if (keyMonomorphic && valueMonomorphic) {
       keySerializer = getOrCreateSerializer(keyTypeRawType);
       valueSerializer = getOrCreateSerializer(valueTypeRawType);
@@ -1174,7 +1201,6 @@ protected Expression writeChunk(
       }
     }
     Expression chunkSize = ofInt("chunkSize", 0);
-
     expressions.add(
         key,
         value,
@@ -1189,21 +1215,38 @@ protected Expression writeChunk(
         valueWriteRef,
         new Invoke(buffer, "putByte", subtract(chunkSizeOffset, Literal.ofInt(1)), chunkHeader),
         chunkSize);
-
     Expression keyWriteRefExpr = keyWriteRef;
     Expression valueWriteRefExpr = valueWriteRef;
     Expression.While writeLoop =
         new Expression.While(
             Literal.ofBoolean(true),
-            () ->
-                new ListExpression(
-                    new If(
-                        or(
-                            isNull(key),
-                            isNull(value),
-                            neq(inlineInvoke(key, "getClass", CLASS_TYPE), keyTypeExpr),
-                            neq(inlineInvoke(value, "getClass", CLASS_TYPE), valueTypeExpr)),
-                        new Break()),
+            () -> {
+              Expression breakCondition;
+              if (keyMonomorphic && valueMonomorphic) {
+                breakCondition = or(eqNull(key), eqNull(value));
+              } else if (keyMonomorphic) {
+                breakCondition =
+                    or(
+                        eqNull(key),
+                        eqNull(value),
+                        neq(inlineInvoke(value, "getClass", CLASS_TYPE), valueTypeExpr));
+              } else if (valueMonomorphic) {
+                breakCondition =
+                    or(
+                        eqNull(key),
+                        eqNull(value),
+                        neq(inlineInvoke(key, "getClass", CLASS_TYPE), keyTypeExpr));
+              } else {
+                breakCondition =
+                    or(
+                        eqNull(key),
+                        eqNull(value),
+                        neq(inlineInvoke(key, "getClass", CLASS_TYPE), keyTypeExpr),
+                        neq(inlineInvoke(value, "getClass", CLASS_TYPE), valueTypeExpr));
+              }
+              Expression writeKey = new Invoke(keySerializer, "write", buffer, key);
+              if (trackingKeyRef) {
+                writeKey =
                     new If(
                         or(
                             not(keyWriteRefExpr),
@@ -1214,7 +1257,11 @@ protected Expression writeChunk(
                                     PRIMITIVE_BOOLEAN_TYPE,
                                     buffer,
                                     key))),
-                        new Invoke(keySerializer, "write", buffer, key)),
+                        writeKey);
+              }
+              Expression writeValue = new Invoke(valueSerializer, "write", buffer, value);
+              if (trackingValueRef) {
+                writeValue =
                     new If(
                         or(
                             not(valueWriteRefExpr),
@@ -1225,24 +1272,43 @@ protected Expression writeChunk(
                                     PRIMITIVE_BOOLEAN_TYPE,
                                     buffer,
                                     value))),
-                        new Invoke(valueSerializer, "write", buffer, value)),
-                    new Assign(add(chunkSize, Literal.ofInt(1)), chunkSize),
-                    new If(eq(chunkSize, Literal.ofInt(MAX_CHUNK_SIZE)), new Break()),
-                    new If(
-                        inlineInvoke(iterator, "hasNext", PRIMITIVE_BOOLEAN_TYPE),
-                        new ListExpression(
-                            new Assign(inlineInvoke(iterator, "next", MAP_ENTRY_TYPE), entry),
-                            new Assign(
-                                tryInlineCast(inlineInvoke(entry, "getKey", OBJECT_TYPE), keyType),
-                                key),
-                            new Assign(
-                                tryInlineCast(
-                                    inlineInvoke(entry, "getValue", OBJECT_TYPE), valueType),
-                                value)),
-                        new ListExpression(
-                            new Assign(new Literal(null, MAP_ENTRY_TYPE), entry), new Break()))));
-    expressions.add(
-        writeLoop, new Invoke(buffer, "putByte", chunkSizeOffset, chunkSize), new Return(entry));
+                        writeValue);
+              }
+              return new ListExpression(
+                  new If(breakCondition, new Break()),
+                  writeKey,
+                  writeValue,
+                  new Assign(chunkSize, add(chunkSize, Literal.ofInt(1))),
+                  new If(eq(chunkSize, Literal.ofInt(MAX_CHUNK_SIZE)), new Break()),
+                  new If(
+                      inlineInvoke(iterator, "hasNext", PRIMITIVE_BOOLEAN_TYPE),
+                      new ListExpression(
+                          new Assign(
+                              entry,
+                              cast(inlineInvoke(iterator, "next", OBJECT_TYPE), MAP_ENTRY_TYPE)),
+                          new Assign(
+                              key,
+                              tryInlineCast(inlineInvoke(entry, "getKey", OBJECT_TYPE), keyType)),
+                          new Assign(
+                              value,
+                              tryInlineCast(
+                                  inlineInvoke(entry, "getValue", OBJECT_TYPE), valueType))),
+                      list(new Assign(entry, new Literal(null, MAP_ENTRY_TYPE)), new Break())));
+            });
+    expressions.add(writeLoop, new Invoke(buffer, "putByte", chunkSizeOffset, chunkSize));
+    if (!inline) {
+      expressions.add(new Return(entry));
+      // method too big, spilt it into a new method.
+      // Generate similar signature as `AbstractMapSerializer.writeJavaChunk`(
+      //   MemoryBuffer buffer,
+      //   Entry<Object, Object> entry,
+      //   Iterator<Entry<Object, Object>> iterator,
+      //   Serializer keySerializer,
+      //   Serializer valueSerializer
+      //  )
+      Set<Expression> params = ofHashSet(buffer, entry, iterator);
+      return invokeGenerated(ctx, params, expressions, "writeChunk", false);
+    }
     return expressions;
   }
 

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

@@ -266,26 +266,29 @@ public ExprCode doGenCode(CodegenContext ctx) {
       if (StringUtils.isNotBlank(targetExprCode.code())) {
         codeBuilder.append(targetExprCode.code()).append('\n');
       }
+      String name = ctx.newName(namePrefix);
       String decl =
-        StringUtils.format(
-          "${type} ${name} = ${from};",
-          "type",
-          ctx.type(type()),
-          "name",
-          ctx.newName(namePrefix),
-          "from",
-          targetExprCode.value());
+          StringUtils.format(
+              "${type} ${name} = ${from};",
+              "type",
+              ctx.type(type()),
+              "name",
+              name,
+              "from",
+              targetExprCode.value());
       codeBuilder.append(decl);
-      return new ExprCode(codeBuilder.toString(), null, null);
+      return new ExprCode(
+          codeBuilder.toString(),
+          targetExprCode.isNull(),
+          Code.variable(type().getRawType(), name));
     }
 
     @Override
     public String toString() {
-      return String.format("%s %s = %s", type(), namePrefix, from);
+      return String.format("%s %s = %s;", type(), namePrefix, from);
     }
   }
 
-
   class Literal implements Expression {
     public static final Literal True = new Literal(true, PRIMITIVE_BOOLEAN_TYPE);
     public static final Literal False = new Literal(false, PRIMITIVE_BOOLEAN_TYPE);
@@ -2595,12 +2598,12 @@ public String toString() {
   }
 
   class Assign implements Expression {
-    private Expression from;
     private Expression to;
+    private Expression from;
 
-    public Assign(Expression from, Expression to) {
-      this.from = from;
+    public Assign(Expression to, Expression from) {
       this.to = to;
+      this.from = from;
     }
 
     @Override
@@ -2622,7 +2625,7 @@ public ExprCode doGenCode(CodegenContext ctx) {
               });
       String assign =
           StringUtils.format(
-              "${from} = ${to};", "from", fromExprCode.value(), "to", toExprCode.value());
+              "${to} = ${from};", "from", fromExprCode.value(), "to", toExprCode.value());
       codeBuilder.append(assign);
       return new ExprCode(codeBuilder.toString(), null, null);
     }

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

@@ -34,6 +34,7 @@
 import java.util.Collections;
 import java.util.List;
 import org.apache.fury.codegen.Expression.Cast;
+import org.apache.fury.codegen.Expression.ListExpression;
 import org.apache.fury.codegen.Expression.LogicalAnd;
 import org.apache.fury.codegen.Expression.LogicalOr;
 import org.apache.fury.codegen.Expression.Null;
@@ -45,6 +46,9 @@
 
 /** Expression utils to create expression and code in a more convenient way. */
 public class ExpressionUtils {
+  public static ListExpression list(Expression... expressions) {
+    return new ListExpression(expressions);
+  }
 
   public static Expression newObjectArray(Expression... expressions) {
     return new NewArray(TypeRef.of(Object[].class), expressions);
@@ -71,6 +75,11 @@ public static Expression eqNull(Expression target) {
     return eq(target, new Null(target.type()));
   }
 
+  public static Expression neqNull(Expression target) {
+    Preconditions.checkArgument(!target.type().isPrimitive());
+    return neq(target, new Null(target.type()));
+  }
+
   public static LogicalAnd and(Expression left, Expression right, String name) {
     return new LogicalAnd(false, left, right);
   }
@@ -153,7 +162,7 @@ public static Arithmetic subtract(Expression left, Expression right) {
   }
 
   public static Arithmetic subtract(Expression left, Expression right, String valuePrefix) {
-    Arithmetic arithmetic = new Arithmetic(true, "-", left, right);
+    Arithmetic arithmetic = new Arithmetic(false, "-", left, right);
     arithmetic.valuePrefix = valuePrefix;
     return arithmetic;
   }