package p.studio.compiler.backend.irvm;

import org.junit.jupiter.api.Test;
import p.studio.compiler.backend.bytecode.BytecodeEmitter;
import p.studio.compiler.backend.bytecode.BytecodeModule;
import p.studio.compiler.backend.bytecode.BytecodePreloadVerifierService;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.ArrayList;
import java.util.List;

import static org.junit.jupiter.api.Assertions.*;

class OptimizeIRVMServiceTest {

    @Test
    void optimizeMustExecuteEnabledPassesInDeclaredOrder() {
        final var order = new ArrayList<String>();
        final var service = new OptimizeIRVMService(new IRVMValidator(), List.of(
                namedPass("A", order, true),
                namedPass("B", order, true)));

        final var output = service.optimize(validProgram());

        assertEquals(validProgram(), output);
        assertEquals(List.of("A", "B"), order);
        assertEquals(List.of("A", "B"), service.lastExecutedPassNames());
    }

    @Test
    void optimizeMustSkipDisabledPasses() {
        final var order = new ArrayList<String>();
        final var service = new OptimizeIRVMService(new IRVMValidator(), List.of(
                namedPass("A", order, false),
                namedPass("B", order, true)));

        service.optimize(validProgram());

        assertEquals(List.of("B"), order);
        assertEquals(List.of("B"), service.lastExecutedPassNames());
    }

    @Test
    void optimizeMustRejectPassThatMutatesVmProfile() {
        final var service = new OptimizeIRVMService(new IRVMValidator(), List.of(new OptimizeIRVMService.IRVMPass() {
            @Override
            public String name() {
                return "MutateProfile";
            }

            @Override
            public IRVMProgram apply(final IRVMProgram input) {
                return new IRVMProgram(new IRVMModule(
                        "experimental",
                        input.module().functions()));
            }
        }));

        final var thrown = assertThrows(IllegalArgumentException.class, () -> service.optimize(validProgram()));
        assertTrue(thrown.getMessage().contains("vm profile"));
    }

    @Test
    void optimizeDefaultPassesMustRemoveUnreachableInstructions() {
        final var service = new OptimizeIRVMService();
        final var input = singleFunctionProgram(
                ReadOnlyList.from(
                        new IRVMInstruction(IRVMOp.RET, null),
                        new IRVMInstruction(IRVMOp.HALT, null)),
                ReadOnlyList.from(
                        BytecodeEmitter.Operation.ret(),
                        BytecodeEmitter.Operation.halt()));

        final var optimized = service.optimize(input);

        assertNotEquals(input, optimized);
        assertEquals(1, optimized.module().functions().getFirst().instructions().size());
        assertEquals(IRVMOp.RET, optimized.module().functions().getFirst().instructions().getFirst().op());
        assertEquals(List.of(
                "UnreachableInstructionEliminationPass",
                "NormalizeRedundantJumpTargetsPass",
                "SimplifyJumpToNextPcPass"), service.lastExecutedPassNames());
    }

    @Test
    void normalizeRedundantJumpTargetsPassMustCollapseJumpChain() {
        final var service = new OptimizeIRVMService(new IRVMValidator(), List.of(
                OptimizeIRVMService.normalizeRedundantJumpTargetsPass()));
        final var input = singleFunctionProgram(
                ReadOnlyList.from(
                        new IRVMInstruction(IRVMOp.JMP, 6),
                        new IRVMInstruction(IRVMOp.JMP, 12),
                        new IRVMInstruction(IRVMOp.RET, null)),
                ReadOnlyList.from(
                        BytecodeEmitter.Operation.jmp(6, null),
                        BytecodeEmitter.Operation.jmp(12, null),
                        BytecodeEmitter.Operation.ret()));

        final var optimized = service.optimize(input);
        final var firstInstruction = optimized.module().functions().getFirst().instructions().getFirst();
        final var firstOperation = optimized.emissionPlan().functions().getFirst().operations().getFirst();

        assertEquals(IRVMOp.JMP, firstInstruction.op());
        assertEquals(12, firstInstruction.immediate());
        assertEquals(BytecodeEmitter.OperationKind.JMP, firstOperation.kind());
        assertEquals(12, firstOperation.immediate());
    }

