From cc3d53d6a7c43923936f06200e295585be4ea477 Mon Sep 17 00:00:00 2001
From: squid233 <60126026+squid233@users.noreply.github.com>
Date: Sun, 1 Sep 2024 14:42:21 +0800
Subject: [PATCH] Fix bug for methods returning MethodHandle

---
 README.md                                     |  2 +-
 gradle.properties                             |  2 +-
 src/main/java/overrun/marshal/Downcall.java   | 20 ++++++++++++-------
 .../marshal/test/downcall/IDowncall.java      |  8 +++++++-
 4 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 3473014..720f36c 100644
--- a/README.md
+++ b/README.md
@@ -88,7 +88,7 @@ Import as a Gradle dependency:
 
 ```groovy
 dependencies {
-    implementation("io.github.over-run:marshal:0.1.0-alpha.30-jdk23")
+    implementation("io.github.over-run:marshal:0.1.0-alpha.31-jdk23")
 }
 ```
 
diff --git a/gradle.properties b/gradle.properties
index bf0dafd..98c0864 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -12,7 +12,7 @@ projGroupId=io.github.over-run
 projArtifactId=marshal
 # The project name should only contain lowercase letters, numbers and hyphen.
 projName=marshal
-projVersion=0.1.0-alpha.30-jdk23
+projVersion=0.1.0-alpha.31-jdk23
 projDesc=Marshaler of native libraries
 # Uncomment them if you want to publish to maven repository.
 projUrl=https://github.com/Over-Run/marshal
diff --git a/src/main/java/overrun/marshal/Downcall.java b/src/main/java/overrun/marshal/Downcall.java
index 2a80a66..192a036 100644
--- a/src/main/java/overrun/marshal/Downcall.java
+++ b/src/main/java/overrun/marshal/Downcall.java
@@ -635,12 +635,16 @@ private static DowncallData generateData(
 
             // function descriptor
             final FunctionDescriptor get = descriptorMap.get(entrypoint);
-            final FunctionDescriptor descriptor = get != null ?
-                get :
-                DescriptorTransformer.getInstance().process(new DescriptorTransformer.Context(method,
-                    methodData.descriptorSkipFirstParameter(),
-                    methodData.parameters()));
-
+            final FunctionDescriptor descriptor;
+            if (method.getReturnType() == MethodHandle.class) {
+                descriptor = get;
+            } else {
+                descriptor = get != null ?
+                    get :
+                    DescriptorTransformer.getInstance().process(new DescriptorTransformer.Context(method,
+                        methodData.descriptorSkipFirstParameter(),
+                        methodData.parameters()));
+            }
             descriptorMap1.put(entrypoint, descriptor);
 
             final Optional<MemorySegment> optional = lookup.find(entrypoint);
@@ -664,7 +668,9 @@ private static DowncallData generateData(
                     throw new NoSuchElementException("Symbol not found: " + entrypoint + " (" + descriptor + "): " + methodData.signatureString());
                 }
             }
-            map.putIfAbsent(entrypoint, handle);
+            if (!map.containsKey(entrypoint) || map.get(entrypoint) == null) {
+                map.put(entrypoint, handle);
+            }
         }
         return new DowncallData(Collections.unmodifiableMap(descriptorMap1),
             Collections.unmodifiableMap(map),
diff --git a/src/test/java/overrun/marshal/test/downcall/IDowncall.java b/src/test/java/overrun/marshal/test/downcall/IDowncall.java
index 0f83d50..1e41182 100644
--- a/src/test/java/overrun/marshal/test/downcall/IDowncall.java
+++ b/src/test/java/overrun/marshal/test/downcall/IDowncall.java
@@ -38,7 +38,10 @@
  * @since 0.1.0
  */
 public interface IDowncall {
-    Map<String, FunctionDescriptor> MAP = Map.of("testDefault", FunctionDescriptor.of(ValueLayout.JAVA_INT));
+    Map<String, FunctionDescriptor> MAP = Map.of(
+        "testDefault", FunctionDescriptor.of(ValueLayout.JAVA_INT),
+        "testReturnInt", FunctionDescriptor.of(ValueLayout.JAVA_INT)
+    );
 
     static IDowncall getInstance(boolean testDefaultNull) {
         ProcessorTypes.registerStruct(Vector3.class, Vector3.OF);
@@ -46,6 +49,9 @@ static IDowncall getInstance(boolean testDefaultNull) {
         return Downcall.load(MethodHandles.lookup(), DowncallProvider.lookup(testDefaultNull), DowncallOption.descriptors(MAP));
     }
 
+    @Entrypoint("testReturnInt")
+    MethodHandle mh_testReturnInt();
+
     void test();
 
     @Entrypoint("test")