/*
 * Copyright (c) 2026, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

/*
 * @test id=AVX2
 * @bug 8369699
 * @key randomness
 * @summary Test the Template Library's expression generation for the Vector API.
 * @requires os.family == "linux" & os.simpleArch == "x64"
 * @modules jdk.incubator.vector
 * @modules java.base/jdk.internal.misc
 * @library /test/lib /
 * @compile ../../compiler/lib/verify/Verify.java
 * @run driver ${test.main.class} -XX:UseAVX=2
 */

// TODO: remove the x64 and linux restriction above. I added that for now so we are not flooded
//       with failures in the CI. We should remove these restriction once more bugs are fixed.
//       x64 linux is the easiest to debug on for me, that's why I picked it.
//       In addition, I put a UseAVX=2 restriction below, to avoid AVX512 bugs for now.
//
//       A trick to extend this to other platforms: create a new run block, so you have full
//       freedom to restrict it as necessary for platform and vector features.

package compiler.vectorapi;

import java.util.List;
import java.util.Arrays;
import java.util.ArrayList;
import java.util.Set;
import java.util.Random;
import jdk.test.lib.Utils;
import java.util.stream.IntStream;

import compiler.lib.compile_framework.CompileFramework;

import compiler.lib.template_framework.Template;
import compiler.lib.template_framework.TemplateToken;
import static compiler.lib.template_framework.Template.scope;
import static compiler.lib.template_framework.Template.let;
import static compiler.lib.template_framework.Template.$;
import compiler.lib.template_framework.library.CodeGenerationDataNameType;
import compiler.lib.template_framework.library.Expression;
import compiler.lib.template_framework.library.Expression.Nesting;
import compiler.lib.template_framework.library.Operations;
import compiler.lib.template_framework.library.TestFrameworkClass;
import compiler.lib.template_framework.library.PrimitiveType;
import compiler.lib.template_framework.library.VectorType;

/**
 * Basic Expression Fuzzer for Vector API.
 */
public class VectorExpressionFuzzer {
    private static final Random RANDOM = Utils.getRandomInstance();

    public static void main(String[] args) {
        // Create a new CompileFramework instance.
        CompileFramework comp = new CompileFramework();

        // Add a java source file.
        comp.addJavaSourceCode("compiler.vectorapi.templated.Templated", generate(comp));

        // Compile the source file.
        comp.compile("--add-modules=jdk.incubator.vector");

        List<String> vmArgs = new ArrayList<>(List.of(
            "--add-modules=jdk.incubator.vector",
            "--add-opens", "jdk.incubator.vector/jdk.incubator.vector=ALL-UNNAMED",
            "--add-opens", "java.base/java.lang=ALL-UNNAMED"
        ));
        vmArgs.addAll(Arrays.asList(args)); // Forward args
        String[] vmArgsArray = vmArgs.toArray(new String[0]);

        // compiler.vectorapi.templated.InnterTest.main(new String[] {});
        comp.invoke("compiler.vectorapi.templated.Templated", "main", new Object[] { vmArgsArray } );
    }

    // For Example 2 below.
    record TestArgument(
        Object defineAndFill,         // used once: to define and fill the input value
        Object passMethodArgument,    // used twice: to pass argument to test and reference
        Object receiveMethodArgument, // used twice: to receive argument in test and reference
        Object use) {}                // used twice: to use the value as expression argument

