path: root/src/main/java/org/shedlang/compiler/shidtopatio/ShidToPatioContext.java

package org.shedlang.compiler.shidtopatio;

import org.pcollections.OrderedPMap;
import org.pcollections.PMap;
import org.pcollections.PStack;
import org.pcollections.PVector;
import org.shedlang.compiler.backends.NameAllocator;
import org.shedlang.compiler.patio.*;
import org.shedlang.compiler.patio.instructions.PatioInstruction;
import org.shedlang.compiler.patio.instructions.PatioLocalStore;
import org.shedlang.compiler.patio.operands.PatioLocalRef;
import org.shedlang.compiler.patio.operands.PatioOperand;
import org.shedlang.compiler.patio.types.PatioDataValue;
import org.shedlang.compiler.patio.types.PatioType;
import org.shedlang.compiler.shidtopatio.errors.DuplicateCompilationUnitElementName;
import org.shedlang.compiler.sources.Source;
import org.shedlang.compiler.typechecker.TypeInfos;
import org.shedlang.compiler.types.*;
import org.shedlang.compiler.util.CheckReturnValue;
import org.shedlang.compiler.util.P;

import java.util.*;
import java.util.function.Function;

public record ShidToPatioContext(
    GlobalContext globalContext,
    CompilationUnitContext compilationUnitContext,
    PStack<FunctionContext> blockStack,
    FunctionContext functionContext
) {
    public record GlobalContext(
        TypeInfos typeInfos,
        PMap<SimpleStructType, Integer> structTags,
        PMap<UserDefinedEffect, Integer> effectTags,
        OrderedPMap<RuntimeType, Integer> runtimeTypeTags
    ) {
        public static final GlobalContext EMPTY = forTypeInfos(TypeInfos.EMPTY);

        @CheckReturnValue
        public static GlobalContext forTypeInfos(TypeInfos typeInfos) {
            var context = new GlobalContext(
                typeInfos,
                P.map(),
                P.map(),
                P.orderedPMap()
            );

            // Ensure that scalars have tag of 0: this is relied on in the
            // runtime (for instance, when allocating ints and strings).
            context = context.addRuntimeTypeTag(ScalarRuntimeType.INSTANCE);

            // Ensure that arrays have tag of 1.
            context = context.addRuntimeTypeTag(ArrayRuntimeType.INSTANCE);

            // Ensure that array builders have tag of 2.
            context = context.addRuntimeTypeTag(ArrayBuilderRuntimeType.INSTANCE);

            context = generateStructTags(context);
            context = generateEffectTags(context);

            return context;
        }

        @CheckReturnValue
        private static GlobalContext generateStructTags(GlobalContext context) {
            var structTags = P.<SimpleStructType, Integer>map();

            for (var structType : context.typeInfos.structTypes()) {
                // TODO: get struct types with infos to avoid potential throw.
                var structTypeInfo = context.typeInfos.structTypeInfo(structType).orElseThrow();
                var structRuntimeType = ShidToPatioStructs.structRuntimeType(structType, structTypeInfo);

                var runtimeTypeTagResult = context.runtimeTypeTag(structRuntimeType);
                var runtimeTypeTag = runtimeTypeTagResult.value();
                context = runtimeTypeTagResult.context();

                structTags = structTags.plus(structType, runtimeTypeTag);
            }

            return new GlobalContext(
                context.typeInfos,
                structTags,
                context.effectTags,
                context.runtimeTypeTags
            );
        }

        @CheckReturnValue
        private static GlobalContext generateEffectTags(GlobalContext context) {
            var effectTags = P.<UserDefinedEffect, Integer>map();

            for (var effect : context.typeInfos.userDefinedEffects()) {
                // TODO: get effects with infos to avoid potential throw.
                var effectInfo = context.typeInfos.userDefinedEffectInfo(effect).orElseThrow();
                var effectHandlerRuntimeType = ShidToPatioEffects.effectHandlerRuntimeType(effect, effectInfo);

                var runtimeTypeTagResult = context.runtimeTypeTag(effectHandlerRuntimeType);
                var runtimeTypeTag = runtimeTypeTagResult.value();
                context = runtimeTypeTagResult.context();

                effectTags = effectTags.plus(effect, runtimeTypeTag);
            }

            return new GlobalContext(
                context.typeInfos,
                context.structTags,
                effectTags,
                context.runtimeTypeTags
            );
        }

        @CheckReturnValue
        public NamespaceTypeInfo namespaceTypeInfo(NamespaceName namespaceName) {
            return typeInfos.namespaceTypeInfo(namespaceName)
                // TODO: throw better error
                .orElseThrow(() -> new RuntimeException("No type info for: " + namespaceName));
        }

        @CheckReturnValue
        public StructTypeInfo structTypeInfo(SimpleStructType type) {
            return typeInfos.structTypeInfo(type)
                // TODO: throw better error
                .orElseThrow(() -> new RuntimeException("No type info for: " + type));
        }

        @CheckReturnValue
        public int structTag(SimpleStructType type) {
            var tag = this.structTags.get(type);
            if (tag == null) {
                // TODO: throw better error
                throw new RuntimeException("No tag for: " + type);
            } else {
                return tag;
            }
        }

        @CheckReturnValue
        public UserDefinedEffectInfo userDefinedEffectInfo(UserDefinedEffect effect) {
            return typeInfos.userDefinedEffectInfo(effect)
                // TODO: throw better error
                .orElseThrow(() -> new RuntimeException("No effect info for: " + effect));
        }

        @CheckReturnValue
        public int effectTag(UserDefinedEffect effect) {
            var tag = this.effectTags.get(effect);
            if (tag == null) {
                // TODO: throw better error
                throw new RuntimeException("No tag for: " + effect);
            } else {
                return tag;
            }
        }

        @CheckReturnValue
        private GlobalContext addRuntimeTypeTag(RuntimeType runtimeType) {
            return new GlobalContext(
                this.typeInfos,
                this.structTags,
                this.effectTags,
                this.runtimeTypeTags.plus(runtimeType, this.runtimeTypeTags.size())
            );
        }

        @CheckReturnValue
        public Result<Integer> runtimeTypeTag(RuntimeType runtimeType) {
            var globalContext = this.runtimeTypeTags.containsKey(runtimeType)
                ? this
                : this.addRuntimeTypeTag(runtimeType);

            return globalContext.toResult(globalContext.runtimeTypeTags.get(runtimeType));
        }

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

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

    public record CompilationUnitContext(
        Optional<NamespaceName> namespaceName,
        OrderedPMap<String, PatioCompilationUnitElement> elementsByName
    ) {
        public static CompilationUnitContext EMPTY = initial(Optional.empty());

        public static CompilationUnitContext initial(Optional<NamespaceName> namespaceName) {
            return new CompilationUnitContext(
                namespaceName,
                OrderedPMap.empty()
            );
        }

        @CheckReturnValue
        public Optional<NamespaceName> namespaceName() {
            return this.namespaceName;
        }

        @CheckReturnValue
        public CompilationUnitContext addElement(PatioCompilationUnitElement element) {
            // TODO: handle existing element with name.
            return new CompilationUnitContext(
                this.namespaceName,
                this.elementsByName.plus(element.name(), element)
            );
        }

        @CheckReturnValue
        private <T> Result<T> toResult(T value) {
            return new Result<>(value, this);
        }

        @CheckReturnValue
        public Result<String> allocateCompilationUnitName(
            ShidToPatioCompilationUnitElementSpace space,
            Optional<String> name
        ) {
            var fullName = NameAllocator.allocateName(
                elementsByName::containsKey,
                "shed__" + space,
                name
            );

            return toResult(fullName);
        }

        @CheckReturnValue
        public CompilationUnitContext addCompilationUnitElement(
            PatioCompilationUnitElement element
        ) {
            if (this.elementsByName.containsKey(element.name())) {
                throw new DuplicateCompilationUnitElementName(element.name(), element.source());
            }

            return new CompilationUnitContext(
                this.namespaceName,
                this.elementsByName.plus(element.name(), element)
            );
        }

        @CheckReturnValue
        public CompilationUnitContext ensureCompilationUnitElement(
            PatioCompilationUnitElement element
        ) {
            var existingElement = this.elementsByName.get(element.name());

            if (existingElement == null) {
                return new CompilationUnitContext(
                    this.namespaceName,
                    this.elementsByName.plus(element.name(), element)
                );
            } else if (existingElement.equals(element)) {
                return this;
            } else {
                throw new DuplicateCompilationUnitElementName(element.name(), element.source());
            }

        }

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

    public record FunctionContext(
        PStack<PatioOperand> temporaryStack,
        PVector<PatioInstruction> instructions,
        OrderedPMap<String, PatioLocal> localsByName,
        PMap<String, PatioLocalRef> shidNameToLocal,
        PMap<String, PVector<Function<ShidToPatioContext, ShidToPatioContext>>> waitingForShidLocal,
        Optional<PatioOperand> stashedEffectHandler,
        Optional<PatioOperand> stashedGcRootStackIndex,
        PStack<PatioInstruction> onReturn
    ) {
        public static final FunctionContext EMPTY = new FunctionContext(
            P.stack(),
            P.vector(),
            OrderedPMap.empty(),
            P.map(),
            P.map(),
            Optional.empty(),
            Optional.empty(),
            P.stack()
        );

        @CheckReturnValue
        public Result<PatioOperand> temporaryPop() {
            return new FunctionContext(
                this.temporaryStack.minus(0),
                this.instructions,
                this.localsByName,
                this.shidNameToLocal,
                this.waitingForShidLocal,
                this.stashedEffectHandler,
                this.stashedGcRootStackIndex,
                this.onReturn
            ).toResult(this.temporaryStack.get(0));
        }

        @CheckReturnValue
        public Result<List<PatioOperand>> temporaryPopAll() {
            return new FunctionContext(
                P.stack(),
                this.instructions,
                this.localsByName,
                this.shidNameToLocal,
                this.waitingForShidLocal,
                this.stashedEffectHandler,
                this.stashedGcRootStackIndex,
                this.onReturn
            ).toResult(this.temporaryStack);
        }

        @CheckReturnValue
        public FunctionContext temporaryPush(PatioOperand value) {
            return new FunctionContext(
                this.temporaryStack.plus(value),
                this.instructions,
                this.localsByName,
                this.shidNameToLocal,
                this.waitingForShidLocal,
                this.stashedEffectHandler,
                this.stashedGcRootStackIndex,
                this.onReturn
            );
        }

        @CheckReturnValue
        public FunctionContext addInstruction(PatioInstruction instruction) {
            return new FunctionContext(
                this.temporaryStack,
                this.instructions.plus(instruction),
                this.localsByName,
                this.shidNameToLocal,
                this.waitingForShidLocal,
                this.stashedEffectHandler,
                this.stashedGcRootStackIndex,
                this.onReturn
            );
        }

        @CheckReturnValue
        public Result<PatioLocalRef> allocateLocal(
            ShidToPatioLocalSpace space,
            Optional<String> name,
            Function<String, PatioLocal> createLocal
        ) {
            var fullName = NameAllocator.allocateName(
                this.localsByName::containsKey,
                space,
                name
            );

            return new FunctionContext(
                this.temporaryStack,
                this.instructions,
                this.localsByName.plus(fullName, createLocal.apply(fullName)),
                this.shidNameToLocal,
                this.waitingForShidLocal,
                this.stashedEffectHandler,
                this.stashedGcRootStackIndex,
                this.onReturn
            ).toResult(new PatioLocalRef(fullName));
        }

        @CheckReturnValue
        public FunctionContext.Result<PatioLocalRef> declareShidLocal(
            String name,
            PatioType type,
            PatioLocalType localType
        ) {
            var context = this;

            if (this.shidNameToLocal.containsKey(name)) {
                // TODO: better error
                throw new RuntimeException("Shid local already exists: " + name);
            }

            var fullName = NameAllocator.allocateFixedName(
                this.localsByName::containsKey,
                ShidToPatioLocalSpace.LOCAL,
                name
            );

            var localRef = new PatioLocalRef(fullName);

            context = new FunctionContext(
                this.temporaryStack,
                this.instructions,
                this.localsByName.plus(
                    fullName,
                    new PatioLocal(fullName, type, localType)
                ),
                // TODO: the mapping is redundant given we fix the mapping
                this.shidNameToLocal.plus(name, localRef),
                this.waitingForShidLocal.plus(name, P.vector()),
                this.stashedEffectHandler,
                this.stashedGcRootStackIndex,
                this.onReturn
            );

            return context.toResult(localRef);
        }

        @CheckReturnValue
        public FunctionContext.Result<PVector<Function<ShidToPatioContext, ShidToPatioContext>>> initShidLocal(String name) {
            var context = new FunctionContext(
                this.temporaryStack,
                this.instructions,
                this.localsByName,
                this.shidNameToLocal,
                this.waitingForShidLocal.minus(name),
                this.stashedEffectHandler,
                this.stashedGcRootStackIndex,
                this.onReturn
            );

            return context.toResult(this.waitingForShidLocal.getOrDefault(name, P.vector()));
        }

        @CheckReturnValue
        public PatioLocalRef shidLocalNameToLocal(String name) {
            var local = this.shidNameToLocal.get(name);

            if (local == null) {
                // TODO: better error
                throw new RuntimeException("Unknown Shid local: " + name);
            }

            return local;
        }

        @CheckReturnValue
        public Optional<PatioLocalRef> shidLocalNameToLocalOrNone(String name) {
            var local = this.shidNameToLocal.get(name);

            return Optional.ofNullable(local);
        }

        @CheckReturnValue
        public boolean isShidLocalDefined(String name) {
            return !this.waitingForShidLocal.containsKey(name);
        }

        @CheckReturnValue
        public PatioType localType(PatioLocalRef local){
            // TODO: handle missing key.
            return this.localsByName.get(local.name()).type();
        }

        @CheckReturnValue
        public FunctionContext addWaitingForShidLocal(
            String name,
            Function<ShidToPatioContext, ShidToPatioContext> func
        ) {
            return new FunctionContext(
                this.temporaryStack,
                this.instructions,
                this.localsByName,
                this.shidNameToLocal,
                this.waitingForShidLocal.plus(
                    name,
                    this.waitingForShidLocal.getOrDefault(name, P.vector()).plus(func)
                ),
                this.stashedEffectHandler,
                this.stashedGcRootStackIndex,
                this.onReturn
            );
        }

        @CheckReturnValue
        public FunctionContext stashEffectHandler(PatioOperand effectHandler) {
            return new FunctionContext(
                this.temporaryStack,
                this.instructions,
                this.localsByName,
                this.shidNameToLocal,
                this.waitingForShidLocal,
                // TODO: check current value is empty
                Optional.of(effectHandler),
                this.stashedGcRootStackIndex,
                this.onReturn
            );
        }

        @CheckReturnValue
        public PatioOperand stashedEffectHandlerOrThrow() {
            // TODO: better error
            return this.stashedEffectHandler
                .orElseThrow(() -> new RuntimeException("No stashed effect handler"));
        }

        @CheckReturnValue
        public FunctionContext stashGcRootStackIndex(PatioOperand gcRootStackIndex) {
            return new FunctionContext(
                this.temporaryStack,
                this.instructions,
                this.localsByName,
                this.shidNameToLocal,
                this.waitingForShidLocal,
                this.stashedEffectHandler,
                // TODO: check current value is empty
                Optional.of(gcRootStackIndex),
                this.onReturn
            );
        }

        @CheckReturnValue
        public PatioOperand stashedGcRootStackIndexOrThrow() {
            // TODO: better error
            return this.stashedGcRootStackIndex
                .orElseThrow(() -> new RuntimeException("No stashed GC root stack index"));
        }

        public FunctionContext pushOnReturn(PatioInstruction onReturn) {
            return new FunctionContext(
                this.temporaryStack,
                this.instructions,
                this.localsByName,
                this.shidNameToLocal,
                this.waitingForShidLocal,
                this.stashedEffectHandler,
                this.stashedGcRootStackIndex,
                this.onReturn.plus(onReturn)
            );
        }

        public FunctionContext popOnReturn() {
            return new FunctionContext(
                this.temporaryStack,
                this.instructions,
                this.localsByName,
                this.shidNameToLocal,
                this.waitingForShidLocal,
                this.stashedEffectHandler,
                this.stashedGcRootStackIndex,
                this.onReturn.minus(0)
            );
        }

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

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

    public static final ShidToPatioContext EMPTY = new ShidToPatioContext(
        GlobalContext.EMPTY,
        CompilationUnitContext.EMPTY,
        P.stack(),
        FunctionContext.EMPTY
    );

    @CheckReturnValue
    public static ShidToPatioContext initialForCompilationUnit(
        Optional<NamespaceName> namespaceName,
        GlobalContext globalContext
    ) {
        return new ShidToPatioContext(
            globalContext,
            CompilationUnitContext.initial(namespaceName),
            P.stack(),
            FunctionContext.EMPTY
        );
    }

    @CheckReturnValue
    public NamespaceTypeInfo namespaceTypeInfo(NamespaceName namespaceName) {
        return this.globalContext.namespaceTypeInfo(namespaceName);
    }

    @CheckReturnValue
    public StructTypeInfo structTypeInfo(SimpleStructType type) {
        return this.globalContext.structTypeInfo(type);
    }

    @CheckReturnValue
    public int structTag(SimpleStructType type) {
        return globalContext.structTag(type);
    }

    @CheckReturnValue
    public UserDefinedEffectInfo userDefinedEffectInfo(UserDefinedEffect effect) {
        return this.globalContext.userDefinedEffectInfo(effect);
    }

    @CheckReturnValue
    public int effectTag(UserDefinedEffect effect) {
        return globalContext.effectTag(effect);
    }

    @CheckReturnValue
    public ShidToPatioResult<Integer> runtimeTypeTag(RuntimeType runtimeType) {
        return withGlobalContextResult(this.globalContext.runtimeTypeTag(runtimeType));
    }

    @CheckReturnValue
    public Set<RuntimeType> runtimeTypes() {
        return this.globalContext.runtimeTypeTags.keySet();
    }

    @CheckReturnValue
    public Optional<NamespaceName> namespaceName() {
        return this.compilationUnitContext.namespaceName();
    }

    @CheckReturnValue
    public ShidToPatioResult<String> allocateCompilationUnitElementName(
        ShidToPatioCompilationUnitElementSpace space,
        String name
    ) {
        return withCompilationUnitContextResult(
            this.compilationUnitContext.allocateCompilationUnitName(
                space,
                Optional.of(name)
            )
        );
    }

    @CheckReturnValue
    public ShidToPatioResult<String> allocateCompilationUnitElementName(
        ShidToPatioCompilationUnitElementSpace space
    ) {
        return withCompilationUnitContextResult(
            this.compilationUnitContext.allocateCompilationUnitName(
                space,
                Optional.empty()
            )
        );
    }

    @CheckReturnValue
    public ShidToPatioContext addDataElementMutable(
        String name,
        PatioVisibility visibility,
        PatioDataValue initialValue,
        Source source
    ) {
        var dataElement = new PatioDataElementMutableDefinition(
            name,
            visibility,
            initialValue,
            source
        );
        return addCompilationUnitElement(dataElement);
    }

    @CheckReturnValue
    public ShidToPatioContext addCompilationUnitElement(
        PatioCompilationUnitElement element
    ) {
        return withCompilationUnitContext(
            this.compilationUnitContext.addCompilationUnitElement(element)
        );
    }

    @CheckReturnValue
    public ShidToPatioContext ensureCompilationUnitElement(
        PatioCompilationUnitElement element
    ) {
        return withCompilationUnitContext(
            this.compilationUnitContext.ensureCompilationUnitElement(element)
        );
    }

    @CheckReturnValue
    public ShidToPatioResult<PatioOperand> temporaryPop() {
        var result = this.functionContext.temporaryPop();

        return withFunctionContextResult(result);
    }

    @CheckReturnValue
    public ShidToPatioContext temporaryPopIntoLocal(PatioLocalRef destination, Source source) {
        var result = this.functionContext.temporaryPop();

        return new ShidToPatioContext(this.globalContext, this.compilationUnitContext, this.blockStack, result.context)
            .addInstruction(new PatioLocalStore(destination, result.value, source));
    }

    @CheckReturnValue
    public ShidToPatioResult<List<PatioOperand>> temporaryPopAll() {
        var result = this.functionContext.temporaryPopAll();

        return withFunctionContextResult(result);
    }

    @CheckReturnValue
    public ShidToPatioContext temporaryPush(PatioOperand value) {
        var result = this.functionContext.temporaryPush(value);

        return withFunctionContext(result);
    }

    @CheckReturnValue
    public ShidToPatioContext addInstruction(PatioInstruction instruction) {
        return withFunctionContext(this.functionContext.addInstruction(instruction));
    }

    @CheckReturnValue
    public ShidToPatioResult<PatioLocalRef> allocateParam(
        ShidToPatioLocalSpace space,
        PatioType type
    ) {
        var result = this.functionContext.allocateLocal(
            space,
            Optional.empty(),
            (patioName) -> new PatioLocal(patioName, type, PatioLocalType.PARAM)
        );

        return withFunctionContextResult(result);
    }

    @CheckReturnValue
    public ShidToPatioResult<PatioLocalRef> allocateVal(
        ShidToPatioLocalSpace space,
        PatioType type
    ) {
        var result = this.functionContext.allocateLocal(
            space,
            Optional.empty(),
            (patioName) -> new PatioLocal(patioName, type, PatioLocalType.VAL)
        );

        return withFunctionContextResult(result);
    }

    @CheckReturnValue
    public ShidToPatioResult<PatioLocalRef> allocateVal(
        ShidToPatioLocalSpace space,
        String name,
        PatioType type
    ) {
        var result = this.functionContext.allocateLocal(
            space,
            Optional.of(name),
            (patioName) -> new PatioLocal(patioName, type, PatioLocalType.VAL)
        );

        return withFunctionContextResult(result);
    }

    @CheckReturnValue
    public ShidToPatioContext defineShidParam(String name) {
        return declareShidParam(name).initShidLocal(name);
    }

    @CheckReturnValue
    private ShidToPatioContext declareShidParam(String name) {
        return declareShidLocal(
            name,
            ShidToPatioTypes.VALUE_TYPE,
            PatioLocalType.PARAM
        ).context();
    }

    @CheckReturnValue
    public ShidToPatioResult<PatioLocalRef> declareShidVal(String name, PatioType type) {
        return declareShidLocal(
            name,
            type,
            PatioLocalType.VAL
        );
    }

    @CheckReturnValue
    public ShidToPatioResult<PatioLocalRef> declareShidLocal(
        String name,
        PatioType type,
        PatioLocalType patioLocalType
    ) {
        return withFunctionContextResult(this.functionContext.declareShidLocal(
            name,
            type,
            patioLocalType
        ));
    }

    @CheckReturnValue
    public ShidToPatioContext initShidLocal(String name) {
        var result = this.functionContext.initShidLocal(name);
        var context = withFunctionContext(result.context);

        for (var func : result.value) {
            context = func.apply(context);
        }

        return context;
    }

    @CheckReturnValue
    public ShidToPatioContext onceShidLocalDefined(
        String name,
        Function<ShidToPatioContext, ShidToPatioContext> func
    ) {
        var local = this.functionContext.shidLocalNameToLocalOrNone(name);

        if (local.isEmpty()) {
            // TODO: better error
            throw new RuntimeException("Unknown Shid local: " + name);
        } else if (!this.functionContext.isShidLocalDefined(name)) {
            return withFunctionContext(
                this.functionContext.addWaitingForShidLocal(name, func)
            );
        } else {
            return func.apply(this);
        }
    }

    @CheckReturnValue
    public PatioLocalRef shidLocalNameToLocal(String name) {
        return this.functionContext.shidLocalNameToLocal(name);
    }

    @CheckReturnValue
    public PatioType localType(PatioLocalRef local) {
        return this.functionContext.localType(local);
    }

    @CheckReturnValue
    public PatioType shidLocalType(String name) {
        var local = shidLocalNameToLocal(name);
        return localType(local);
    }

    @CheckReturnValue
    public ShidToPatioContext createFunction(
        String name,
        PatioVisibility visibility,
        List<PatioType> returnTypes,
        Source source,
        Function<ShidToPatioContext, ShidToPatioContext> createBody
    ) {
        var bodyContext = new ShidToPatioContext(
            this.globalContext,
            this.compilationUnitContext,
            P.stack(),
            FunctionContext.EMPTY
        );

        bodyContext = createBody.apply(bodyContext);

        var locals = bodyContext.functionContext.localsByName.values().stream().toList();

        return new ShidToPatioContext(
            bodyContext.globalContext,
            bodyContext.compilationUnitContext.addElement(new PatioFunctionDefinition(
                name,
                visibility,
                locals,
                returnTypes,
                bodyContext.functionContext.instructions,
                source
            )),
            this.blockStack,
            this.functionContext
        );
    }

    @CheckReturnValue
    public ShidToPatioContext enterBlock() {
        return new ShidToPatioContext(
            this.globalContext,
            this.compilationUnitContext,
            this.blockStack.plus(this.functionContext),
            new FunctionContext(
                P.stack(),
                P.vector(),
                this.functionContext.localsByName,
                this.functionContext.shidNameToLocal,
                this.functionContext.waitingForShidLocal,
                this.functionContext.stashedEffectHandler,
                this.functionContext.stashedGcRootStackIndex,
                this.functionContext.onReturn
            )
        );
    }

    @CheckReturnValue
    public ShidToPatioResult<List<PatioInstruction>> exitBlock() {
        if (!this.functionContext.temporaryStack.isEmpty()) {
            throw new RuntimeException("temporary stack must be empty at the end of a block");
        }

        return new ShidToPatioContext(
            this.globalContext,
            this.compilationUnitContext,
            this.blockStack.minus(0),
            new FunctionContext(
                this.blockStack.get(0).temporaryStack,
                this.blockStack.get(0).instructions,
                this.functionContext.localsByName,
                this.functionContext.shidNameToLocal,
                this.functionContext.waitingForShidLocal,
                this.blockStack.get(0).stashedEffectHandler,
                this.blockStack.get(0).stashedGcRootStackIndex,
                this.blockStack.get(0).onReturn
            )
        ).toResult(this.functionContext.instructions);
    }

    @CheckReturnValue
    public ShidToPatioContext stashEffectHandler(PatioOperand effectHandler) {
        return withFunctionContext(this.functionContext.stashEffectHandler(effectHandler));
    }

    @CheckReturnValue
    public PatioOperand stashedEffectHandler() {
        return this.functionContext.stashedEffectHandlerOrThrow();
    }

    @CheckReturnValue
    public String exceptionTypeName(UserDefinedEffect effect) {
        // TODO: Allocate properly
        return "shed__effect__exception__" + String.join("_", effect.namespaceName().parts()) + "__" + effect.name();
    }

    @CheckReturnValue
    public ShidToPatioContext stashGcRootStackIndex(PatioOperand gcRootStackIndex) {
        return withFunctionContext(this.functionContext.stashGcRootStackIndex(gcRootStackIndex));
    }

    @CheckReturnValue
    public PatioOperand stashedGcRootStackIndex() {
        return this.functionContext.stashedGcRootStackIndexOrThrow();
    }

    @CheckReturnValue
    public ShidToPatioContext pushOnReturn(PatioInstruction onReturn) {
        return withFunctionContext(this.functionContext.pushOnReturn(onReturn));
    }

    @CheckReturnValue
    public ShidToPatioContext popOnReturn() {
        return withFunctionContext(this.functionContext.popOnReturn());
    }

    @CheckReturnValue
    public List<PatioInstruction> onReturn() {
        return this.functionContext.onReturn();
    }

    @CheckReturnValue
    private ShidToPatioContext withGlobalContext(GlobalContext globalContext) {
        return new ShidToPatioContext(
            globalContext,
            this.compilationUnitContext,
            this.blockStack,
            this.functionContext
        );
    }


    @CheckReturnValue
    private <T> ShidToPatioResult<T> withGlobalContextResult(GlobalContext.Result<T> result) {
        return withGlobalContext(result.context).toResult(result.value);
    }

    @CheckReturnValue
    private ShidToPatioContext withCompilationUnitContext(CompilationUnitContext compilationUnitContext) {
        return new ShidToPatioContext(this.globalContext, compilationUnitContext, this.blockStack, this.functionContext);
    }

    @CheckReturnValue
    private <T> ShidToPatioResult<T> withCompilationUnitContextResult(CompilationUnitContext.Result<T> result) {
        return withCompilationUnitContext(result.context).toResult(result.value);
    }

    @CheckReturnValue
    private ShidToPatioContext withFunctionContext(FunctionContext functionContext) {
        return new ShidToPatioContext(this.globalContext, this.compilationUnitContext, this.blockStack, functionContext);
    }

    @CheckReturnValue
    private <T> ShidToPatioResult<T> withFunctionContextResult(FunctionContext.Result<T> result) {
        return withFunctionContext(result.context).toResult(result.value);
    }

    @CheckReturnValue
    public <T> ShidToPatioResult<T> toResult(T value) {
        return new ShidToPatioResult<>(value, this);
    }

    @CheckReturnValue
    public GlobalContext.Result<PatioCompilationUnit> toCompilationUnit(Source source) {
        return this.globalContext.toResult(new PatioCompilationUnit(
            new ArrayList<>(this.compilationUnitContext.elementsByName.values()),
            source
        ));
    }
}