path: root/src/main/java/org/shedlang/compiler/backends/wasm/PatioToWasm.java

package org.shedlang.compiler.backends.wasm;

import org.pcollections.PMap;
import org.pcollections.PStack;
import org.pcollections.PVector;
import org.shedlang.compiler.backends.targets.WasmTarget;
import org.shedlang.compiler.backends.wasm.structure.ast.instructions.*;
import org.shedlang.compiler.backends.wasm.structure.ast.linking.WasmMemoryAddress;
import org.shedlang.compiler.backends.wasm.structure.ast.modules.*;
import org.shedlang.compiler.backends.wasm.structure.ast.types.*;
import org.shedlang.compiler.patio.*;
import org.shedlang.compiler.patio.instructions.*;
import org.shedlang.compiler.patio.operands.*;
import org.shedlang.compiler.patio.types.*;
import org.shedlang.compiler.patio.values.*;
import org.shedlang.compiler.util.Lists;
import org.shedlang.compiler.util.P;

import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.List;

public class PatioToWasm {
    private static final String INDIRECT_FUNCTION_TABLE_NAME = "__indirect_function_table";

    public static WasmModule generateCompilationUnit(
        PatioCompilationUnit patioCompilationUnit,
        WasmTarget target
    ) {
        var module = WasmModule.EMPTY.addImport(
            new WasmImport(
                "env",
                INDIRECT_FUNCTION_TABLE_NAME,
                new WasmImportDescriptor.Table(
                    INDIRECT_FUNCTION_TABLE_NAME,
                    new WasmTableType(WasmReferenceType.FUNC_REF, WasmLimits.min32(0))
                )
            )
        );

        var wasmMemoryBuilder = WasmMemoryBuilder.empty(target.indexType());

        for (var patioCompilationUnitElement : patioCompilationUnit.elements()) {
            switch (patioCompilationUnitElement) {
                case PatioDataElementImmutableDefinition patioDataElementImmutableDefinition -> {
                    wasmMemoryBuilder = wasmMemoryBuilder.add(
                        patioDataElementImmutableDefinition.name(),
                        patioDataValueToBytes(patioDataElementImmutableDefinition.value(), target),
                        WasmAlignment.ALIGN_32_BIT
                    ).memory();
                }
                case PatioDataElementMutableDeclaration patioDataElementMutableDeclaration -> {
                    // TODO: how to test this?
                    module = module.addUndefinedDataElement(patioDataElementMutableDeclaration.name());
                }
                case PatioDataElementMutableDefinition patioDataElementMutableDefinition -> {
                    wasmMemoryBuilder = wasmMemoryBuilder.add(
                        patioDataElementMutableDefinition.name(),
                        patioDataValueToBytes(patioDataElementMutableDefinition.initialValue(), target),
                        WasmAlignment.ALIGN_32_BIT
                    ).memory();
                }
                case PatioFunctionDeclaration patioFunctionDeclaration -> {
                    var name = patioFunctionDeclaration.name();

                    var wasmFuncType = new WasmFuncType(
                        patioFunctionDeclaration.params().stream()
                            .map(param -> patioTypeToWasmValueType(param.type(), target))
                            .toList(),
                        patioFunctionDeclaration.returnTypes().stream()
                            .map(returnType -> patioTypeToWasmValueType(returnType, target))
                            .toList()
                    );

                    module = module.addImport(new WasmImport(
                        "env",
                        name,
                        new WasmImportDescriptor.Func(name, wasmFuncType)
                    ));
                }
                case PatioExceptionTypeDefinition patioExceptionTypeDefinition -> {
                    module = module.addTag(new WasmTag(
                        WasmTagType.EXCEPTION,
                        patioExceptionTypeDefinition.name(),
                        List.of(target.patioIntType())
                    ));
                }
                case PatioFunctionDefinition patioFunctionDefinition -> {
                    var name = patioFunctionDefinition.name();

                    var wasmParams = new ArrayList<WasmLocal>();
                    var wasmLocals = new ArrayList<WasmLocal>();

                    for (var patioLocal : patioFunctionDefinition.locals()) {
                        var wasmLocal = patioLocalToWasmLocal(patioLocal, target);
                        switch (patioLocal.localType()) {
                            case PARAM -> {
                                wasmParams.add(wasmLocal);
                            }
                            case VAL, VAR -> {
                                wasmLocals.add(wasmLocal);
                            }
                        }
                    }

                    var wasmResultTypes = patioFunctionDefinition.returnTypes().stream()
                        .map(returnType -> patioTypeToWasmValueType(returnType, target))
                        .toList();

                    var context = generateInstructions(
                        patioFunctionDefinition.body(),
                        Context.forFunction(target, module, Lists.concat(wasmParams, wasmLocals))
                    );

                    module = context.module;

                    module = module.addFunction(new WasmFunction(
                        name,
                        wasmParams,
                        wasmLocals,
                        wasmResultTypes,
                        new WasmExpression(context.instructions)
                    ));
                }
            }
        }

        module = wasmMemoryBuilder.toModule(module);

        return module;
    }

