Skip to content

jextract/jni: Overhaul enum associated value retrieval#727

Draft
sidepelican wants to merge 17 commits intoswiftlang:mainfrom
sidepelican:enum_case_getter
Draft

jextract/jni: Overhaul enum associated value retrieval#727
sidepelican wants to merge 17 commits intoswiftlang:mainfrom
sidepelican:enum_case_getter

Conversation

@sidepelican
Copy link
Copy Markdown
Contributor

This PR overhauls the implementation of enum associated value retrieval.
By leveraging the recent improvements in tuple handling, I can now support complex patterns (such as arrays and generic types).

The Problem

The current getAsCase implementation relies on specialized, hard-coded logic.

let getAsCaseFunction = TranslatedFunctionDecl(
name: getAsCaseName,
isStatic: false,
isThrowing: false,
isAsync: false,
nativeFunctionName: "$\(getAsCaseName)",
parentName: SwiftQualifiedTypeName(enumName),
functionTypes: [],
translatedFunctionSignature: TranslatedFunctionSignature(
selfParameter: TranslatedParameter(
parameter: JavaParameter(name: "selfPointer", type: .long),
conversion: .aggregate(
[
.ifStatement(
.constant("getDiscriminator() != Discriminator.\(caseName.uppercased())"),
thenExp: .constant("return Optional.empty();"),
),
.valueMemoryAddress(.placeholder),
]
),
),
selfTypeParameter: !isGenericParent
? nil
: .init(
parameter: JavaParameter(name: "selfTypePointer", type: .long),
conversion: .typeMetadataAddress(.placeholder),
),
parameters: [],
result: TranslatedResult(
javaType: .optional(caseType),
nativeJavaType: nativeParametersType,
outParameters: conversions.flatMap(\.translated.outParameters),
conversion: enumCase.parameters.isEmpty
? constructRecordConversion
: .aggregate(variable: ("$nativeParameters", nativeParametersType), [constructRecordConversion]),
),
exceptions: exceptions,
),
nativeFunctionSignature: NativeFunctionSignature(
selfParameter: NativeParameter(
parameters: [JavaParameter(name: "selfPointer", type: .long)],
conversion: .extractSwiftValue(.placeholder, swiftType: .nominal(enumCase.enumType), allowNil: false),
indirectConversion: nil,
conversionCheck: nil,
),
selfTypeParameter: !isGenericParent
? nil
: .init(
parameters: [JavaParameter(name: "selfTypePointer", type: .long)],
conversion: .extractMetatypeValue(.placeholder),
indirectConversion: nil,
conversionCheck: nil,
),
parameters: [],
result: NativeResult(
javaType: nativeParametersType,
conversion: .placeholder,
outParameters: conversions.flatMap(\.native.outParameters),
),
),
)

It only works for very simple patterns.
Adding support for complex types (tuples, arrays, etc.) within the current structure is not sustainable and would lead to code duplication.

The Solution

Instead of a specialized path, I have unified the logic with our existing tuple infrastructure.

In Swift, associated values are extracted as an optional tuple (e.g., (Foo, String)?).
In Java, the glue code wraps it into the appropriate Case class.

Example

Original:

enum MyEnum {
  case foo(Foo, String)
}

Generated:

extension MyEnum {
  fileprivate func getAsFoo() -> (Foo, String)? {
    if case let .foo(_0, _1) = self {
      return (_0, _1)
    }
    return nil
  }
}

// native $getAsFoo implementation
@_cdecl(...)
public func Java_com_example_swift_MyEnum__00024getAsFoo...
public Optional<Case.Foo> getAsFoo() {
  MyEnum.$getAsFoo(...);
  ...
  Optional<Tuple2<Foo, String>> associatedValues$ = ...; // get results as optional tuple
  // Wrap it in a Case class
  return associatedValues$.map((t) -> {
    return new Case.Foo(t.$0, t.$1);
  }
}
private static native void $getAsFoo(long selfPointer, byte[] result_discriminator$, ...);

Benchmarks

I have tested via ./gradlew jmh -PjmhIncludes='.EnumBenchmark*' in Samples/SwiftJavaExtractJNISampleApp/.

  • before
Benchmark                          Mode  Cnt    Score   Error  Units
EnumBenchmark.getAssociatedValues  avgt   30  974.888 ± 8.462  ns/op
  • after
Benchmark                          Mode  Cnt    Score   Error  Units
EnumBenchmark.getAssociatedValues  avgt   30  750.577 ± 8.338  ns/op

The performance improvement is likely a direct result of reducing the number of JNI calls.

Note

This PR is currently a draft as it depends on #725. I will mark it as ready for review once the dependency is merged.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant