Allow specifying multiple @Code and @Library annotations

.gitignore

@@ -1,3 +1,4 @@
 /out/
 /.idea/
 *.iml
+/nalim.jar

example/one/nalim/example/CpuId.java

@@ -0,0 +1,25 @@
+package one.nalim.example;
+
+import one.nalim.Arch;
+import one.nalim.Code;
+import one.nalim.Linker;
+import one.nalim.Os;
+
+public class CpuId {
+
+    @Code(
+            os = Os.WINDOWS,
+            arch = Arch.AMD64,
+            value = "4989D9 89D0 0FA2 418900 41895804 41894808 4189500C 4C89CB C3"
+    )
+    @Code(
+            os = Os.LINUX,
+            arch = Arch.AMD64,
+            value = "4889D7 89F0 4889DE 0FA2 8907 895F04 894F08 89570C 4889F3 C3"
+    )
+    private static native void cpuid(int func, int[] out);
+
+    static {
+        Linker.linkClass(CpuId.class);
+    }
+}

src/one/nalim/AnnotationUtil.java

@@ -0,0 +1,112 @@
+package one.nalim;
+
+import java.lang.annotation.Annotation;
+import java.lang.reflect.AnnotatedElement;
+import java.lang.reflect.Array;
+import java.lang.reflect.Method;
+import java.util.function.Function;
+
+final class AnnotationUtil {
+
+    private static final int MAX_SCORE = 3;
+    private static final int NO_MATCH = -1;
+
+    static Library findLibrary(AnnotatedElement element, Os expectedOs, Arch expectedArch) {
+        return find(
+                Library.class,
+                element.getAnnotation(LibrarySet.class),
+                element.getAnnotation(Library.class),
+                LibrarySet::value, Library::os, Library::arch,
+                expectedOs, expectedArch);
+    }
+
+    static Code findCode(Method method, Os expectedOs, Arch expectedArch) {
+        return find(
+                Code.class,
+                method.getAnnotation(CodeSet.class),
+                method.getAnnotation(Code.class),
+                CodeSet::value, Code::os, Code::arch,
+                expectedOs, expectedArch);
+    }
+
+    private static <T extends Annotation, U extends Annotation> U find(
+            Class<U> annotationType,
+            T containerAnnotation,
+            U annotation,
+            Function<T, U[]> annotationsExtractor,
+            Function<U, Os> osExtractor, Function<U, Arch> archExtractor,
+            Os expectedOs, Arch expectedArch) {
+
+        final U[] annotations;
+        if (containerAnnotation == null) {
+            if (annotation == null) {
+                return null;
+            }
+
+            annotations = (U[]) Array.newInstance(annotationType, 1);
+            annotations[0] = annotation;
+        } else {
+            annotations = annotationsExtractor.apply(containerAnnotation);
+            if (annotations == null || annotations.length == 0) {
+                return null;
+            }
+        }
+
+        // Find the best-matching annotation by comparing their scores.
+        U match = null;
+        int matchScore = -1;
+        for (final U a : annotations) {
+            final int score = score(
+                    expectedOs, expectedArch,
+                    osExtractor.apply(a), archExtractor.apply(a));
+
+            if (score > matchScore) {
+                match = a;
+                matchScore = score;
+                if (score >= MAX_SCORE) {
+                    break;
+                }
+            }
+        }
+
+        return match;
+    }
+
+    private static int score(Os expectedOs, Arch expectedArch, Os os, Arch arch) {
+        if (os == expectedOs) {
+            if (arch == expectedArch) {
+                // Both OS and arch were specified and both match.
+                return MAX_SCORE;
+            }
+
+            if (arch == Arch.UNSPECIFIED) {
+                // Only OS was specified and it matches.
+                return MAX_SCORE - 1;
+            }
+
+            // Both OS and arch were specified, but arch doesn't match.
+            return NO_MATCH;
+        }
+
+        if (os == Os.UNSPECIFIED) {
+            if (arch == expectedArch) {
+                // Only Arch was specified, and it matches.
+                return MAX_SCORE - 2;
+            }
+
+            if (arch == Arch.UNSPECIFIED) {
+                // Neither OS nor arch were specified.
+                return MAX_SCORE - 3;
+            }
+
+            // Only Arch was specified, but it doesn't match.
+            return NO_MATCH;
+        }
+
+        // OS was specified, but it doesn't match.
+        return NO_MATCH;
+    }
+
+    private AnnotationUtil() {
+    }
+}