    private static Context generateInstructions(
        List<PatioInstruction> patioInstructions,
        Context context
    ) {
        for (var patioInstruction : patioInstructions) {
            context = generateInstruction(patioInstruction, context);
        }

        return context;
    }

    private static Context generateInstruction(
        PatioInstruction patioInstruction,
        Context context
    ) {
        var target = context.target;

        return switch (patioInstruction) {
            case PatioBlock patioBlock ->
                generateInstructions(patioBlock.body(), context);

            case PatioBoolToInt patioBoolToInt -> {
                context = generateOperand(patioBoolToInt.value(), context);

                if (target == WasmTarget.WASM64) {
                    context = context.addInstruction(new WasmI64ExtendI32Unsigned());
                }

                context = popIntoPatioLocal(patioBoolToInt.destination(), context);

                yield context;
            }

            case PatioCall patioCall -> {
                for (var patioArg : patioCall.args()) {
                    context = generateOperand(patioArg, context);
                }

                context = context.addInstruction(new WasmCall(new WasmFuncName(patioCall.functionName())));

                for (var patioDestination : patioCall.destinations().reversed()) {
                    context = popIntoPatioLocal(patioDestination, context);
                }

                yield context;
            }

            case PatioCallPointer patioCallPointer -> {
                for (var patioArg : patioCallPointer.args()) {
                    context = generateOperand(patioArg, context);
                }

                context = generateOperand(patioCallPointer.functionPointer(), context);

                var wasmFuncType = new WasmFuncType(
                    patioOperandsToWasmValueTypes(patioCallPointer.args(), context),
                    patioOperandsToWasmValueTypes(patioCallPointer.destinations(), context)
                );

                context = context.addType(wasmFuncType);

                context = context.addInstruction(new WasmCallIndirect(
                    new WasmTableName(INDIRECT_FUNCTION_TABLE_NAME),
                    wasmFuncType
                ));

                for (var patioDestination : patioCallPointer.destinations().reversed()) {
                    context = popIntoPatioLocal(patioDestination, context);
                }

                yield context;
            }

            case PatioDataPtrToInt patioDataPtrToInt ->
                generateOperandToPatioLocal(
                    patioDataPtrToInt.destination(),
                    patioDataPtrToInt.value(),
                    context
                );

            case PatioIf patioIf -> {
                context = generateOperand(patioIf.condition(), context);

                var wasmBlockTypeResult = generateBlockType(patioIf.destinations(), context);
                var wasmBlockType = wasmBlockTypeResult.value();
                context = wasmBlockTypeResult.context();

                context = context.enterIfInstruction(wasmBlockType);
                context = generateBlockBody(patioIf.trueBody(), context);

                context = context.enterElse();
                context = generateBlockBody(patioIf.falseBody(), context);

                context = context.exitIfInstruction();

                context = popIntoPatioLocals(patioIf.destinations(), context);

                yield context;
            }

            case PatioIntAdd patioIntAdd ->
                generateBinaryOperation(
                    patioIntAdd,
                    wasmIntAdd(target.patioIntType()),
                    context
                );

            case PatioIntEquals patioIntEquals ->
                generateBinaryOperation(
                    patioIntEquals,
                    wasmIntEq(target.patioIntType()),
                    context
                );

            case PatioIntNotEqual patioIntNotEqual ->
                generateBinaryOperation(
                    patioIntNotEqual,
                    wasmIntNe(target.patioIntType()),
                    context
                );

            case PatioIntSub patioIntSub ->
                generateBinaryOperation(
                    patioIntSub,
                    wasmIntSub(target.patioIntType()),
                    context
                );

            case PatioIntToBool patioIntToBool -> {
                context = generateOperand(patioIntToBool.value(), context);

                if (target == WasmTarget.WASM64) {
                    context = context.addInstruction(new WasmI32WrapI64());
                }

                context = popIntoPatioLocal(patioIntToBool.destination(), context);

                yield context;
            }

            case PatioIntToDataPtr patioIntToDataPtr ->
                generateOperandToPatioLocal(
                    patioIntToDataPtr.destination(),
                    patioIntToDataPtr.value(),
                    context
                );

            case PatioLocalStore patioLocalStore ->
                generateOperandToPatioLocal(
                    patioLocalStore.destination(),
                    patioLocalStore.value(),
                    context
                );

            case PatioLoop patioLoop -> {
                context = context.enterBlockInstruction();
                context = context.enterLoopInstruction();
                context = generateInstructions(patioLoop.body(), context);
                context = context.exitLoopInstruction();
                context = context.exitBlockInstruction();

                yield context;
            }

            case PatioLoopBreak patioLoopBreak -> {
                context = context.addBrToBlockInstruction();

                yield context;
            }

            case PatioLoopContinue patioLoopContinue -> {
                context = context.addBrToLoopInstruction();

                yield context;
            }

            case PatioMemoryLoad patioMemoryLoad -> {
                context = generateOperand(patioMemoryLoad.pointer(), context);

                context = context.addInstruction(wasmIntLoad(
                    context.wasmLocalType(patioMemoryLoad.destination()),
                    wasmMemoryArg(
                        patioMemoryLoad.layout(),
                        patioMemoryLoad.elementOffsets(),
                        target
                    )
                ));

                context = popIntoPatioLocal(patioMemoryLoad.destination(), context);

                yield context;
            }

            case PatioMemoryStore patioMemoryStore -> {
                context = generateOperand(patioMemoryStore.pointer(), context);

                context = generateOperand(patioMemoryStore.value(), context);

                context = context.addInstruction(wasmIntStore(
                    patioOperandToWasmValueType(patioMemoryStore.value(), context),
                    wasmMemoryArg(
                        patioMemoryStore.layout(),
                        patioMemoryStore.elementOffsets(),
                        target
                    )
                ));

                yield context;
            }

            case PatioReturn patioReturn -> {
                for (var value : patioReturn.values()) {
                    context = generateOperand(value, context);
                }

                context = context.addInstruction(new WasmReturn());

                yield context;
            }

            case PatioSwitch patioSwitch -> {
                // TODO: use br_table?

                var wasmBlockTypeResult = generateBlockType(patioSwitch.destinations(), context);
                var wasmBlockType = wasmBlockTypeResult.value();
                context = wasmBlockTypeResult.context();

                for (var patioSwitchCase : patioSwitch.cases()) {
                    context = generateOperand(patioSwitch.value(), context);
                    context = context.addInstruction(wasmIntConst(
                        patioOperandToWasmValueType(patioSwitch.value(), context),
                        patioSwitchCase.value()
                    ));
                    context = context.addInstruction(wasmIntEq(
                        patioOperandToWasmValueType(patioSwitch.value(), context)
                    ));
                    context = context.enterIfInstruction(wasmBlockType);
                    context = generateBlockBody(patioSwitchCase.body(), context);
                    context = context.enterElse();
                }

                context = generateBlockBody(patioSwitch.defaultBody(), context);

                for (var patioSwitchCase : patioSwitch.cases()) {
                    context = context.exitIfInstruction();
                }

                context = popIntoPatioLocals(patioSwitch.destinations(), context);

                yield context;
            }

            case PatioThrow throw_ -> {
                context = generateOperand(throw_.value(), context);

                context = context.addInstruction(new WasmThrow(new WasmTagName(throw_.exceptionTypeName())));

                yield context;
            }

            case PatioTry try_ -> {
                var wasmBlockTypeResult = generateBlockType(try_.destinations(), context);
                var wasmBlockType = wasmBlockTypeResult.value();
                context = wasmBlockTypeResult.context();

                context = context.enterTryInstruction(wasmBlockType);

                context = generateBlockBody(try_.body(), context);

                for (var catchClause : try_.catchClauses()) {
                    context = context.enterCatch(new WasmTagName(catchClause.exceptionTypeName()));

                    context = popIntoPatioLocal(catchClause.destination(), context);
                    context = generateBlockBody(catchClause.body(), context);
                }

                context = context.exitTryInstruction();

                context = popIntoPatioLocals(try_.destinations(), context);

                yield context;
            }

            case PatioUnreachable patioUnreachable ->
                context.addInstruction(new WasmUnreachable());

            case PatioVarLoad patioVarLoad -> {
                yield generateOperandToPatioLocal(
                    patioVarLoad.destination(),
                    patioVarLoad.var(),
                    context
                );
            }

            case PatioVarStore patioVarStore -> {
                yield generateOperandToPatioLocal(
                    patioVarStore.var(),
                    patioVarStore.value(),
                    context
                );
            }
        };
    }

