package p.studio.compiler.pbs;

import p.studio.compiler.pbs.ast.PbsAst;

public final class PbsCallableSlotCounter {
    private PbsCallableSlotCounter() {
    }

    public static int returnSlots(
            final PbsAst.ReturnKind returnKind,
            final PbsAst.TypeRef returnType) {
        return switch (returnKind) {
            case INFERRED_UNIT, EXPLICIT_UNIT -> 0;
            case RESULT -> 1;
            case PLAIN -> Math.max(0, typeSlots(returnType));
        };
    }

    private static int typeSlots(final PbsAst.TypeRef typeRef) {
        final var unwrapped = unwrapGroup(typeRef);
        if (unwrapped == null) {
            return 0;
        }
        return switch (unwrapped.kind()) {
            case UNIT, ERROR -> 0;
            case NAMED_TUPLE -> unwrapped.fields().stream()
                    .mapToInt(field -> Math.max(1, typeSlots(field.typeRef())))
                    .sum();
            default -> 1;
        };
    }

    private static PbsAst.TypeRef unwrapGroup(final PbsAst.TypeRef typeRef) {
        if (typeRef == null) {
            return null;
        }
        if (typeRef.kind() != PbsAst.TypeRefKind.GROUP) {
            return typeRef;
        }
        return unwrapGroup(typeRef.inner());
    }
}