    @Test
    void simplifyJumpToNextPcPassMustRemoveDirectFallthroughJump() {
        final var service = new OptimizeIRVMService(new IRVMValidator(), List.of(
                OptimizeIRVMService.simplifyJumpToNextPcPass()));
        final var input = singleFunctionProgram(
                ReadOnlyList.from(
                        new IRVMInstruction(IRVMOp.JMP, 6),
                        new IRVMInstruction(IRVMOp.RET, null)),
                ReadOnlyList.from(
                        BytecodeEmitter.Operation.jmp(6, null),
                        BytecodeEmitter.Operation.ret()));

        final var optimized = service.optimize(input);

        assertEquals(1, optimized.module().functions().getFirst().instructions().size());
        assertEquals(IRVMOp.RET, optimized.module().functions().getFirst().instructions().getFirst().op());
        assertEquals(1, optimized.emissionPlan().functions().getFirst().operations().size());
        assertEquals(BytecodeEmitter.OperationKind.RET, optimized.emissionPlan().functions().getFirst().operations().getFirst().kind());
    }

    @Test
    void optimizeDefaultPassesMustEliminateUnreachableInternalExtensionBeforeEmission() {
        final var service = new OptimizeIRVMService();
        final var input = new IRVMProgram(
                new IRVMModule(
                        "experimental-v1",
                        ReadOnlyList.from(new IRVMFunction(
                                "main",
                                0,
                                0,
                                0,
                                1,
                                ReadOnlyList.from(
                                        new IRVMInstruction(IRVMOp.JMP, 8),
                                        new IRVMInstruction(IRVMOp.INTERNAL_EXT, null),
                                        new IRVMInstruction(IRVMOp.RET, null))))),
                new BytecodeEmitter.EmissionPlan(
                        0,
                        ReadOnlyList.empty(),
                        ReadOnlyList.empty(),
                        ReadOnlyList.from(new BytecodeEmitter.FunctionPlan(
                                "main",
                                0,
                                0,
                                0,
                                1,
                                ReadOnlyList.from(
                                        BytecodeEmitter.Operation.jmp(8, null),
                                        BytecodeEmitter.Operation.halt(),
                                        BytecodeEmitter.Operation.ret())))));

        final var optimized = service.optimize(input);
        final var function = optimized.module().functions().getFirst();
        final var firstInstruction = function.instructions().getFirst();
        final var firstOperation = optimized.coherentEmissionPlan().functions().getFirst().operations().getFirst();

        assertEquals(1, function.instructions().size());
        assertFalse(optimized.hasInternalOpcodes());
        assertEquals(IRVMOp.RET, firstInstruction.op());
        assertEquals(BytecodeEmitter.OperationKind.RET, firstOperation.kind());

        final var emitted = new BytecodeEmitter().emit(optimized.coherentEmissionPlan());
        new BytecodePreloadVerifierService().verify(emitted);
    }

    @Test
    void unreachableInstructionEliminationPassMustPreserveJumpSpanWhenTargetIsRemapped() {
        final var jumpSpan = new BytecodeModule.SourceSpan(7, 10, 12);
        final var removedSpan = new BytecodeModule.SourceSpan(7, 13, 14);
        final var retSpan = new BytecodeModule.SourceSpan(7, 15, 16);
        final var service = new OptimizeIRVMService(new IRVMValidator(), List.of(
                OptimizeIRVMService.unreachableInstructionEliminationPass()));
        final var input = singleFunctionProgram(
                ReadOnlyList.from(
                        new IRVMInstruction(IRVMOp.JMP, 8),
                        new IRVMInstruction(IRVMOp.HALT, null),
                        new IRVMInstruction(IRVMOp.RET, null)),
                ReadOnlyList.from(
                        BytecodeEmitter.Operation.jmp(8, jumpSpan),
                        BytecodeEmitter.Operation.halt(removedSpan),
                        BytecodeEmitter.Operation.ret(retSpan)));

        final var optimized = service.optimize(input);
        final var firstOperation = optimized.coherentEmissionPlan().functions().getFirst().operations().getFirst();
        final var secondOperation = optimized.coherentEmissionPlan().functions().getFirst().operations().get(1);

        assertEquals(2, optimized.module().functions().getFirst().instructions().size());
        assertEquals(BytecodeEmitter.OperationKind.JMP, firstOperation.kind());
        assertEquals(6, firstOperation.immediate());
        assertEquals(jumpSpan, firstOperation.span());
        assertEquals(retSpan, secondOperation.span());
    }