    private static Context.Result<WasmBlockType> generateBlockType(
        List<PatioLocalRef> destinations,
        Context context
    ) {
        if (destinations.isEmpty()) {
            return context.toResult(WasmBlockType.EMPTY);
        } else if (destinations.size() == 1) {
            var valueType = patioOperandToWasmValueType(destinations.get(0), context);
            return context.toResult(WasmBlockType.of(valueType));
        } else {
            var resultTypes = destinations.stream()
                .map(patioTarget -> patioOperandToWasmValueType(patioTarget, context))
                .toList();

            var funcType = new WasmFuncType(List.of(), resultTypes);

            return context.addType(funcType).toResult(WasmBlockType.of(funcType));
        }

    }

    private static Context generateBlockBody(
        PatioBlock block,
        Context context
    ) {
        context = generateInstructions(block.body(), context);

        for (var patioResult : block.results()) {
            context = generateOperand(patioResult, context);
        }

        return context;
    }

    private static Context generateBinaryOperation(
        PatioInstructionBinaryOperation patioInstruction,
        WasmInstruction wasmInstruction,
        Context context
    ) {
        context = generateOperand(patioInstruction.left(), context);

        context = generateOperand(patioInstruction.right(), context);

        context = context.addInstruction(wasmInstruction);

        context = popIntoPatioLocal(patioInstruction.destination(), context);

        return context;
    }

    private static Context generateOperandToPatioLocal(
        PatioLocalRef destination,
        PatioOperand value,
        Context context
    ) {
        context = generateOperand(value, context);

        context = popIntoPatioLocal(destination, context);

        return context;
    }

    private static Context popIntoPatioLocal(PatioLocalRef destination, Context context) {
        var destinationWasmLocalIndex = context.wasmLocalIndex(destination);
        return context.addInstruction(new WasmLocalSet(destinationWasmLocalIndex));
    }

    private static Context popIntoPatioLocals(
        List<PatioLocalRef> destinations,
        Context context
    ) {
        for (var patioTarget : destinations.reversed()) {
            context = popIntoPatioLocal(patioTarget, context);
        }
        return context;
    }

    private static Context generateOperand(
        PatioOperand patioOperand,
        Context context
    ) {
        return switch (patioOperand) {
            case PatioBoolConst patioBoolConst ->
                context.addInstruction(new WasmI32Const(patioBoolConst.value() ? 1 : 0));

            case PatioByteConst patioByteConst ->
                context.addInstruction(new WasmI32Const(Byte.toUnsignedInt(patioByteConst.value())));

            case PatioDataPtrConst patioDataPtrConst ->
                context.addInstruction(generateDataPtrValue(
                    patioDataPtrConst.value(),
                    context.target
                ));

            case PatioFunctionPtrConst patioFunctionPtrConst -> {
                var wasmFuncName = new WasmFuncName(patioFunctionPtrConst.name());
                context = context
                    .addInstruction(new WasmTableElementIndexConst(wasmFuncName));

                if (context.target == WasmTarget.WASM64) {
                    context = context.addInstruction(new WasmI64ExtendI32Unsigned());
                }

                yield context;
            }

            case PatioIntConst patioIntConst ->
                context.addInstruction(wasmIntConst(
                    context.target.patioIntType(),
                    generateIntValue(patioIntConst.value(), context.target)
                ));

            case PatioInt32Const patioInt32Const ->
                context.addInstruction(new WasmI32Const(
                    generateIntValue(patioInt32Const.value(), context.target)
                ));

            case PatioLocalRef patioLocalRef -> {
                var wasmLocalIndex = context.wasmLocalIndex(patioLocalRef);
                yield context.addInstruction(new WasmLocalGet(wasmLocalIndex));
            }
        };
    }

    private static WasmInstruction generateDataPtrValue(PatioDataPtrValue value, WasmTarget target) {
        return switch (value) {
            case PatioDataPtrValueDataElement dataElement -> {
                var wasmMemoryAddress = new WasmMemoryAddress(dataElement.name());
                yield new WasmMemoryAddressConst(
                    target.indexType(),
                    wasmMemoryAddress
                );
            }

            case PatioDataPtrValueNull patioDataPtrValueNull -> {
                yield wasmIntConst(target.dataPtrType(), 0);
            }
        };
    }

    private static int generateIntValue(PatioIntValue value, WasmTarget target) {
        return switch (value) {
            case PatioIntValueByteOffsetOf byteOffsetOf ->
                offsetOf(
                    byteOffsetOf.layout(),
                    byteOffsetOf.elementOffsets(),
                    target
                );

            case PatioIntValueByteSizeOf byteSizeOf ->
                sizeOf(
                    byteSizeOf.layout(),
                    target
                );

            case PatioIntValueLiteral literal ->
                literal.value();
        };
    }

    private static int sizeOf(PatioDataLayout layout, WasmTarget target) {
        return switch (layout) {
            case PatioDataLayoutArray patioDataLayoutArray ->
                sizeOf(patioDataLayoutArray.element(), target) * patioDataLayoutArray.size();

            case PatioDataLayoutBool patioDataLayoutBool ->
                1;

            case PatioDataLayoutByte patioDataLayoutByte ->
                1;

            case PatioDataLayoutDataPtr patioDataLayoutDataPtr ->
                target.dataPtrType().sizeBytes();

            case PatioDataLayoutFunctionPtr patioDataLayoutFunctionPtr ->
                target.functionPtrType().sizeBytes();

            case PatioDataLayoutInt patioDataLayoutInt ->
                target.patioIntType().sizeBytes();

            case PatioDataLayoutInt32 patioDataLayoutInt32 ->
                WasmValueType.I32.sizeBytes();

            case PatioDataLayoutStruct patioDataLayoutStruct ->
                patioDataLayoutStruct.elements().stream()
                    .map(element -> sizeOf(element, target))
                    .reduce(0, Integer::sum);
        };
    }