    // Generate a source Java file as String
    public static String generate(CompileFramework comp) {
        // Generate a list of test methods.
        List<TemplateToken> tests = new ArrayList<>();

        // We are going to use some random numbers in our tests, so import some good methods for that.
        tests.add(PrimitiveType.generateLibraryRNG());

        // Example 1:
        // To start simple, we just call the expression with the same constant arguments
        // every time. We can still compare the "gold" value obtained via interpreter with
        // later results obtained from the compiled method.
        var template1 = Template.make("type", (VectorType.Vector type) -> {
            // The depth determines roughly how many operations are going to be used in the expression.
            int depth = RANDOM.nextInt(1, 10);
            Expression expression = Expression.nestRandomly(type, Operations.ALL_OPERATIONS, depth, Nesting.EXACT);
            List<Object> expressionArguments = expression.argumentTypes.stream().map(CodeGenerationDataNameType::con).toList();
            return scope(
                """
                // --- $test start ---
                // Using $GOLD
                // type: #type

                static final Object $GOLD = $test();

                @Test
                public static Object $test() {
                    try {
                """,
                "        return ", expression.asToken(expressionArguments), ";\n",
                expression.info.exceptions.stream().map(exception ->
                    "        } catch (" + exception + " e) { return e;\n"
                ).toList(),
                """
                    } finally {
                        // Just so javac is happy if there are no exceptions to catch.
                    }
                }

                @Check(test = "$test")
                public static void $check(Object result) {
                """,
                expression.info.isResultDeterministic
                    ? "    Verify.checkEQ(result, $GOLD);\n"
                    : "    // result not deterministic - don't verify.\n",
                """
                }

                // --- $test end   ---
                """
            );
        });

        // Example 2:
        // Now a more complicated case. We want:
        // - For every invocation of the test method, we want to have different inputs for the arguments.
        // - We check correctness with a reference method that does the same but runs in the interpreter.
        // - Input values are delivered via fields or array loads.
        // - The final vector is written into an array, and that array is returned.
        var template2Body = Template.make("expression", "arguments", (Expression expression, List<Object> arguments) -> scope(
            let("elementType", ((VectorType.Vector)expression.returnType).elementType),
            """
            try {
            #elementType[] out = new #elementType[1000];
            """,
            expression.asToken(arguments), ".intoArray(out, 0);\n",
            "return out;\n",
            expression.info.exceptions.stream().map(exception ->
                "} catch (" + exception + " e) { return e;\n"
            ).toList(),
            """
            } finally {
                // Just javac is happy if there are no exceptions to catch.
            }
            """
        ));

        var template2 = Template.make("type", (VectorType.Vector type) -> {
            // The depth determines roughly how many operations are going to be used in the expression.
            int depth = RANDOM.nextInt(1, 10);
            Expression expression = Expression.nestRandomly(type, Operations.ALL_OPERATIONS, depth, Nesting.EXACT);

            List<TestArgument> arguments = new ArrayList<>();
            for (int i = 0; i < expression.argumentTypes.size(); i++) {
                String name = "arg_" + i;
                CodeGenerationDataNameType argumentType = expression.argumentTypes.get(i);
                switch(RANDOM.nextInt(4)) {
                    case 0 -> {
                        // Use the constant directly, no argument passing needed.
                        // To make the logic of passing arguments easy, we just pass null and receive an unused argument anyway.
                        arguments.add(new TestArgument(
                            "",
                            "null",
                            "Object unused_" + i,
                            argumentType.con()
                        ));
                    }
                    case 1 -> {
                        // Create the constant outside, and pass it.
                        arguments.add(new TestArgument(
                            List.of(argumentType.name(), " ", name, " = ", argumentType.con(), ";\n"),
                            name,
                            List.of(argumentType.name(), " ", name),
                            name
                        ));
                    }
                    default -> {
                        if (argumentType instanceof PrimitiveType t) {
                            // We can use the LibraryRGN to create a new value for the primitive in each
                            // invocation. We have to make sure to call the LibraryRNG in the "defineAndFill",
                            // so we get the same value for both test and reference. If we called LibraryRNG
                            // for "use", we would get separate values, which is not helpful.
                            arguments.add(new TestArgument(
                                List.of(t.name(), " ", name, " = ", t.callLibraryRNG(), ";\n"),
                                name,
                                List.of(t.name(), " ", name),
                                name
                            ));
                        } else if (argumentType instanceof VectorType.Vector t) {
                            PrimitiveType et = t.elementType;
                            arguments.add(new TestArgument(
                                List.of(et.name(), "[] ", name, " = new ", et.name(), "[1000];\n",
                                        "LibraryRNG.fill(", name,");\n"),
                                name,
                                List.of(et.name(), "[] ", name),
                                List.of(t.name(), ".fromArray(", t.speciesName, ", ", name, ", 0)")
                            ));
                        } else if (argumentType instanceof VectorType.Mask t) {
                            arguments.add(new TestArgument(
                                List.of("boolean[] ", name, " = new boolean[1000];\n",
                                        "LibraryRNG.fill(", name,");\n"),
                                name,
                                List.of("boolean[] ", name),
                                List.of("VectorMask.fromArray(", t.vectorType.speciesName, ", ", name, ", 0)")
                            ));
                        } else if (argumentType instanceof VectorType.Shuffle t) {
                            arguments.add(new TestArgument(
                                List.of("int[] ", name, " = new int[1000];\n",
                                        "LibraryRNG.fill(", name,");\n"),
                                name,
                                List.of("int[] ", name),
                                List.of("VectorShuffle.fromArray(", t.vectorType.speciesName, ", ", name, ", 0)")
                            ));
                        } else {
                            // We don't know anything special how to create different values
                            // each time, so let's just crate the same constant each time.
                            arguments.add(new TestArgument(
                                List.of(argumentType.name(), " ", name, " = ", argumentType.con(), ";\n"),
                                name,
                                List.of(argumentType.name(), " ", name),
                                name
                            ));
                        }
                    }
                }
            }

            // "join" together the arguments to make a comma separated list.
            List<Object> passArguments = IntStream.range(0, arguments.size() * 2 - 1).mapToObj(i ->
                (i % 2 == 0) ? arguments.get(i / 2).passMethodArgument() : ", "
            ).toList();
            List<Object> receiveArguments = IntStream.range(0, arguments.size() * 2 - 1).mapToObj(i ->
                (i % 2 == 0) ? arguments.get(i / 2).receiveMethodArgument() : ", "
            ).toList();

            List<Object> useArguments = arguments.stream().map(TestArgument::use).toList();

            return scope(
                """
                // --- $test start ---
                // type: #type

                @Run(test = "$test")
                public void $run() {
                """,
                arguments.stream().map(TestArgument::defineAndFill).toList(),
                """
                Object v0 = $test(
                """,
                passArguments,
                """
                );
                Object v1 = $reference(
                """,
                passArguments,
                """
                );
                """,
                expression.info.isResultDeterministic
                    ? "Verify.checkEQ(v0, v1);\n"
                    : "// result not deterministic - don't verify.\n",
                """
                }

                @Test
                public static Object $test(
                """,
                receiveArguments,
                """
                ) {
                """,
                template2Body.asToken(expression, useArguments),
                """
                }

                @DontCompile
                public static Object $reference(
                """,
                receiveArguments,
                """
                ) {
                """,
                template2Body.asToken(expression, useArguments),
                """
                }
                // --- $test end   ---
                """
            );
        });

        for (VectorType.Vector type : CodeGenerationDataNameType.VECTOR_VECTOR_TYPES) {
            tests.add(template1.asToken(type));
            tests.add(template2.asToken(type));
        }

        // Create the test class, which runs all tests.
        return TestFrameworkClass.render(
            // package and class name.
            "compiler.vectorapi.templated", "Templated",
            // Set of imports.
            Set.of("compiler.lib.verify.*",
                   "java.util.Random",
                   "jdk.test.lib.Utils",
                   "compiler.lib.generators.*",
                   "jdk.incubator.vector.*"),
            // classpath, so the Test VM has access to the compiled class files.
            comp.getEscapedClassPathOfCompiledClasses(),
            // The list of tests.
            tests);
    }
}