libadalang/testsuite/drivers/java_driver.py

from __future__ import annotations

import os
import os.path as P
import subprocess

from e3.fs import mkdir, rm
from e3.testsuite import logger
from e3.testsuite.driver.classic import TestAbortWithFailure, TestSkip

from drivers.base_driver import BaseDriver


class JavaDriver(BaseDriver):
    """
    Driver to build and run tests for the Java bindings.
    """

    ni_main: str
    """
    Path to the native-image pre-compiled Java test main (see the
    build_ni_main method).
    """

    # Template for the Java test main source
    ni_main_template = """import java.util.Arrays;
    :IMPORTS:

    public final class NativeImageMain {
        private static String[] toForward(String[] srcArgs) {
            return Arrays.copyOfRange(srcArgs, 1, srcArgs.length);
        }

        public static void main(String[] args) {
            if(args.length < 1)
                throw new RuntimeException("Usage: main [test_class] ...");
            switch(args[0]) {
                :TEST_CASES:
            }
        }
    }
    """

    @property
    def mode(self) -> str:
        """
        Return the Java bindings execution mode for this test ("graal_c_api" or
        "jni").
        """
        result = self.test_env["mode"]
        assert result in ("graal_c_api", "jni")
        return result

    @staticmethod
    def in_java_bindings(java_bindings_dir: str, *args) -> str | None:
        """
        Return the absolute path composed of parts provided in `args` from the
        directory provided in `java_bindings_dir`. This function checks that
        the result is an existing file, otherwise it returns `None`.
        """
        bindings_dir = os.path.abspath(java_bindings_dir)
        requested_file = os.path.join(bindings_dir, *args)
        return requested_file if os.path.isfile(requested_file) else None

    @staticmethod
    def jar_file(java_bindings_dir: str) -> str | None:
        """
        Return the path to Libadalang's JAR file, or None if it cannot be
        found.
        """
        return JavaDriver.in_java_bindings(
            java_bindings_dir,
            "target",
            "libadalang.jar",
        )

    @staticmethod
    def lib_jar(java_bindings_dir: str, lib_name: str) -> str | None:
        """
        Return the path to the requested Java library JAR file that is located
        in the tree of the provided Java bindings directory.
        If the required lib file doesn't exist, this function returns None.
        """
        return JavaDriver.in_java_bindings(
            java_bindings_dir,
            "target",
            "lib",
            lib_name,
        )

    @staticmethod
    def java_exec() -> str:
        """
        Return the path to the "java" executable.
        """
        return os.path.realpath(os.path.join(
            os.environ['JAVA_HOME'], 'bin', 'java'
        ))

    @staticmethod
    def javac_exec() -> str:
        """
        Return the path to the "javac" executable.
        """
        return os.path.realpath(os.path.join(
            os.environ['JAVA_HOME'], 'bin', 'javac'
        ))

    @staticmethod
    def ni_compiler_exec() -> str:
        """
        Return the path to the "native-image" compiler executable.
        """
        return os.path.realpath(os.path.join(
            os.environ['GRAAL_HOME'],
            'bin',
            ('native-image.cmd' if os.name == 'nt' else 'native-image')
        ))

    @classmethod
    def build_ni_main(
        cls,
        working_dir: str,
        java_bindings_dir: str,
        drivers: list[JavaDriver],
    ) -> None:
        """
        Build a single main for all the given Graal C API test.

        Also set the "ni_main" attribute for all these drivers.
        """
        # If there are no Graal C API tests to run, do nothing
        if not drivers:
            return

        logger.info("Pre-compile Graal C API tests...")

        # Ensure Libadalang's Java bindings are available
        libadalang_jar = cls.jar_file(java_bindings_dir)
        if libadalang_jar is None:
            raise RuntimeError(
                "Cannot find the Java bindings JAR archive, make sure you"
                " built libadalang with the '--enable-java' flag"
            )

        # Create a build directory for the native-image build
        ni_dir = os.path.join(working_dir, "_native_image_build")
        ni_bin_dir = os.path.join(ni_dir, "bin")
        ni_bin = os.path.join(ni_bin_dir, "main")
        if os.path.isdir(ni_dir):
            rm(ni_dir)
        mkdir(ni_dir)
        mkdir(ni_bin_dir)

        # Locate the javac compiler and build a class path to give access to
        # Libadalangs' Java bindings, as well as test classes compiled in
        # "ni_dir".
        javac_exec = cls.javac_exec()
        class_path = os.pathsep.join([
            libadalang_jar,
            cls.lib_jar(java_bindings_dir, "langkit_support.jar"),
            cls.lib_jar(java_bindings_dir, "truffle-api.jar"),
            cls.lib_jar(java_bindings_dir, "polyglot.jar"),
            ni_dir,
        ])

        def javac(java_filename: str) -> None:
            """
            Shortcut to compile a Java class.
            """
            subprocess.check_call(
                [
                    javac_exec,
                    '-cp', class_path,
                    '-encoding', 'utf8',
                    '-d', ni_dir,
                    java_filename,
                ]
            )

        # Populate it with the Java sources to build
        import_stmts = []
        case_stmts = []
        for test in drivers:
            test_env = test.test_env

            # Copy the test Java source file and add a "package" statement
            # at the top of it.
            main_class = test_env.get('main_class', 'Main')
            src_java_file_name = os.path.join(
                test_env['test_dir'],
                test_env.get('java_path', '.'),
                f"{main_class}.java"
            )
            java_file_name = os.path.join(ni_dir, f"{main_class}.java")
            with open(src_java_file_name) as f:
                contents = f.read()
            with open(java_file_name, 'w') as f:
                print("package tests;", file=f)
                f.write(contents)

            # Compile the test Java file
            javac(java_file_name)

            # Add the test class to the native-image main Java content
            import_stmts.append(f"import tests.{main_class};")
            case_stmts.append(f'case "{main_class}": {main_class}.main'
                              '(toForward(args)); break;')

            # Communicate the path to the final native-image executable to
            # test drivers.
            test.ni_main = ni_bin

        # Create the native-image Java main file content, write it and
        # compile it.
        ni_main_content = cls.ni_main_template.replace(
            ':IMPORTS:',
            '\n'.join(import_stmts)
        ).replace(
            ':TEST_CASES:',
            '\n'.join(case_stmts)
        )
        ni_main_file_name = os.path.join(ni_dir, "NativeImageMain.java")
        with open(ni_main_file_name, 'w') as f:
            print(ni_main_content, file=f)
        javac(ni_main_file_name)

        os_specific_options = []
        if os.name != "nt":
            def find_file_in_env(file_name, env_var) -> str | None:
                for dir in os.environ.get(env_var, "").split(os.pathsep):
                    if os.path.isfile(os.path.join(dir, file_name)):
                        return dir
                return None

            # Find the directory that contains the "libadalang.h" header file
            header_dir = find_file_in_env("libadalang.h", "C_INCLUDE_PATH")
            assert header_dir is not None

            # Find the directory that contains the LAL shared library
            lib_dirs = [
                dir for dir in [
                    find_file_in_env("libadalang.so", "LIBRARY_PATH"),
                    find_file_in_env("libz.so", "LIBRARY_PATH"),
                ]
                if dir is not None
            ]

            # We also need to provide rpath-links to the compiler to allow it
            # to find libraries during linking phase.
            ld_library_path = os.environ.get('LD_LIBRARY_PATH')
            rpaths = (
                [
                    f"-Wl,-rpath-link={p}"
                    for p in ld_library_path.split(os.pathsep)
                ]
                if ld_library_path else
                []
            )

            # Create native-image options to provide required information when
            # spawning GCC.
            os_specific_options.extend([
                f"--native-compiler-options=-I{header_dir}",
                *[
                    f"--native-compiler-options=-L{lib_dir}"
                    for lib_dir in lib_dirs
                ],
                *[f"--native-compiler-options={rp}" for rp in rpaths],
            ])
        else:
            # Ensure the compiler isn't emitting warnings about CPU features
            os_specific_options.append("-march=native")

        # Run the native-image compiler
        subprocess.check_call([
            cls.ni_compiler_exec(),
            '-cp', class_path,
            '--no-fallback',
            "-Ob",
            "--silent",
            "--enable-native-access=org.graalvm.truffle",
            "-J--sun-misc-unsafe-memory-access=allow",
            "-H:+UnlockExperimentalVMOptions",
            "-H:-StrictQueryCodeCompilation",
            *os_specific_options,
            'NativeImageMain',
            ni_bin,
        ])

        logger.info("Done")

    def run(self):
        # Verify the Java bindings directory
        bindings_dir = self.env.java_bindings
        if bindings_dir is None:
            raise TestSkip('Test requires Java bindings')

        # Verify that the libadalang bindings are compiled
        libadalang_jar = self.jar_file(bindings_dir)
        if libadalang_jar is None:
            raise TestAbortWithFailure(
                "Cannot find the Java bindings JAR archive, make sure you"
                " built libadalang with the '--enable-java' flag"
            )

        # Get the class name and source file for the Java main
        main_class = self.test_env.get("main_class", "Main")
        main_java = self.test_dir(
            self.test_env.get("java_path", "."), f"{main_class}.java"
        )

        # Get the project path
        project_path = self.test_dir(self.test_env.get("projects_path", "."))

        if self.mode == 'graal_c_api':
            # Run the pre-compiled native-image produced executable with the
            # main Java class as first argument.
            assert self.ni_main, (
                "Test driver cannot access the pre-compiled Graal C API tests"
            )
            args = [self.ni_main, main_class]

        else:
            # Run the Java main class as Java program

            # Give access to the Libadalang Java bindings and to the tests'
            # Java sources.
            class_path = P.pathsep.join([
                libadalang_jar,
                self.lib_jar(bindings_dir, "langkit_support.jar"),
                self.lib_jar(bindings_dir, "truffle-api.jar"),
                self.test_env['working_dir']
            ])

            # Get the java.library.path from LD_LIBRARY_PATH and compiled JNI
            # stubs.
            java_library_path = P.pathsep.join(
                [P.join(bindings_dir, 'jni'), os.environ['LD_LIBRARY_PATH']]
            )

            # Prepare the command to run the Java main
            args = [
                self.java_exec(),
                '-cp', class_path,
                '-Dfile.encoding=UTF-8',
                # Enable the Java application to load and use native libraries
                "--enable-native-access=ALL-UNNAMED",
                f"-Djava.library.path={java_library_path}",
            ]
            args.append(main_java)

        # Run the test main. Mains expect a GPR project path as their only
        # argument.
        self.run_and_check(args + [project_path])