    private static int offsetOf(
        PatioDataLayout patioDataLayout,
        List<Integer> patioOffsets,
        WasmTarget target
    ) {
        var wasmOffset = 0;

        for (var patioOffset : patioOffsets) {
            switch (patioDataLayout) {
                case PatioDataLayoutArray patioDataLayoutArray -> {
                    var elementSize = sizeOf(patioDataLayoutArray.element(), target);
                    wasmOffset += elementSize * patioOffset;
                    patioDataLayout = patioDataLayoutArray.element();
                }
                case PatioDataLayoutBool patioDataLayoutBool -> {
                    throw new RuntimeException("offset not allowed for bools");
                }
                case PatioDataLayoutByte patioDataLayoutByte -> {
                    throw new RuntimeException("offset not allowed for bytes");
                }
                case PatioDataLayoutDataPtr patioDataLayoutDataPtr -> {
                    throw new RuntimeException("offset not allowed for data ptrs");
                }
                case PatioDataLayoutFunctionPtr patioDataLayoutFunctionPtr -> {
                    throw new RuntimeException("offset not allowed for function ptrs");
                }
                case PatioDataLayoutInt patioDataLayoutInt -> {
                    throw new RuntimeException("offset not allowed for ints");
                }
                case PatioDataLayoutInt32 patioDataLayoutInt32 -> {
                    throw new RuntimeException("offset not allowed for int32s");
                }
                case PatioDataLayoutStruct patioDataLayoutStruct -> {
                    for (var i = 0; i < patioOffset; i++) {
                        var element = patioDataLayoutStruct.elements().get(i);
                        wasmOffset += sizeOf(element, target);
                    }
                    patioDataLayout = patioDataLayoutStruct.elements().get(patioOffset);
                }
            }
        }

        return wasmOffset;
    }

    private static WasmMemoryArg wasmMemoryArg(
        PatioDataLayout patioDataLayout,
        List<Integer> patioOffsets,
        WasmTarget target
    ) {
        var wasmOffset = offsetOf(patioDataLayout, patioOffsets, target);

        // TODO: alignment
        return new WasmMemoryArg(wasmOffset, WasmAlignment.NONE);
    }

    private static byte[] patioDataValueToBytes(
        PatioDataValue value,
        WasmTarget target
    ) {
        var bytes = ByteBuffer.allocate(sizeOf(value, target));
        bytes.order(ByteOrder.LITTLE_ENDIAN);

        patioDataValueToBytes(value, bytes, target);

        return bytes.array();
    }

    private static void patioDataValueToBytes(
        PatioDataValue value,
        ByteBuffer bytes,
        WasmTarget target
    ) {
        switch (value) {
            case PatioDataValueArray patioDataValueArray -> {
                for (var element : patioDataValueArray.elements()) {
                    patioDataValueToBytes(element, bytes, target);
                }
            }
            case PatioDataValueBool patioDataValueBool -> {
                bytes.put(patioDataValueBool.value() ? (byte) 1 : (byte) 0);
            }
            case PatioDataValueByte patioDataValueByte -> {
                bytes.put(patioDataValueByte.value());
            }
            case PatioDataValueDataPtr patioDataValueDataPtr -> {
                switch (patioDataValueDataPtr.value()) {
                    case PatioDataPtrValueDataElement dataElement -> {
                        throw new UnsupportedOperationException("TODO");
                    }
                    case PatioDataPtrValueNull null_ -> {
                        bytes.putInt(0);
                    }
                };
            }
            case PatioDataValueInt patioDataValueInt -> {
                var intValue = generateIntValue(patioDataValueInt.value(), target);
                switch (target) {
                    case WASM32 -> {
                        bytes.putInt(intValue);
                    }
                    case WASM64 -> {
                        bytes.putLong(intValue);
                    }
                }
            }
            case PatioDataValueInt32 patioDataValueInt32 -> {
                var intValue = generateIntValue(patioDataValueInt32.value(), target);
                bytes.putInt(intValue);
            }
            case PatioDataValueStruct patioDataValueStruct -> {
                for (var element : patioDataValueStruct.elements()) {
                    patioDataValueToBytes(element, bytes, target);
                }
            }
        }
    }

    private static int sizeOf(PatioDataValue value, WasmTarget target) {
        return sizeOf(value.layout(), target);
    }

    private static WasmLocal patioLocalToWasmLocal(
        PatioLocal local,
        WasmTarget target
    ) {
        return new WasmLocal(local.name(), patioTypeToWasmValueType(local.type(), target));
    }

    private static List<WasmValueType> patioOperandsToWasmValueTypes(
        List<? extends PatioOperand> operands,
        Context context
    ) {
        return operands.stream()
            .map(operand -> patioOperandToWasmValueType(operand, context))
            .toList();
    }

    private static WasmValueType patioOperandToWasmValueType(
        PatioOperand operand,
        Context context
    ) {
        return switch (operand) {
            case PatioOperandConst patioOperandConst ->
                patioTypeToWasmValueType(patioOperandConst.type(), context.target);
            case PatioLocalRef patioLocalRef ->
                // TODO: push into PatioOperand?
                context.wasmLocalType(patioLocalRef);
        };
    }

    private static WasmValueType patioTypeToWasmValueType(
        PatioType type,
        WasmTarget target
    ) {
        return switch (type) {
            case PatioBoolType patioBoolType ->
                WasmValueType.I32;
            case PatioByteType patioByteType ->
                WasmValueType.I32;
            case PatioDataPtrType patioDataPtrType ->
                target.dataPtrType();
            case PatioFunctionPtrType patioFunctionPtrType ->
                target.functionPtrType();
            case PatioIntType patioIntType ->
                target.patioIntType();
            case PatioInt32Type patioInt32Type ->
                WasmValueType.I32;
        };
    }

    private static WasmInstruction wasmIntConst(WasmValueType wasmValueType, int value) {
        return switch (wasmValueType) {
            case WasmValueType.Number number -> {
                yield switch (number.type()) {
                    case I32 -> new WasmI32Const(value);
                    case I64 -> new WasmI64Const(value);
                    case F32, F64 -> {
                        throw new IllegalArgumentException("Cannot create int const for " + number);
                    }
                };
            }
        };
    }

    private static WasmInstruction wasmIntEq(WasmValueType wasmValueType) {
        if (wasmValueType == WasmValueType.I32) {
            return new WasmI32Eq();
        } else if (wasmValueType == WasmValueType.I64) {
            return new WasmI64Eq();
        } else {
            throw new IllegalArgumentException("Cannot create int eq for " + wasmValueType);
        }
    }

    private static WasmInstruction wasmIntNe(WasmValueType wasmValueType) {
        if (wasmValueType == WasmValueType.I32) {
            return new WasmI32Ne();
        } else if (wasmValueType == WasmValueType.I64) {
            return new WasmI64Ne();
        } else {
            throw new IllegalArgumentException("Cannot create int ne for " + wasmValueType);
        }
    }

    private static WasmInstruction wasmIntAdd(WasmValueType wasmValueType) {
        if (wasmValueType == WasmValueType.I32) {
            return new WasmI32Add();
        } else if (wasmValueType == WasmValueType.I64) {
            return new WasmI64Add();
        } else {
            throw new IllegalArgumentException("Cannot create int add for " + wasmValueType);
        }
    }

    private static WasmInstruction wasmIntSub(WasmValueType wasmValueType) {
        if (wasmValueType == WasmValueType.I32) {
            return new WasmI32Sub();
        } else if (wasmValueType == WasmValueType.I64) {
            return new WasmI64Sub();
        } else {
            throw new IllegalArgumentException("Cannot create int sub for " + wasmValueType);
        }
    }

    private static WasmInstruction wasmIntLoad(
        WasmValueType wasmValueType,
        WasmMemoryArg memoryArg
    ) {
        if (wasmValueType == WasmValueType.I32) {
            return new WasmI32Load(memoryArg);
        } else if (wasmValueType == WasmValueType.I64) {
            return new WasmI64Load(memoryArg);
        } else {
            throw new IllegalArgumentException("Cannot create int load for " + wasmValueType);
        }
    }

    private static WasmInstruction wasmIntStore(
        WasmValueType wasmValueType,
        WasmMemoryArg memoryArg
    ) {
        if (wasmValueType == WasmValueType.I32) {
            return new WasmI32Store(memoryArg);
        } else if (wasmValueType == WasmValueType.I64) {
            return new WasmI64Store(memoryArg);
        } else {
            throw new IllegalArgumentException("Cannot create int store for " + wasmValueType);
        }
    }

    private static class Context {
        private record PatioLocalInfo(
            WasmLocalIndex wasmLocalIndex,
            WasmValueType wasmType
        ) {}

        public static Context forFunction(
            WasmTarget target,
            WasmModule module,
            List<WasmLocal> locals
        ) {
            var patioLocalNameToInfo = P.<String, PatioLocalInfo>map();

            for (var local : locals) {
                var wasmLocalIndex = new WasmLocalIndex(patioLocalNameToInfo.size());
                var localInfo = new PatioLocalInfo(wasmLocalIndex, local.type());
                patioLocalNameToInfo = patioLocalNameToInfo.plus(
                    local.name(),
                    localInfo
                );
            }

            return new Context(
                target,
                module,
                P.vector(),
                P.stack(),
                patioLocalNameToInfo,
                new NullStructuredInstructionBuilder(),
                P.stack()
            );
        }

        private final WasmTarget target;
        private final WasmModule module;
        private final PVector<WasmInstruction> instructions;
        private final PStack<PVector<WasmInstruction>> instructionsStack;
        private final PMap<String, PatioLocalInfo> patioLocalNameToInfo;
        private final WasmStructuredInstructionBuilder structuredInstructionBuilder;
        private final PStack<WasmStructuredInstructionBuilder> structuredInstructionBuilderStack;

        private Context(
            WasmTarget target,
            WasmModule module,
            PVector<WasmInstruction> instructions,
            PStack<PVector<WasmInstruction>> instructionsStack,
            PMap<String, PatioLocalInfo> patioLocalNameToInfo,
            WasmStructuredInstructionBuilder structuredInstructionBuilder,
            PStack<WasmStructuredInstructionBuilder> structuredInstructionBuilderStack
        ) {
            this.target = target;
            this.module = module;
            this.instructions = instructions;
            this.instructionsStack = instructionsStack;
            this.patioLocalNameToInfo = patioLocalNameToInfo;
            this.structuredInstructionBuilder = structuredInstructionBuilder;
            this.structuredInstructionBuilderStack = structuredInstructionBuilderStack;
        }

        private Context addType(WasmFuncType type) {
            return new Context(
                this.target,
                this.module.addType(type),
                this.instructions,
                this.instructionsStack,
                this.patioLocalNameToInfo,
                this.structuredInstructionBuilder,
                this.structuredInstructionBuilderStack
            );
        }

        private Context addInstruction(WasmInstruction instruction) {
            return new Context(
                this.target,
                this.module,
                this.instructions.plus(instruction),
                this.instructionsStack,
                this.patioLocalNameToInfo,
                this.structuredInstructionBuilder,
                this.structuredInstructionBuilderStack
            );
        }

        private WasmLocalIndex wasmLocalIndex(PatioLocalRef patioLocalRef) {
            return this.patioLocalNameToInfo.get(patioLocalRef.name()).wasmLocalIndex;
        }

        public WasmValueType wasmLocalType(PatioLocalRef patioLocalRef) {
            return this.patioLocalNameToInfo.get(patioLocalRef.name()).wasmType;
        }


        public Context enterBlockInstruction() {
            return new Context(
                this.target,
                this.module,
                P.vector(),
                this.instructionsStack.plus(this.instructions),
                this.patioLocalNameToInfo,
                new WasmBlockInstructionBuilder(),
                this.structuredInstructionBuilderStack.plus(this.structuredInstructionBuilder)
            );
        }

        public Context exitBlockInstruction() {
            // TODO: verify matching enter/exit?
            return exitStructuredInstruction();
        }

        public Context addBrToBlockInstruction() {
            return addBr(WasmBlockInstructionBuilder.class);
        }

        public Context enterLoopInstruction() {
            return new Context(
                this.target,
                this.module,
                P.vector(),
                this.instructionsStack.plus(this.instructions),
                this.patioLocalNameToInfo,
                new WasmLoopInstructionBuilder(),
                this.structuredInstructionBuilderStack.plus(this.structuredInstructionBuilder)
            );
        }

        public Context exitLoopInstruction() {
            // TODO: verify matching enter/exit?
            return exitStructuredInstruction();
        }

        public Context addBrToLoopInstruction() {
            return addBr(WasmLoopInstructionBuilder.class);
        }

        public Context enterIfInstruction(WasmBlockType blockType) {
            return new Context(
                this.target,
                this.module,
                P.vector(),
                this.instructionsStack.plus(this.instructions),
                this.patioLocalNameToInfo,
                new WasmIfInstructionBuilder(blockType),
                this.structuredInstructionBuilderStack.plus(this.structuredInstructionBuilder)
            );
        }

        public Context enterElse() {
            return new Context(
                this.target,
                this.module,
                P.vector(),
                this.instructionsStack,
                this.patioLocalNameToInfo,
                this.structuredInstructionBuilder.enterElse(this.instructions),
                this.structuredInstructionBuilderStack
            );
        }

        public Context exitIfInstruction() {
            // TODO: verify matching enter/exit?
            return exitStructuredInstruction();
        }

        public Context enterTryInstruction(WasmBlockType blockType) {
            return new Context(
                this.target,
                this.module,
                P.vector(),
                this.instructionsStack.plus(this.instructions),
                this.patioLocalNameToInfo,
                new WasmTryInstructionBuilder(blockType),
                this.structuredInstructionBuilderStack.plus(this.structuredInstructionBuilder)
            );
        }

        public Context enterCatch(WasmTagName tagName) {
            return new Context(
                this.target,
                this.module,
                P.vector(),
                this.instructionsStack,
                this.patioLocalNameToInfo,
                this.structuredInstructionBuilder.enterCatch(tagName, this.instructions),
                this.structuredInstructionBuilderStack
            );
        }

        public Context exitTryInstruction() {
            // TODO: verify matching enter/exit?
            return exitStructuredInstruction();
        }