    @Test
    void normalizeRedundantJumpTargetsPassMustPreserveJumpSpanAfterRewrite() {
        final var firstJumpSpan = new BytecodeModule.SourceSpan(9, 20, 24);
        final var secondJumpSpan = new BytecodeModule.SourceSpan(9, 25, 28);
        final var service = new OptimizeIRVMService(new IRVMValidator(), List.of(
                OptimizeIRVMService.normalizeRedundantJumpTargetsPass()));
        final var input = singleFunctionProgram(
                ReadOnlyList.from(
                        new IRVMInstruction(IRVMOp.JMP, 6),
                        new IRVMInstruction(IRVMOp.JMP, 12),
                        new IRVMInstruction(IRVMOp.RET, null)),
                ReadOnlyList.from(
                        BytecodeEmitter.Operation.jmp(6, firstJumpSpan),
                        BytecodeEmitter.Operation.jmp(12, secondJumpSpan),
                        BytecodeEmitter.Operation.ret()));

        final var optimized = service.optimize(input);
        final var firstOperation = optimized.coherentEmissionPlan().functions().getFirst().operations().getFirst();

        assertEquals(BytecodeEmitter.OperationKind.JMP, firstOperation.kind());
        assertEquals(12, firstOperation.immediate());
        assertEquals(firstJumpSpan, firstOperation.span());
    }

    @Test
    void simplifyJumpToNextPcPassMustPreserveSpanOnSurvivingOperation() {
        final var retSpan = new BytecodeModule.SourceSpan(11, 40, 41);
        final var service = new OptimizeIRVMService(new IRVMValidator(), List.of(
                OptimizeIRVMService.simplifyJumpToNextPcPass()));
        final var input = singleFunctionProgram(
                ReadOnlyList.from(
                        new IRVMInstruction(IRVMOp.JMP, 6),
                        new IRVMInstruction(IRVMOp.RET, null)),
                ReadOnlyList.from(
                        BytecodeEmitter.Operation.jmp(6, new BytecodeModule.SourceSpan(11, 35, 39)),
                        BytecodeEmitter.Operation.ret(retSpan)));

        final var optimized = service.optimize(input);
        final var onlyOperation = optimized.coherentEmissionPlan().functions().getFirst().operations().getFirst();

        assertEquals(1, optimized.module().functions().getFirst().instructions().size());
        assertEquals(BytecodeEmitter.OperationKind.RET, onlyOperation.kind());
        assertEquals(retSpan, onlyOperation.span());
    }

    private OptimizeIRVMService.IRVMPass namedPass(
            final String name,
            final List<String> order,
            final boolean enabled) {
        return new OptimizeIRVMService.IRVMPass() {
            @Override
            public String name() {
                return name;
            }

            @Override
            public boolean enabled() {
                return enabled;
            }

            @Override
            public IRVMProgram apply(final IRVMProgram input) {
                order.add(name);
                return input;
            }
        };
    }

    private IRVMProgram validProgram() {
        return singleFunctionProgram(
                ReadOnlyList.from(new IRVMInstruction(IRVMOp.HALT, null)),
                ReadOnlyList.from(BytecodeEmitter.Operation.halt()));
    }

    private IRVMProgram singleFunctionProgram(
            final ReadOnlyList<IRVMInstruction> instructions,
            final ReadOnlyList<BytecodeEmitter.Operation> operations) {
        return new IRVMProgram(
                new IRVMModule(
                        "core-v1",
                        ReadOnlyList.from(new IRVMFunction(
                                "main",
                                0,
                                0,
                                0,
                                1,
                                instructions))),
                new BytecodeEmitter.EmissionPlan(
                        0,
                        ReadOnlyList.empty(),
                        ReadOnlyList.empty(),
                        ReadOnlyList.from(new BytecodeEmitter.FunctionPlan(
                                "main",
                                0,
                                0,
                                0,
                                1,
                                operations))));
    }
}
