Skip to content

Commit

Permalink
support ref tracking for map chunk encoding
Browse files Browse the repository at this point in the history
  • Loading branch information
chaokunyang committed Jan 27, 2025
1 parent 8424b69 commit 4e2a1b8
Show file tree
Hide file tree
Showing 4 changed files with 95 additions and 53 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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))));
Expand Down Expand Up @@ -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 =
Expand All @@ -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,
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
}
Expand All @@ -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);
Expand All @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
}

0 comments on commit 4e2a1b8

Please sign in to comment.