        public Context exitStructuredInstruction() {
            var context = new Context(
                this.target,
                this.module,
                this.instructionsStack.get(0),
                this.instructionsStack.minus(0),
                this.patioLocalNameToInfo,
                this.structuredInstructionBuilderStack.get(0),
                this.structuredInstructionBuilderStack.minus(0)
            );

            var instruction = this.structuredInstructionBuilder
                .exit(this.instructions);

            return context.addInstruction(instruction);
        }

        private Context addBr(Class<? extends WasmStructuredInstructionBuilder> clazz) {
            var labelIndex = findLabelIndex(clazz);

            return this.addInstruction(new WasmBr(labelIndex));
        }

        private int findLabelIndex(Class<? extends WasmStructuredInstructionBuilder> clazz) {
            if (clazz.isInstance(this.structuredInstructionBuilder)) {
                return 0;
            }

            for (var i = 0; i < this.structuredInstructionBuilderStack.size(); i++) {
                if (clazz.isInstance(this.structuredInstructionBuilderStack.get(i))) {
                    return i + 1;
                }
            }

            throw new IllegalStateException("cannot find label index");
        }

        record Result<T>(T value, Context context) {}

        public <T> Result<T> toResult(T value) {
            return new Result<>(value, this);
        }
    }

    private interface WasmStructuredInstructionBuilder {
        WasmInstruction exit(PVector<WasmInstruction> instructions);

        default WasmStructuredInstructionBuilder enterElse(PVector<WasmInstruction> instructions) {
            throw new UnsupportedOperationException();
        }

        default WasmStructuredInstructionBuilder enterCatch(
            WasmTagName tagName,
            PVector<WasmInstruction> instructions
        ) {
            throw new UnsupportedOperationException();
        }
    }

    private static class NullStructuredInstructionBuilder implements WasmStructuredInstructionBuilder {
        @Override
        public WasmInstruction exit(PVector<WasmInstruction> instructions) {
            throw new UnsupportedOperationException();
        }
    }

    private static class WasmBlockInstructionBuilder implements WasmStructuredInstructionBuilder {
        @Override
        public WasmInstruction exit(PVector<WasmInstruction> instructions) {
            return new WasmBlock(WasmBlockType.EMPTY, instructions);
        }
    }

    private static class WasmLoopInstructionBuilder implements WasmStructuredInstructionBuilder {
        @Override
        public WasmInstruction exit(PVector<WasmInstruction> instructions) {
            return new WasmLoop(WasmBlockType.EMPTY, instructions);
        }
    }

    private record WasmIfInstructionBuilder(
        WasmBlockType blockType
    ) implements WasmStructuredInstructionBuilder {
        @Override
        public WasmInstruction exit(PVector<WasmInstruction> instructions) {
            throw new IllegalStateException("if instruction without else not supported");
        }

        @Override
        public WasmStructuredInstructionBuilder enterElse(PVector<WasmInstruction> instructions) {
            return new WasmIfElseInstructionBuilder(blockType, instructions);
        }
    }

    private record WasmIfElseInstructionBuilder(
        WasmBlockType blockType,
        PVector<WasmInstruction> ifTrue
    ) implements WasmStructuredInstructionBuilder {
        @Override
        public WasmInstruction exit(PVector<WasmInstruction> instructions) {
            return new WasmIf(blockType, ifTrue, instructions);
        }
    }

    private record WasmTryInstructionBuilder(
        WasmBlockType blockType
    ) implements WasmStructuredInstructionBuilder {
        @Override
        public WasmInstruction exit(PVector<WasmInstruction> instructions) {
            return new WasmTry(blockType, List.of(), instructions);
        }

        @Override
        public WasmStructuredInstructionBuilder enterCatch(
            WasmTagName tagName,
            PVector<WasmInstruction> instructions
        ) {
            return new WasmTryCatchInstructionBuilder(
                blockType,
                instructions,
                tagName,
                P.vector()
            );
        }
    }

    private record WasmTryCatchInstructionBuilder(
        WasmBlockType blockType,
        List<WasmInstruction> tryBodyInstructions,
        WasmTagName tagName,
        PVector<WasmCatchBlock> catchBlocks
    ) implements WasmStructuredInstructionBuilder {
        @Override
        public WasmInstruction exit(PVector<WasmInstruction> instructions) {
            return new WasmTry(
                this.blockType,
                this.catchBlocks.plus(new WasmCatchBlock(this.tagName, instructions)),
                this.tryBodyInstructions
            );
        }

        @Override
        public WasmStructuredInstructionBuilder enterCatch(
            WasmTagName tagName,
            PVector<WasmInstruction> instructions
        ) {
            return new WasmTryCatchInstructionBuilder(
                this.blockType,
                this.tryBodyInstructions,
                tagName,
                this.catchBlocks.plus(new WasmCatchBlock(this.tagName, instructions))
            );
        }
    }
}