src/one/nalim/Arch.java

@@ -0,0 +1,29 @@
+package one.nalim;
+
+import java.util.Locale;
+
+public enum Arch {
+    UNSPECIFIED,
+    AMD64,
+    AARCH64,
+    RISCV64,
+    UNKNOWN;
+
+    static final Arch CURRENT;
+
+    static {
+        String arch = System.getProperty("os.arch").toLowerCase(Locale.US);
+        if (!arch.contains("64")) {
+            // Non-64bit architectures are unsupported.
+            CURRENT = UNKNOWN;
+        } else if (arch.contains("x86") || arch.contains("amd64")) {
+            CURRENT = AMD64;
+        } else if (arch.contains("aarch") || arch.contains("arm")) {
+            CURRENT = AARCH64;
+        } else if (arch.contains("riscv")) {
+            CURRENT = RISCV64;
+        } else {
+            CURRENT = UNKNOWN;
+        }
+    }
+}

src/one/nalim/CallingConvention.java

@@ -25,28 +25,24 @@
 
 import java.lang.annotation.Annotation;
 import java.nio.ByteBuffer;
+import java.util.Objects;
 
 abstract class CallingConvention {
 
     static CallingConvention getInstance() {
-        String arch = System.getProperty("os.arch").toLowerCase();
-        if (!arch.contains("64")) {
-            throw new IllegalStateException("Unsupported architecture: " + arch);
-        }
-
-        if (arch.contains("aarch") || arch.contains("arm")) {
-            return new AArch64CallingConvention();
-        }
-
-        if (arch.contains("riscv")) {
-            return new RISCV64CallingConvention();
-        }
-
-        String os = System.getProperty("os.name").toLowerCase();
-        if (os.contains("windows")) {
-            return new AMD64WindowsCallingConvention();
-        } else {
-            return new AMD64LinuxCallingConvention();
+        switch (Arch.CURRENT) {
+            case AMD64:
+                if (Objects.requireNonNull(Os.CURRENT) == Os.WINDOWS) {
+                    return new AMD64WindowsCallingConvention();
+                }
+                return new AMD64LinuxCallingConvention();
+            case AARCH64:
+                return new AArch64CallingConvention();
+            case RISCV64:
+                return new RISCV64CallingConvention();
+            default:
+                throw new IllegalStateException(
+                        "Unsupported architecture: " + System.getProperty("os.arch"));
         }
     }
 

src/one/nalim/Code.java

@@ -19,6 +19,7 @@
 package one.nalim;
 
 import java.lang.annotation.ElementType;
+import java.lang.annotation.Repeatable;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
@@ -28,11 +29,26 @@
  */
 @Retention(RetentionPolicy.RUNTIME)
 @Target(ElementType.METHOD)
+@Repeatable(CodeSet.class)
 public @interface Code {
     /**
      * Machine code for the method implementation
      * encoded as a hex string, possibly with spaces.
      * Typically, the code should end with a return instruction.
      */
     String value();
+
+    /**
+     * The operating system that will run the code specified in this annotation.
+     * When there are multiple {@link Code} annotations, the one with matching
+     * `os` attribute will have precedence over the one without it.
+     */
+    Os os() default Os.UNSPECIFIED;
+
+    /**
+     * The CPU architecture that will run the code specified in this annotation.
+     * When there are multiple {@link Code} annotations, the one with matching
+     * `arch` attribute will have precedence over the one without it.
+     */
+    Arch arch() default Arch.UNSPECIFIED;
 }

src/one/nalim/CodeSet.java

@@ -0,0 +1,12 @@
+package one.nalim;
+
+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.METHOD)
+public @interface CodeSet {
+    Code[] value();
+}

src/one/nalim/Library.java

@@ -19,6 +19,7 @@
 package one.nalim;
 
 import java.lang.annotation.ElementType;
+import java.lang.annotation.Repeatable;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
@@ -28,6 +29,21 @@
  */
 @Retention(RetentionPolicy.RUNTIME)
 @Target({ElementType.TYPE, ElementType.METHOD})
+@Repeatable(LibrarySet.class)
 public @interface Library {
     String value();
+
+    /**
+     * The operating system that will run the code specified in this annotation.
+     * When there are multiple {@link Library} annotations, the one with matching
+     * `os` attribute will have precedence over the one without it.
+     */
+    Os os() default Os.UNSPECIFIED;
+
+    /**
+     * The CPU architecture that will run the code specified in this annotation.
+     * When there are multiple {@link Library} annotations, the one with matching
+     * `arch` attribute will have precedence over the one without it.
+     */
+    Arch arch() default Arch.UNSPECIFIED;
 }

src/one/nalim/LibrarySet.java

@@ -0,0 +1,12 @@
+package one.nalim;
+
+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.TYPE, ElementType.METHOD})
+public @interface LibrarySet {
+    Library[] value();
+}

