Add annotation support for string parameters

dylibso · Sep 2, 2024 · d9a84dc · d9a84dc
1 parent fffe966
commit d9a84dc
Show file tree

Hide file tree

Showing 9 changed files with 149 additions and 62 deletions.
diff --git a/function-annotations/src/main/java/com/dylibso/chicory/function/annotations/Buffer.java b/function-annotations/src/main/java/com/dylibso/chicory/function/annotations/Buffer.java
@@ -0,0 +1,10 @@
+package com.dylibso.chicory.function.annotations;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.PARAMETER)
+public @interface Buffer {}
diff --git a/function-annotations/src/main/java/com/dylibso/chicory/function/annotations/CString.java b/function-annotations/src/main/java/com/dylibso/chicory/function/annotations/CString.java
@@ -0,0 +1,10 @@
+package com.dylibso.chicory.function.annotations;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.PARAMETER)
+public @interface CString {}
diff --git a/...ion-processor/src/main/java/com/dylibso/chicory/function/processor/FunctionProcessor.java b/...ion-processor/src/main/java/com/dylibso/chicory/function/processor/FunctionProcessor.java
@@ -6,6 +6,8 @@
 import static javax.tools.Diagnostic.Kind.ERROR;
 import static javax.tools.Diagnostic.Kind.NOTE;
 
+import com.dylibso.chicory.function.annotations.Buffer;
+import com.dylibso.chicory.function.annotations.CString;
 import com.dylibso.chicory.function.annotations.HostModule;
 import com.dylibso.chicory.function.annotations.WasmExport;
 import com.github.javaparser.ast.ArrayCreationLevel;
