JspecifyChecker.java

/*
 * SPDX-FileCopyrightText: 2025 kaumei.io
 * SPDX-License-Identifier: Apache-2.0
 */
package io.kaumei.jdbc.anno.ctx;

import io.kaumei.jdbc.anno.model.OptionalFlag;

import javax.lang.model.element.Element;
import javax.lang.model.type.TypeMirror;
import java.util.function.BiFunction;
import java.util.function.Predicate;

class JspecifyChecker implements BiFunction<Element, TypeMirror, OptionalFlag> {

    private final Predicate<TypeMirror> nonNull;
    private final Predicate<TypeMirror> nullMarked;
    private final Predicate<TypeMirror> nullUnmarked;
    private final Predicate<TypeMirror> nullable;

    JspecifyChecker(Context ctx) {
        this.nonNull = predicate(ctx, "org.jspecify.annotations.NonNull");
        this.nullMarked = predicate(ctx, "org.jspecify.annotations.NullMarked");
        this.nullUnmarked = predicate(ctx, "org.jspecify.annotations.NullUnmarked");
        this.nullable = predicate(ctx, "org.jspecify.annotations.Nullable");
    }

    private static Predicate<TypeMirror> predicate(Context ctx, String fcq) {
        var element = ctx.elements.getTypeElement(fcq);
        if (element != null) {
            return (given) -> ctx.isSameType(element.asType(), given);
        }
        return (given) -> false;
    }


    /**
     * Primitives → NON_NULL
     * - Optional<T> → OPTIONAL_TYPE
     * - @Nullable on element → OPTIONAL
     * - @NonNull on element → NON_NULL
     * - In @NullMarked scope → unannotated types are NON_NULL
     * - In @NullUnmarked scope → unannotated types are UNSPECIFIED
     * - No markers found → UNSPECIFIED
     */
    @Override
    public OptionalFlag apply(Element context, TypeMirror type) {
        boolean nullable = false;
        boolean nonnull = false;

        // Check for @Nullable or @NonNull annotation on the element (e.g., method, parameter, field)
        for (var annotation : type.getAnnotationMirrors()) {
            var annoType = annotation.getAnnotationType();
            if (!nullable && this.nullable.test(annoType)) {
                nullable = true;
            } else if (!nonnull && this.nonNull.test(annoType)) {
                nonnull = true;
            }
        }

        Element current = context;
        while (current != null && !nullable && !nonnull) {
            // Check for @Nullable or @NonNull annotation on the enclosing elements (bottom up)
            // stop if we found nullable or nonnull
            for (var annotation : current.getAnnotationMirrors()) {
                var annoType = annotation.getAnnotationType();
                if (this.nullMarked.test(annoType)) {
                    nonnull = true;
                } else if (this.nullUnmarked.test(annoType)) {
                    return OptionalFlag.UNSPECIFIED;
                }
            }
            current = current.getEnclosingElement();
        }

        if (nonnull && !nullable) {
            return OptionalFlag.NON_NULL;
        } else if (!nonnull && nullable) {
            return OptionalFlag.NULLABLE;
        }
        return OptionalFlag.UNSPECIFIED;
    }
}