src/one/nalim/Link.java

@@ -19,6 +19,7 @@
 package one.nalim;
 
 import java.lang.annotation.ElementType;
+import java.lang.annotation.Repeatable;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;

src/one/nalim/Linker.java

@@ -67,13 +67,13 @@ public static long findAddress(String symbol) {
     }
 
     public static void linkClass(Class<?> c) {
-        Library library = c.getAnnotation(Library.class);
+        Library library = AnnotationUtil.findLibrary(c, Os.CURRENT, Arch.CURRENT);
         if (library != null) {
             loadLibrary(library.value());
         }
 
         for (Method m : c.getDeclaredMethods()) {
-            if (m.getAnnotation(Link.class) != null || m.getAnnotation(Code.class) != null) {
+            if (isStaticNative(m)) {
                 linkMethod(m);
             }
         }
@@ -82,7 +82,7 @@ public static void linkClass(Class<?> c) {
     public static void linkMethod(Method m) {
         checkMethodType(m);
 
-        Code code = m.getAnnotation(Code.class);
+        Code code = AnnotationUtil.findCode(m, Os.CURRENT, Arch.CURRENT);
         if (code != null) {
             installCode(m, parseHex(code.value()));
             return;
@@ -99,7 +99,7 @@ public static void linkMethod(Method m) {
     public static void linkMethod(Method m, String symbol, boolean naked) {
         checkMethodType(m);
 
-        Library library = m.getAnnotation(Library.class);
+        Library library = AnnotationUtil.findLibrary(m, Os.CURRENT, Arch.CURRENT);
         if (library != null) {
             loadLibrary(library.value());
         }
@@ -122,12 +122,16 @@ public static void linkMethod(Method m, String symbol, boolean naked) {
     }
 
     private static void checkMethodType(Method m) {
-        int modifiers = m.getModifiers();
-        if (!Modifier.isStatic(modifiers) || !Modifier.isNative(modifiers)) {
+        if (!isStaticNative(m)) {
             throw new IllegalArgumentException("Method must be static native: " + m);
         }
     }
 
+    private static boolean isStaticNative(Method m) {
+        int modifiers = m.getModifiers();
+        return Modifier.isStatic(modifiers) && Modifier.isNative(modifiers);
+    }
+
     private static byte[] parseHex(String hex) {
         hex = hex.replaceAll("\\s+", "");
 

src/one/nalim/Os.java

@@ -0,0 +1,26 @@
+package one.nalim;
+
+import java.util.Locale;
+
+public enum Os {
+    UNSPECIFIED,
+    LINUX,
+    MACOS,
+    WINDOWS,
+    UNKNOWN;
+
+    static final Os CURRENT;
+
+    static {
+        final String os = System.getProperty("os.name").toLowerCase(Locale.US);
+        if (os.contains("linux")) {
+            CURRENT = Os.LINUX;
+        } else if (os.contains("windows")) {
+            CURRENT = Os.WINDOWS;
+        } else if (os.contains("mac") || os.contains("darwin") || os.contains("os x")) {
+            CURRENT = Os.MACOS;
+        } else {
+            CURRENT = UNKNOWN;
+        }
+    }
+}