@@ -137,10 +139,7 @@ private Expression processMethod(
         NodeList<Expression> paramTypes = new NodeList<>();
         NodeList<Expression> arguments = new NodeList<>();
         for (VariableElement parameter : executable.getParameters()) {
-            var argExpr =
-                    new ArrayAccessExpr(
-                            new NameExpr("args"),
-                            new IntegerLiteralExpr(String.valueOf(paramTypes.size())));
+            var argExpr = argExpr(paramTypes.size());
             switch (parameter.asType().toString()) {
                 case "int":
                     paramTypes.add(valueType("I32"));
@@ -158,6 +157,31 @@ private Expression processMethod(
                     paramTypes.add(valueType("F64"));
                     arguments.add(new MethodCallExpr(argExpr, "asDouble"));
                     break;
+                case "java.lang.String":
+                    if (annotatedWith(parameter, Buffer.class)) {
+                        var lenExpr = argExpr(paramTypes.size() + 1);
+                        paramTypes.add(valueType("I32"));
+                        paramTypes.add(valueType("I32"));
+                        arguments.add(
+                                new MethodCallExpr(
+                                        new MethodCallExpr(new NameExpr("instance"), "memory"),
+                                        "readString",
+                                        new NodeList<>(
+                                                new MethodCallExpr(argExpr, "asInt"),
+                                                new MethodCallExpr(lenExpr, "asInt"))));
+                    } else if (annotatedWith(parameter, CString.class)) {
+                        paramTypes.add(valueType("I32"));
+                        paramTypes.add(valueType("I32"));
+                        arguments.add(
+                                new MethodCallExpr(
+                                        new MethodCallExpr(new NameExpr("instance"), "memory"),
+                                        "readCString",
+                                        new NodeList<>(new MethodCallExpr(argExpr, "asInt"))));
+                    } else {
+                        log(ERROR, "Missing annotation for WASM type: java.lang.String", parameter);
+                        throw new AbortProcessingException();
+                    }
+                    break;
                 case "com.dylibso.chicory.runtime.Instance":
                     arguments.add(new NameExpr("instance"));
                     break;
@@ -266,6 +290,10 @@ private static boolean annotatedWith(Element element, Class<? extends Annotation
                 .anyMatch(annotationName::equals);
     }
 
+    private static Expression argExpr(int n) {
+        return new ArrayAccessExpr(new NameExpr("args"), new IntegerLiteralExpr(String.valueOf(n)));
+    }
+
     private static Expression valueType(String type) {
         return new FieldAccessExpr(new NameExpr("ValueType"), type);
     }

diff --git a/...processor/src/test/java/com/dylibso/chicory/function/processor/FunctionProcessorTest.java b/...processor/src/test/java/com/dylibso/chicory/function/processor/FunctionProcessorTest.java
@@ -29,16 +29,30 @@ void generateModules() {
     }
 
     @Test
-    void invalidParameterType() {
+    void invalidParameterTypeUnsupported() {
         Compilation compilation =
                 javac().withProcessors(new FunctionProcessor())
-                        .compile(JavaFileObjects.forResource("InvalidParameter.java"));
+                        .compile(JavaFileObjects.forResource("InvalidParameterUnsupported.java"));
 
         assertThat(compilation).failed();
 
         assertThat(compilation)
-                .hadErrorContaining("Unsupported WASM type: java.lang.String")
-                .inFile(JavaFileObjects.forResource("InvalidParameter.java"))
+                .hadErrorContaining("Unsupported WASM type: java.math.BigDecimal")
+                .inFile(JavaFileObjects.forResource("InvalidParameterUnsupported.java"))
+                .onLineContaining("public long square(BigDecimal x) {");
+    }
+
+    @Test
+    void invalidParameterTypeString() {
+        Compilation compilation =
+                javac().withProcessors(new FunctionProcessor())
+                        .compile(JavaFileObjects.forResource("InvalidParameterString.java"));
+
+        assertThat(compilation).failed();
+
+        assertThat(compilation)
+                .hadErrorContaining("Missing annotation for WASM type: java.lang.String")
+                .inFile(JavaFileObjects.forResource("InvalidParameterString.java"))
                 .onLineContaining("public long concat(int a, String s) {");
     }
 

diff --git a/.../src/test/resources/InvalidParameter.java → ...est/resources/InvalidParameterString.java b/.../src/test/resources/InvalidParameter.java → ...est/resources/InvalidParameterString.java
@@ -4,7 +4,7 @@
 import com.dylibso.chicory.function.annotations.WasmExport;
 
 @HostModule("bad_param")
-public final class InvalidParameter {
+public final class InvalidParameterString {
 
     @WasmExport
     public long concat(int a, String s) {

diff --git a/function-processor/src/test/resources/InvalidParameterUnsupported.java b/function-processor/src/test/resources/InvalidParameterUnsupported.java
@@ -0,0 +1,14 @@
+package chicory.testing;
+
+import com.dylibso.chicory.function.annotations.HostModule;
+import com.dylibso.chicory.function.annotations.WasmExport;
+import java.math.BigDecimal;
+
+@HostModule("bad_param")
+public final class InvalidParameterUnsupported {
+
+    @WasmExport
+    public long square(BigDecimal x) {
+        return x.pow(2);
+    }
+}
diff --git a/function-processor/src/test/resources/Simple.java b/function-processor/src/test/resources/Simple.java
@@ -1,17 +1,32 @@
 package chicory.testing;
 
+import com.dylibso.chicory.function.annotations.Buffer;
+import com.dylibso.chicory.function.annotations.CString;
 import com.dylibso.chicory.function.annotations.HostModule;
 import com.dylibso.chicory.function.annotations.WasmExport;
 import com.dylibso.chicory.runtime.HostFunction;
 import com.dylibso.chicory.runtime.Memory;
 import com.dylibso.chicory.wasm.exceptions.ChicoryException;
+import java.security.SecureRandom;
 
 @HostModule("simple")
 public final class Simple {
 
     @WasmExport
-    public void print(Memory memory, int ptr, int len) {
-        System.out.println(memory.readString(ptr, len));
+    public void print(@Buffer String data) {
+        System.out.println(data);
+    }
+
+    @WasmExport
+    public void printx(@CString String data) {
+        System.out.println(data);
+    }
+
+    @WasmExport
+    public void randomGet(Memory memory, int ptr, int len) {
+        byte[] data = new byte[len];
+        new SecureRandom().nextBytes(data);
+        memory.write(ptr, data);
     }
 
     @WasmExport

diff --git a/function-processor/src/test/resources/SimpleGenerated.java b/function-processor/src/test/resources/SimpleGenerated.java
@@ -16,14 +16,35 @@ public static HostFunction[] toHostFunctions(Simple functions) {
         return new HostFunction[] {
             new HostFunction(
                     (Instance instance, Value... args) -> {
-                        functions.print(instance.memory(), args[0].asInt(), args[1].asInt());
+                        functions.print(
+                                instance.memory().readString(args[0].asInt(), args[1].asInt()));
                         return null;
                     },
                     "simple",
                     "print",
                     List.of(ValueType.I32, ValueType.I32),
                     List.of()
             ),
+            new HostFunction(
+                    (Instance instance, Value... args) -> {
+                        functions.printx(instance.memory().readCString(args[0].asInt()));
+                        return null;
+                    },
+                    "simple",
+                    "printx",
+                    List.of(ValueType.I32, ValueType.I32),
+                    List.of()
+            ),
+            new HostFunction(
+                    (Instance instance, Value... args) -> {
+                        functions.randomGet(instance.memory(), args[0].asInt(), args[1].asInt());
+                        return null;
+                    },
+                    "simple",
+                    "random_get",
+                    List.of(ValueType.I32, ValueType.I32),
+                    List.of()
+            ),
             new HostFunction(
                     (Instance instance, Value... args) -> {
                         functions.exit();