public enum InstrumentationType {

    /**
     * Before:
     * <pre>
     *     public void test(&#064;NotNull String s) {
     *         // body
     *     }
     * </pre>
     * After:
     * <pre>
     *     public void test(&#064;NotNull String s) {
     *         if (s == null) {
     *             throw new NullPointerException("[problem details]");
     *         }
     *         // body
     *     }
     * </pre>
     */
    METHOD_PARAMETER("parameter"),

    /**
     * Before:
     * <pre>
     *     &#064;NotNull
     *     public String compute() {
     *         return doCompute();
     *     }
     * </pre>
     * After:
     * <pre>
     *     &#064;NotNull
     *     public String compute() {
     *         String tmpVar = doCompute();
     *         if (tmpVar == null) {
     *             throw new NullPointerException("[problem details]");
     *         }
     *         return tmpVar;
     *     }
     * </pre>
     */
    METHOD_RETURN("return");

    private static Map<String, InstrumentationType> BY_SHORT_NAME = new HashMap<>();
    static {
        for (InstrumentationType type : values()) {
            BY_SHORT_NAME.put(type.getShortName(), type);
        }
    }

    @NotNull private final String shortName;

    InstrumentationType(@NotNull String shortName) {
        this.shortName = shortName;
    }

    @Nullable
    public static InstrumentationType byShortName(@NotNull String shortName) {
        return BY_SHORT_NAME.get(shortName);
    }

    @NotNull
    public String getShortName() {
        return shortName;
    }
} public class TrautePluginSettings {

    public static final String DEFAULT_EXCEPTION_TO_THROW = NullPointerException.class.getSimpleName();

    private final Set<String>                           notNullAnnotations          = new HashSet<>();
    private final Set<String>                           nullableAnnotations         = new HashSet<>();
    private final Set<InstrumentationType>              instrumentationsToApply     = new HashSet<>();
    private final Map<InstrumentationType, String>      exceptionsToThrow           = new HashMap<>();
    private final Map<InstrumentationType, String>      exceptionTextPatterns       = new HashMap<>();
    private final Map<InstrumentationType, Set<String>> notNullByDefaultAnnotations = new HashMap<>();

    @Nullable private final File logFile;

    private final boolean verboseMode;

    public TrautePluginSettings(@NotNull Set<String> notNullAnnotations,
                                @NotNull Set<String> nullableAnnotations,
                                @NotNull Set<InstrumentationType> instrumentationsToApply,
                                @NotNull Map<InstrumentationType, String> exceptionsToThrow,
                                @NotNull Map<InstrumentationType, String> exceptionTextPatterns,
                                @NotNull Map<InstrumentationType, Set<String>> notNullByDefaultAnnotations,
                                @Nullable File logFile,
                                boolean verboseMode)
    {
        this.logFile = logFile;
        this.notNullAnnotations.addAll(notNullAnnotations);
        this.nullableAnnotations.addAll(nullableAnnotations);
        this.instrumentationsToApply.addAll(instrumentationsToApply);
        this.exceptionsToThrow.putAll(exceptionsToThrow);
        this.exceptionTextPatterns.putAll(exceptionTextPatterns);
        this.notNullByDefaultAnnotations.putAll(notNullByDefaultAnnotations);
        this.verboseMode = verboseMode;
    }

    @NotNull
    public Set<String> getNotNullAnnotations() {
        return notNullAnnotations;
    }

    @NotNull
    public Set<String> getNullableAnnotations() {
        return nullableAnnotations;
    }

    public boolean isEnabled(@NotNull InstrumentationType type) {
        return instrumentationsToApply.contains(type);
    }

    @NotNull
    public Set<InstrumentationType> getInstrumentationsToApply() {
        return instrumentationsToApply;
    }

    @NotNull
    public String getExceptionToThrow(@NotNull InstrumentationType type) {
        return exceptionsToThrow.getOrDefault(type, DEFAULT_EXCEPTION_TO_THROW);
    }

    @NotNull
    public Map<InstrumentationType, String> getExceptionsToThrow() {
        return exceptionsToThrow;
    }

    @Nullable
    public String getExceptionTextPattern(@NotNull InstrumentationType type) {
        return exceptionTextPatterns.get(type);
    }

    @NotNull
    public Map<InstrumentationType, String> getExceptionTextPatterns() {
        return exceptionTextPatterns;
    }

    @NotNull
    public Set<String> getNotNullByDefaultAnnotations(@NotNull InstrumentationType type) {
        return notNullByDefaultAnnotations.get(type);
    }

    @NotNull
    public Map<InstrumentationType, Set<String>> getNotNullByDefaultAnnotations() {
        return notNullByDefaultAnnotations;
    }

    @NotNull
    public Optional<File> getLogFile() {
        return Optional.ofNullable(logFile);
    }

    public boolean isVerboseMode() {
        return verboseMode;
    }
} @SuppressWarnings("UnusedReturnValue")
public class TrautePluginSettingsBuilder {

    public static final Set<String> DEFAULT_NOT_NULL_ANNOTATIONS = new HashSet<>(asList(
            // Used by IntelliJ by default - https://www.jetbrains.com/help/idea/nullable-and-notnull-annotations.html
            "org.jetbrains.annotations.NotNull",

            // JSR-305 - status=dormant - https://jcp.org/en/jsr/detail?id=305
            "javax.annotation.Nonnull",

            // JavaEE - https://docs.oracle.com/javaee/7/api/javax/validation/constraints/NotNull.html
            "javax.validation.constraints.NotNull",

            // FindBugs - http://findbugs.sourceforge.net/api/edu/umd/cs/findbugs/annotations/NonNull.html
            "edu.umd.cs.findbugs.annotations.NonNull",

            // Android - https://developer.android.com/reference/android/support/annotation/NonNull.html
            "android.support.annotation.NonNull",

            // Eclipse - http://help.eclipse.org/oxygen/index.jsp?topic=%2Forg.eclipse.jdt.doc.user%2Ftasks%2Ftask-using_null_annotations.htm
            "org.eclipse.jdt.annotation.NonNull",

            // Lombok - https://projectlombok.org/api/lombok/NonNull.html
            "lombok.NonNull",

            // Spring - https://docs.spring.io/spring-framework/docs/current/javadoc-api/org/springframework/lang/NonNull.html
            "org.springframework.lang.NonNull",

            // Checker - https://checkerframework.org/api/org/checkerframework/checker/nullness/qual/NonNull.html
            "org.checkerframework.checker.nullness.qual.NonNull",

            // Checker - https://checkerframework.org/api/org/checkerframework/checker/nullness/compatqual/NonNullDecl.html
            "org.checkerframework.checker.nullness.compatqual.NonNullDecl",

            // Checker - https://checkerframework.org/api/org/checkerframework/checker/nullness/compatqual/NonNullType.html
            "org.checkerframework.checker.nullness.compatqual.NonNullType"
    ));

    /**
     * It's possible to specify that method parameters should be treated as {@code not-null} by default.
     * Current collection holds default annotations to use.
     */
    public static final Set<String> DEFAULT_PARAMETERS_NOT_NULL_BY_DEFAULT_ANNOTATIONS = new HashSet<>(asList(
            // Eclipse - https://help.eclipse.org/mars/index.jsp?topic=%2Forg.eclipse.jdt.doc.isv%2Freference%2Fapi%2Forg%2Feclipse%2Fjdt%2Fannotation%2FNonNullByDefault.html
            "org.eclipse.jdt.annotation.NonNullByDefault",

            // JSR-305 - https://static.javadoc.io/com.google.code.findbugs/jsr305/3.0.1/javax/annotation/ParametersAreNonnullByDefault.html
            "javax.annotation.ParametersAreNonnullByDefault",

            // Spring - https://docs.spring.io/spring-framework/docs/current/javadoc-api/org/springframework/lang/NonNullApi.html
            "org.springframework.lang.NonNullApi"
    ));

    public static final Set<String> DEFAULT_RETURN_NOT_NULL_BY_DEFAULT_ANNOTATIONS = new HashSet<>(
            // Spring - https://docs.spring.io/spring-framework/docs/current/javadoc-api/org/springframework/lang/NonNullApi.html
            singleton("org.springframework.lang.NonNullApi")
    );

    /**
     * This annotations work in conjunction with {@link #DEFAULT_PARAMETERS_NOT_NULL_BY_DEFAULT_ANNOTATIONS}
     * annotations, e.g. particular package/class/method might be marked by, say,
     * {@code javax.annotation.ParametersAreNonnullByDefault} and we want to specify that particular method parameter
     * might be {@code null}.
     */
    public static final Set<String> DEFAULT_NULLABLE_ANNOTATIONS = new HashSet<>(asList(
            // IntelliJ - https://www.jetbrains.com/help/idea/nullable-and-notnull-annotations.html#nullable
            "org.jetbrains.annotations.Nullable",

            // JSR-305 - status=dormant - https://jcp.org/en/jsr/detail?id=305
            "javax.annotation.Nullable",
            "javax.annotation.CheckForNull",

            // JavaEE - https://docs.oracle.com/javaee/7/api/javax/validation/constraints/Null.html
            "javax.validation.constraints.Null",

            // FindBugs - http://findbugs.sourceforge.net/api/edu/umd/cs/findbugs/annotations/Nullable.html
            "edu.umd.cs.findbugs.annotations.Nullable",

            // FindBugs - http://findbugs.sourceforge.net/api/edu/umd/cs/findbugs/annotations/CheckForNull.html
            "edu.umd.cs.findbugs.annotations.CheckForNull",

            // Android - https://developer.android.com/reference/android/support/annotation/Nullable.html
            "android.support.annotation.Nullable",

            // Eclipse - http://help.eclipse.org/oxygen/index.jsp?topic=%2Forg.eclipse.jdt.doc.user%2Ftasks%2Ftask-using_null_annotations.htm
            "org.eclipse.jdt.annotation.Nullable",

            // Spring - https://docs.spring.io/spring-framework/docs/current/javadoc-api/org/springframework/lang/Nullable.html
            "org.springframework.lang.Nullable",

            // Checker - https://checkerframework.org/api/org/checkerframework/checker/nullness/qual/Nullable.html
            "org.checkerframework.checker.nullness.qual.Nullable",

            // Checker - https://checkerframework.org/api/org/checkerframework/checker/nullness/compatqual/NullableDecl.html
            "org.checkerframework.checker.nullness.compatqual.NullableDecl",

            // Checker - https://checkerframework.org/api/org/checkerframework/checker/nullness/compatqual/NullableType.html
            "org.checkerframework.checker.nullness.compatqual.NullableType"
    ));

    public static final Set<InstrumentationType> DEFAULT_INSTRUMENTATIONS_TO_APPLY =
            EnumSet.allOf(InstrumentationType.class);

    public static final boolean DEFAULT_VERBOSE_MODE = false;

    private final Set<String>              notNullAnnotations      = new HashSet<>();
    private final Set<String>              nullableAnnotations     = new HashSet<>();
    private final Set<InstrumentationType> instrumentationsToApply = EnumSet.noneOf(InstrumentationType.class);

    private final Map<InstrumentationType, String>      exceptionsToThrow           = new HashMap<>();
    private final Map<InstrumentationType, String>      exceptionTextPatterns       = new HashMap<>();
    private final Map<InstrumentationType, Set<String>> notNullByDefaultAnnotations = new HashMap<>();

    @Nullable private File    logFile;
    @Nullable private Boolean verbose;

    @NotNull
    public static TrautePluginSettingsBuilder settingsBuilder() {
        return new TrautePluginSettingsBuilder();
    }

    @NotNull
    public TrautePluginSettingsBuilder withNotNullAnnotations(@NotNull String... notNullAnnotations) {
        this.notNullAnnotations.addAll(Arrays.asList(notNullAnnotations));
        return this;
    }

    @NotNull
    public TrautePluginSettingsBuilder withNullableAnnotations(@NotNull String... nullableAnnotations) {
        this.nullableAnnotations.addAll(Arrays.asList(nullableAnnotations));
        return this;
    }

    @NotNull
    public TrautePluginSettingsBuilder withInstrumentationToApply(@NotNull InstrumentationType type) {
        instrumentationsToApply.add(type);
        return this;
    }

    @NotNull
    public TrautePluginSettingsBuilder withExceptionToThrow(@NotNull InstrumentationType type,
                                                            @NotNull String exceptionToThrow)
    {
        exceptionsToThrow.put(type, exceptionToThrow);
        return this;
    }

    @NotNull
    public TrautePluginSettingsBuilder withExceptionTextPattern(@NotNull InstrumentationType type,
                                                                @NotNull String pattern)
    {
        exceptionTextPatterns.put(type, pattern);
        return this;
    }

    @NotNull
    public TrautePluginSettingsBuilder withNotNullByDefaultAnnotations(@NotNull InstrumentationType type,
                                                                       @NotNull Collection<String> annotations)
    {
        notNullByDefaultAnnotations.put(type, new HashSet<>(annotations));
        return this;
    }

    @NotNull
    public TrautePluginSettingsBuilder withLogFile(@NotNull File file) {
        logFile = file;
        return this;
    }

    @NotNull
    public TrautePluginSettingsBuilder withVerboseMode(boolean verbose) {
        this.verbose = verbose;
        return this;
    }

    @NotNull
    public TrautePluginSettings build() {
        Set<String> notNullAnnotations = new HashSet<>(this.notNullAnnotations);
        if (notNullAnnotations.isEmpty()) {
            notNullAnnotations.addAll(DEFAULT_NOT_NULL_ANNOTATIONS);
        }

        Set<String> nullableAnnotations = new HashSet<>(this.nullableAnnotations);
        if (nullableAnnotations.isEmpty()) {
            nullableAnnotations.addAll(DEFAULT_NULLABLE_ANNOTATIONS);
        }

        Map<InstrumentationType, Set<String>> notNullByDefaultAnnotations
                = new HashMap<>(this.notNullByDefaultAnnotations);
        if (!notNullByDefaultAnnotations.containsKey(InstrumentationType.METHOD_PARAMETER)) {
            notNullByDefaultAnnotations.put(InstrumentationType.METHOD_PARAMETER,
                                            DEFAULT_PARAMETERS_NOT_NULL_BY_DEFAULT_ANNOTATIONS);
        }
        if (!notNullByDefaultAnnotations.containsKey(InstrumentationType.METHOD_RETURN)) {
            notNullByDefaultAnnotations.put(InstrumentationType.METHOD_RETURN,
                                            DEFAULT_RETURN_NOT_NULL_BY_DEFAULT_ANNOTATIONS);
        }

        Set<InstrumentationType> instrumentationsToApply = EnumSet.copyOf(this.instrumentationsToApply);
        if (instrumentationsToApply.isEmpty()) {
            instrumentationsToApply.addAll(DEFAULT_INSTRUMENTATIONS_TO_APPLY);
        }

        Boolean verbose = this.verbose;
        if (verbose == null) {
            verbose = DEFAULT_VERBOSE_MODE;
        }
        return new TrautePluginSettings(notNullAnnotations,
                                        nullableAnnotations,
                                        instrumentationsToApply,
                                        exceptionsToThrow,
                                        exceptionTextPatterns,
                                        notNullByDefaultAnnotations,
                                        logFile,
                                        verbose);
    }
} public class StatsCollector {

    private final ConcurrentMap<InstrumentationType, Long> stats = new ConcurrentHashMap<>();

    public void increment(@NotNull InstrumentationType type) {
        add(type, 1);
    }

    public void add(@NotNull InstrumentationType type, long count) {
        stats.compute(type, (key, value) -> value == null ? count : value + count);
    }

    @NotNull
    public ConcurrentMap<InstrumentationType, Long> getStats() {
        return stats;
    }

    @Override
    public String toString() {
        return stats.toString();
    }
} public class TrauteConstants {

    public static final String PLUGIN_NAME = "Traute";

    public static final Set<String> PRIMITIVE_TYPES = Collections.unmodifiableSet(new HashSet<>(asList(
            "boolean", "byte", "short", "char", "int", "long", "float", "double"
    )));

    public static final Set<String> METHOD_RETURN_TYPES_TO_SKIP;

    /**
     * <p>
     *     Compiler's option name to use for specifying custom {@code @NotNull} annotations to use
     *     ({@value #SEPARATOR}-separated).
     * </p>
     * <p>
     *     This is not mandatory setting as default annotations are used otherwise. Only given annotations are
     *     checked if this argument is specified.
     * </p>
     * <p>
     *     Example: consider a situation when given parameter's value is
     *     {@code 'org.eclipse.jdt.annotation.NonNull:android.support.annotation.NonNull'} (eclipse and android
     *     annotations, defined as
     *     {@code -Atraute.annotations.not.null=org.eclipse.jdt.annotation.NonNull:android.support.annotation.NonNull}
     *     in the {@code javac} command line). That means that a method which parameter is marked by, say
     *     {@code org.jetbrains.annotations.NotNull} won't trigger {@code null}-check generation by the plugin.
     * </p>
     */
    public static final String OPTION_ANNOTATIONS_NOT_NULL = "traute.annotations.not.null";

    /**
     * <p>
     *     Compiler's option name to use for specifying custom {@code @Nullable} annotations to use
     *     ({@value #SEPARATOR}-separated).
     * </p>
     * <p>
     *     This is not mandatory setting as default annotations are used otherwise. Only given annotations are
     *     checked if this argument is specified.
     * </p>
     * <p>
     *     Example: consider a situation when given parameter's value is
     *     {@code 'org.eclipse.jdt.annotation.Nullable:android.support.annotation.Nullable'} (eclipse and android
     *     annotations, defined as
     *     {@code -Atraute.annotations.nullable=org.eclipse.jdt.annotation.Nullable:android.support.annotation.Nullable}
     *     in the {@code javac} command line). That means that a method which parameter is marked by, say
     *     {@code org.jetbrains.annotations.Nullable} won't be considered as {@code nullable} by the plugin.
     * </p>
     */
    public static final String OPTION_ANNOTATIONS_NULLABLE = "traute.annotations.nullable";

    /**
     * <p>
     *     Prefix for compiler's option prefix for specifying a set of {@code not null by default} annotations
     *     (separated by {@value #SEPARATOR}). Resulting option is constructed from the current prefix
     *     and {@link InstrumentationType#getShortName()}.
     * </p>
     * <p>
     *     E.g. {@code -Atraute.not.null.by.default.parameter=javax.annotation.ParametersAreNonnullByDefault}
     *     instructs the plugin to consider all method parameters not marked as {@link #OPTION_ANNOTATIONS_NULLABLE}
     *     to be treated as {@code not-null} for a package where {@code package-info.java} is marked by the
     *     {@code ParametersAreNonnullByDefault} annotation.
     * </p>
     */
    public static final String OPTION_PREFIX_ANNOTATIONS_NOT_NULL_BY_DEFAULT = "traute.annotations.not.null.by.default.";

    /**
     * <p>
     *     Compiler's option name to use for specifying if the plugin should produce verbose output
     *     for its processing
     * </p>
     * <p>
     *     Verbose output is produced if this option's value is {@code 'true'}, e.g.
     *     {@code '-Atraute.log.verbose=true'} in the {@code javac} command line. Any other value (or now value)
     *     mean that no verbose output should be provided.
     * </p>
     */
    public static final String OPTION_LOG_VERBOSE = "traute.log.verbose";

    /**
     * <p>Compiler's option name for specifying a path to a file to store plugin's logs.</p>
     * <p>The logs are printed to the standard compiler's output by default.</p>
     */
    public static final String OPTION_LOG_FILE = "traute.log.file";

    /**
     * Compiler's option name to use for specifying instrumentation types to use
     *
     * @see InstrumentationType
     * @see InstrumentationType#getShortName()
     */
    public static final String OPTION_INSTRUMENTATIONS_TO_USE = "traute.instrumentations";

    /**
     * <p>
     *     Prefix for compiler's option prefix for specifying an exception to throw on failed
     *     {@code null}-check. Resulting option is constructed from the current prefix and
     *     {@link InstrumentationType#getShortName()}.
     * </p>
     * <p>
     *     E.g. {@code -Atraute.exception.parameter=IllegalArgumentException} instructs the plugin
     *     to generated a parameter check where an {@link IllegalArgumentException} is thrown in case
     *     of failure.
     * </p>
     */
    public static final String OPTION_PREFIX_EXCEPTION_TO_THROW = "traute.exception.";

    /**
     * <p>
     *     Prefix for compiler's option prefix for specifying a text to use in an exception thrown from
     *     a failed {@code null}-check. Resulting option is constructed from the current prefix and
     *     {@link InstrumentationType#getShortName()}.
     * </p>
     * <p>
     *     E.g. {@code "-Atraute.failure.text.parameter=${PARAMETER_NAME} must not be null"} instructs the plugin
     *     to generated a check for a parameter named, say, {@code 'arg'} with exception text
     *     {@code "arg must not be null"}.
     * </p>
     */
    public static final String OPTION_PREFIX_EXCEPTION_TEXT = "traute.failure.text.";

    /**
     * This text is replaced by the actual parameter name in the
     * {@link InstrumentationType#METHOD_PARAMETER parametere check}.
     */
    public static final String VARIABLE_PARAMETER_NAME = "PARAMETER_NAME";

    /**
     * This function capitalizes the variable inside it, e.g. {@code ${capitalize(PARAMETER_NAME)}} for
     * parameter {@code 'arg'} produced {@code 'Arg'}.
     */
    public static final String FUNCTION_CAPITALIZE = "capitalize";

    /**
     * Separator to use for composite properties, e.g. when a user want to specify more than one
     * {@code NonNull} annotation.
     */
    public static final String SEPARATOR = ":";

    public static final String PACKAGE_INFO = "package-info";

    static {
        Set<String> set = new HashSet<>(PRIMITIVE_TYPES);
        set.add("void");
        set.add("Void");
        METHOD_RETURN_TYPES_TO_SKIP = set;
    }

    private TrauteConstants() {
    }
} public class CompilationUnitProcessingContext {

    private final Set<String> imports = new HashSet<>();

    @NotNull private final TrautePluginSettings          pluginSettings;
    @NotNull private final TreeMaker                     astFactory;
    @NotNull private final Names                         symbolsTable;
    @NotNull private final TrautePluginLogger            logger;
    @NotNull private final StatsCollector                statsCollector;
    @NotNull private final ExceptionTextGeneratorManager exceptionTextGeneratorManager;
    @NotNull private final PackageInfoManager            packageInfoManager;

    public CompilationUnitProcessingContext(
            @NotNull TrautePluginSettings pluginSettings,
            @NotNull TreeMaker astFactory,
            @NotNull Names symbolsTable,
            @NotNull TrautePluginLogger logger,
            @NotNull StatsCollector statsCollector,
            @NotNull ExceptionTextGeneratorManager exceptionTextGeneratorManager,
            @NotNull PackageInfoManager packageInfoManager)
    {
        this.pluginSettings = pluginSettings;
        this.statsCollector = statsCollector;
        this.astFactory = astFactory;
        this.symbolsTable = symbolsTable;
        this.logger = logger;
        this.exceptionTextGeneratorManager = exceptionTextGeneratorManager;
        this.packageInfoManager = packageInfoManager;
    }

    public void addImport(@NotNull String importText) {
        imports.add(importText);
    }

    @NotNull
    public TrautePluginSettings getPluginSettings() {
        return pluginSettings;
    }

    @NotNull
    public Set<String> getImports() {
        return imports;
    }

    @NotNull
    public TreeMaker getAstFactory() {
        return astFactory;
    }

    @NotNull
    public Names getSymbolsTable() {
        return symbolsTable;
    }

    @NotNull
    public TrautePluginLogger getLogger() {
        return logger;
    }

    @NotNull
    public StatsCollector getStatsCollector() {
        return statsCollector;
    }

    @NotNull
    public ExceptionTextGeneratorManager getExceptionTextGeneratorManager() {
        return exceptionTextGeneratorManager;
    }

    @NotNull
    public PackageInfoManager getPackageInfoManager() {
        return packageInfoManager;
    }
} public class InstrumentationApplianceFinder extends TreeScanner<Void, Void> {

    private final Stack<Tree>    parents                    = new Stack<>();
    private final Stack<String>  classNames                 = new Stack<>();
    private final Stack<Boolean> processingInterface        = new Stack<>();
    private final Stack<String>  parametersNotNullByDefault = new Stack<>();
    private final Stack<String>  returnNotNullByDefault     = new Stack<>();

    @NotNull private final CompilationUnitProcessingContext          context;
    @NotNull private final Instrumentator<ParameterToInstrumentInfo> parameterInstrumenter;
    @NotNull private final Instrumentator<ReturnToInstrumentInfo>    returnInstrumenter;

    private String              packageName;
    private String              methodName;
    private JCTree.JCExpression methodReturnType;
    private String              methodNotNullAnnotation;
    private int                 tmpVariableCounter;
    private int                 anonymousClassCounter;
    private boolean             instrumentReturnExpression;

    public InstrumentationApplianceFinder(@NotNull CompilationUnitProcessingContext context,
                                          @NotNull Instrumentator<ParameterToInstrumentInfo> parameterInstrumentator,
                                          @NotNull Instrumentator<ReturnToInstrumentInfo> returnInstrumentator)
    {
        this.context = context;
        this.parameterInstrumenter = parameterInstrumentator;
        this.returnInstrumenter = returnInstrumentator;
    }

    @Override
    public Void visitCompilationUnit(CompilationUnitTree node, Void aVoid) {
        ExpressionTree packageName = node.getPackageName();
        this.packageName = packageName == null ? "" : packageName.toString();
        Set<String> packageAnnotations = context.getPackageInfoManager().getPackageAnnotations(this.packageName);
        String location = this.packageName.isEmpty() ? "default package" : this.packageName + " package";
        return withDefaultNotNullAnnotations(packageAnnotations,
                                             location,
                                             () -> super.visitCompilationUnit(node, aVoid));
    }

    @Override
    public Void visitClass(ClassTree node, Void aVoid) {
        String className = node.getSimpleName().toString();
        if (className.isEmpty()) {
            className = "$" + ++anonymousClassCounter;
        }

        ModifiersTree modifiers = node.getModifiers();
        final boolean processingInterface;
        if (modifiers instanceof JCTree.JCModifiers) {
            processingInterface = (((JCTree.JCModifiers) modifiers).flags & Flags.INTERFACE) != 0;
        } else {
            processingInterface = modifiers.toString().contains("interface");
        }
        classNames.push(className);
        this.processingInterface.push(processingInterface);

        try {
            return withDefaultNotNullAnnotations(modifiers,
                                                 className + " class",
                                                 () -> super.visitClass(node, aVoid));
        } finally {
            classNames.pop();
            this.processingInterface.pop();
        }
    }

    private <T> T withDefaultNotNullAnnotations(@Nullable ModifiersTree modifiers,
                                               @NotNull String location,
                                               @NotNull Callable<T> action)
    {
        return withDefaultNotNullAnnotations(extractAnnotations(modifiers), location, action);
    }

    private <T> T withDefaultNotNullAnnotations(@NotNull Set<String> annotations,
                                               @NotNull String location,
                                               @NotNull Callable<T> action)
    {
        TrautePluginSettings settings = context.getPluginSettings();

        Optional<String> parameterNotNullByDefaultAnnotation = findMatch(
                annotations,
                settings.getNotNullByDefaultAnnotations(METHOD_PARAMETER)
        );
        parameterNotNullByDefaultAnnotation.ifPresent(s -> parametersNotNullByDefault.push(String.format(
                "%s annotation on the %s", s, location)));
        Optional<String> returnNotNullByDefaultAnnotation = findMatch(
                annotations,
                settings.getNotNullByDefaultAnnotations(METHOD_RETURN)
        );
        returnNotNullByDefaultAnnotation.ifPresent(s -> returnNotNullByDefault.push(String.format(
                "%s annotation on the %s", s, location)));
        try {
            return action.call();
        } catch (Exception e) {
            throw new IllegalStateException(e);
        } finally {
            if (parameterNotNullByDefaultAnnotation.isPresent()) {
                parametersNotNullByDefault.pop();
            }
            if (returnNotNullByDefaultAnnotation.isPresent()) {
                returnNotNullByDefault.pop();
            }
        }
    }

    @Override
    public Void visitImport(ImportTree node, Void v) {
        if (!node.isStatic()) {
            context.addImport(node.getQualifiedIdentifier().toString());
        }
        return v;
    }

    @Override
    public Void visitMethod(MethodTree method, Void v) {
        methodName = method.getName().toString();
        return withDefaultNotNullAnnotations(
                method.getModifiers(), getQualifiedMethodName() + " method", () -> {
                    instrumentReturnExpression = shouldInstrumentReturnExpression(method);
                    if (shouldInstrumentMethodParameters(method)) {
                        JCTree.JCBlock methodBody = getMethodBody(method);
                        if (methodBody != null) {
                            instrumentMethodParameters(method, methodBody);
                        }
                    }
                    try {
                        return super.visitMethod(method, v);
                    } finally {
                        methodReturnType = null;
                        methodNotNullAnnotation = null;
                        methodName = null;
                        instrumentReturnExpression = false;
                        tmpVariableCounter = 1;
                    }
                });
    }

    @SuppressWarnings("SimplifiableIfStatement")
    private boolean shouldInstrumentReturnExpression(@NotNull MethodTree method) {
        if (!processingInterface.isEmpty()
            && processingInterface.peek()
            && !hasFlag(method.getModifiers(), Modifier.DEFAULT, Modifier.STATIC))
        {
            return false;
        }
        return context.getPluginSettings().isEnabled(METHOD_RETURN) && mayBeInstrumentReturnType(method);
    }

    @SuppressWarnings("SimplifiableIfStatement")
    private boolean shouldInstrumentMethodParameters(@NotNull MethodTree method) {
        if (!processingInterface.isEmpty()
            && processingInterface.peek()
            && !hasFlag(method.getModifiers(), Modifier.DEFAULT, Modifier.STATIC)) {
            return false;
        }
        return context.getPluginSettings().isEnabled(METHOD_PARAMETER);
    }

    @SuppressWarnings("SimplifiableIfStatement")
    private static boolean hasFlag(@Nullable ModifiersTree modifiers, @NotNull Modifier ... targetModifiers) {
        if (modifiers == null) {
            return false;
        }
        Set<Modifier> flags = modifiers.getFlags();
        if (flags == null) {
            return false;
        }
        for (Modifier targetModifier : targetModifiers) {
            if (flags.contains(targetModifier)) {
                return true;
            }
        }
        return false;
    }

    @Nullable
    private JCTree.JCBlock getMethodBody(@NotNull MethodTree method) {
        if (hasFlag(method.getModifiers(), Modifier.ABSTRACT)) {
            return null;
        }
        BlockTree bodyBlock = method.getBody();
        if (bodyBlock == null) {
            return null;
        }
        if (bodyBlock instanceof JCTree.JCBlock) {
            return (JCTree.JCBlock) bodyBlock;
        }
        context.getLogger().reportDetails(String.format(
                "get a %s instance in the method AST but got %s",
                JCTree.JCBlock.class.getName(), bodyBlock.getClass().getName()
        ));
        return null;
    }

    private void instrumentMethodParameters(@NotNull MethodTree method, @NotNull JCTree.JCBlock bodyBlock) {
        SortedSet<ParameterToInstrumentInfo> variablesToCheck = new TreeSet<>(
                // There is a possible case that more than one method parameter is marked by a NotNull annotation.
                // We want to add null-checks in reverse order then, i.e. for the last parameter marked
                // by a NotNull, then for the previous before the last etc
                (o1, o2) -> o2.getMethodParameterIndex() - o1.getMethodParameterIndex()
        );
        int parameterIndex = -1;
        int parametersNumber = method.getParameters().size();
        for (VariableTree variable : method.getParameters()) {
            parameterIndex++;
            if (variable == null) {
                continue;
            }
            Tree type = variable.getType();
            if (type != null && PRIMITIVE_TYPES.contains(type.toString())) {
                continue;
            }
            Annotations annotations = findAnnotation(variable.getModifiers());
            if (annotations.notNull.isPresent()
                || (!parametersNotNullByDefault.isEmpty()) && !annotations.nullable.isPresent())
            {

                String notNullByDefaultAnnotationDescription =
                        parametersNotNullByDefault.isEmpty() ? null : parametersNotNullByDefault.peek();
                variablesToCheck.add(new ParameterToInstrumentInfo(context,
                                                                   annotations.notNull.orElse(null),
                                                                   notNullByDefaultAnnotationDescription,
                                                                   variable,
                                                                   bodyBlock,
                                                                   getQualifiedMethodName(),
                                                                   parameterIndex,
                                                                   parametersNumber,
                                                                   method.getReturnType() == null));
            }
        }

        for (ParameterToInstrumentInfo info : variablesToCheck) {
            mayBeSetPosition(info.getMethodParameter(), context.getAstFactory());
            parameterInstrumenter.instrument(info);
        }
    }

    private boolean mayBeInstrumentReturnType(@NotNull MethodTree method) {
        Tree returnType = method.getReturnType();
        if (returnType == null
            || METHOD_RETURN_TYPES_TO_SKIP.contains(returnType.toString())
            || (!(returnType instanceof JCTree.JCExpression)))
        {
            return false;
        }

        Annotations annotations = findAnnotation(method.getModifiers());
        if (annotations.notNull.isPresent()
            || (!returnNotNullByDefault.isEmpty() && !annotations.nullable.isPresent()))
        {
            methodNotNullAnnotation = annotations.notNull.orElse(null);
            methodReturnType = (JCTree.JCExpression) returnType;
            return true;
        }
        return false;
    }

    @NotNull
    private String getTmpVariableName() {
        return "tmpTrauteVar" + ++tmpVariableCounter;
    }

    private void mayBeSetPosition(@NotNull Tree astNode, @NotNull TreeMaker astFactory) {
        if (astNode instanceof JCTree) {
            // Mark our AST factory with the given AST node's offset in order to see corresponding
            // line in the stack trace when an NPE is thrown.
            astFactory.at(((JCTree) astNode).pos);
        }
    }

    /**
     * Checks if given {@code AST} element's modifiers contain any of the
     * {@link TrautePluginSettings#getNotNullAnnotations() NotNull} or
     * or {@link TrautePluginSettings#getNullableAnnotations() Nullable} annotation.
     *
     * @param modifiers {@code AST} element's modifiers to check
     * @return          annotations lookup result
     */
    @NotNull
    private Annotations findAnnotation(@Nullable ModifiersTree modifiers) {
        Set<String> annotationsInSource = extractAnnotations(modifiers);
        if (annotationsInSource.isEmpty()) {
            return Annotations.EMPTY;
        }
        return new Annotations(findMatch(annotationsInSource, context.getPluginSettings().getNotNullAnnotations()),
                               findMatch(annotationsInSource, context.getPluginSettings().getNullableAnnotations()));
    }

    @NotNull
    private Set<String> extractAnnotations(@Nullable ModifiersTree modifiers) {
        if (modifiers == null) {
            return emptySet();
        }
        java.util.List<? extends AnnotationTree> annotations = modifiers.getAnnotations();
        if (annotations == null) {
            return emptySet();
        }
        Set<String> result = new HashSet<>();
        for (AnnotationTree annotation : annotations) {
            Tree type = annotation.getAnnotationType();
            if (type != null) {
                result.add(type.toString());
            }
        }
        return result;
    }

    /**
     * <p>
     *     Checks if any of the given 'annotations to check' matches any of the
     *     {@code target annotations}
     *     considering {@link CompilationUnitProcessingContext#getImports() available imports}.
     * </p>
     * <p>
     *     Example:
     *     <ul>
     *       <li>annotations to check: [ {@code NotNull} ]</li>
     *       <li>imports: [ {@code org.jetbrains.annotations.NotNull} ]</li>
     *       <li>target annotations: [ {@code org.jetbrains.annotations.NotNull} ]</li>
     *     </ul>
     *     We expect to find a match for the {@code org.jetbrains.annotations.NotNull} then.
     * </p>
     *
     * @param annotationsToCheck    annotations to match against the given 'target annotations'
     * @param targetAnnotations     target annotations to check against the given {@code annotations to check}
     * @return                      a matched annotation (if any)
     */
    @NotNull
    private Optional<String> findMatch(@NotNull Collection<String> annotationsToCheck,
                                       @NotNull Set<String> targetAnnotations)
    {
        for (String annotationInSource : annotationsToCheck) {
            if (targetAnnotations.contains(annotationInSource)) {
                // Qualified annotation, like 'void test(@javax.annotation.Nonnul String s) {}'
                return Optional.of(annotationInSource);
            }
            if (packageName != null) {
                String candidate = String.format("%s.%s", packageName, annotationInSource);
                if (targetAnnotations.contains(candidate)) {
                    return Optional.of(candidate);
                }
            }
            for (String anImport : context.getImports()) {
                // Support an import like 'import org.jetbrains.annotations.*;'
                if (anImport.endsWith(".*")) {
                    String candidate = anImport.substring(0, anImport.length() - 1) + annotationInSource;
                    if (targetAnnotations.contains(candidate)) {
                        return Optional.of(candidate);
                    }
                    continue;
                }
                if (!targetAnnotations.contains(anImport)) {
                    continue;
                }
                if (anImport.endsWith(annotationInSource)) {
                    return Optional.of(anImport);
                }
            }
        }
        return Optional.empty();
    }

    @Nullable
    private String getQualifiedMethodName() {
        StringBuilder buffer = new StringBuilder();
        if (packageName != null) {
            buffer.append(packageName).append(".");
        }
        if (!classNames.isEmpty()) {
            List<String> list = new ArrayList<>(classNames);
            for (String className : list) {
                if (className.startsWith("$")) {
                    // We want to show class name like 'MyCLass$1' instead of 'MyClass$1'
                    buffer.setLength(buffer.length() - 1);
                }
                buffer.append(className).append(".");
            }
        }
        if (methodName == null || buffer.length() == 0) {
            return null;
        }
        buffer.append(methodName);
        return buffer.toString();
    }

    @Override
    public Void visitBlock(BlockTree node, Void aVoid) {
        parents.push(node);
        try {
            return super.visitBlock(node, aVoid);
        } finally {
            parents.pop();
        }
    }

    @Override
    public Void visitIf(IfTree node, Void aVoid) {
        parents.push(node);
        try {
            return super.visitIf(node, aVoid);
        } finally {
            parents.pop();
        }
    }

    @Override
    public Void visitForLoop(ForLoopTree node, Void aVoid) {
        parents.push(node);
        try {
            return super.visitForLoop(node, aVoid);
        } finally {
            parents.pop();
        }
    }

    @Override
    public Void visitEnhancedForLoop(EnhancedForLoopTree node, Void aVoid) {
        parents.push(node);
        try {
            return super.visitEnhancedForLoop(node, aVoid);
        } finally {
            parents.pop();
        }
    }

    @Override
    public Void visitWhileLoop(WhileLoopTree node, Void aVoid) {
        parents.push(node);
        try {
            return super.visitWhileLoop(node, aVoid);
        } finally {
            parents.pop();
        }
    }

    @Override
    public Void visitDoWhileLoop(DoWhileLoopTree node, Void aVoid) {
        parents.push(node);
        try {
            return super.visitDoWhileLoop(node, aVoid);
        } finally {
            parents.pop();
        }
    }

    @Override
    public Void visitCase(CaseTree node, Void aVoid) {
        parents.push(node);
        try {
            return super.visitCase(node, aVoid);
        } finally {
            parents.pop();
        }
    }

    @Override
    public Void visitReturn(ReturnTree node, Void aVoid) {
        if (instrumentReturnExpression
            && (methodNotNullAnnotation != null || !returnNotNullByDefault.isEmpty())
            && methodReturnType != null
            && !parents.isEmpty())
        {
            mayBeSetPosition(node, context.getAstFactory());
            String notNullByDefaultDescription = returnNotNullByDefault.isEmpty() ? null
                                                                                  : returnNotNullByDefault.peek();
            returnInstrumenter.instrument(new ReturnToInstrumentInfo(context,
                                                                     methodNotNullAnnotation,
                                                                     notNullByDefaultDescription,
                                                                     node,
                                                                     methodReturnType,
                                                                     getTmpVariableName(),
                                                                     parents.peek(),
                                                                     getQualifiedMethodName()));
        }
        return super.visitReturn(node, aVoid);
    }

    @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
    private static class Annotations {

        public static final Annotations EMPTY = new Annotations(Optional.empty(), Optional.empty());

        @NotNull public final Optional<String> notNull;
        @NotNull public final Optional<String> nullable;

        public Annotations(@NotNull Optional<String> notNull, @NotNull Optional<String> nullable) {
            this.notNull = notNull;
            this.nullable = nullable;
        }
    }
} public class PackageInfoManager {

    private static final String SUFFIX = TrauteConstants.PACKAGE_INFO + JavaFileObject.Kind.SOURCE.extension;

    private final ConcurrentMap<String/* package name */, Set<String> /* annotations */> packageAnnotations
            = new ConcurrentHashMap<>();

    @NotNull
    public Set<String> getPackageAnnotations(@NotNull String packageName) {
        return packageAnnotations.computeIfAbsent(packageName, p -> {
            try {
                Class<?> clazz = Class.forName(packageName + "." + TrauteConstants.PACKAGE_INFO);
                Annotation[] annotations = clazz.getAnnotations();
                if (annotations != null) {
                    return Arrays.stream(annotations).map(a -> a.annotationType().getName()).collect(toSet());
                }
            } catch (ClassNotFoundException ignore) {
            }
            return emptySet();
        });
    }

    public void onCompilationUnit(@NotNull CompilationUnitTree compilationUnit) {
        JavaFileObject sourceFile = compilationUnit.getSourceFile();
        if (sourceFile == null) {
            return;
        }
        String name = sourceFile.getName();
        if (name == null) {
            return;
        }
        if (!name.endsWith(SUFFIX)) {
            return;
        }
        ExpressionTree packageNameExpression = compilationUnit.getPackageName();
        String packageName = packageNameExpression == null ? "" : packageNameExpression.toString();
        Collection<String> packageAnnotationsInSource = compilationUnit.getPackageAnnotations()
                                                                       .stream()
                                                                       .map(AnnotationTree::getAnnotationType)
                                                                       .filter(Objects::nonNull)
                                                                       .map(Object::toString)
                                                                       .collect(toList());
        Set<String> imports = new HashSet<>();
        compilationUnit.accept(new TreeScanner<Void, Void>() {
            @Override
            public Void visitImport(ImportTree node, Void aVoid) {
                if (!node.isStatic()) {
                    Tree identifier = node.getQualifiedIdentifier();
                    if (identifier != null) {
                        imports.add(identifier.toString());
                    }
                }
                return null;
            }
        }, null);
        Set<String> resultingPackageAnnotations = new HashSet<>();
        // Filter out all exact imports like 'import java.util.List'
        for (String anImport : imports) {
            if (anImport.endsWith(".*")) {
                continue;
            }
            for (String packageAnnotation : packageAnnotationsInSource) {
                if (anImport.endsWith(packageAnnotation)) {
                    resultingPackageAnnotations.add(anImport);
                    packageAnnotationsInSource.remove(packageAnnotation);
                    break;
                }
            }
        }

        // Process wildcard imports
        for (String anImport : imports) {
            if (anImport.endsWith(".*")) {
                String prefix = anImport.substring(0, anImport.length() - 1);
                for (String packageAnnotation : packageAnnotationsInSource) {
                    resultingPackageAnnotations.add(prefix + packageAnnotation);
                }
            }
        }
        resultingPackageAnnotations.addAll(packageAnnotationsInSource);
        packageAnnotations.computeIfAbsent(packageName, p -> new HashSet<>()).addAll(resultingPackageAnnotations);
    }
} public interface Instrumentator<T extends InstrumentationInfo> {

    /**
     * Performs instrumentation for the given data.
     *
     * @param instrumentationInfo   instrumentation info
     */
    void instrument(@NotNull T instrumentationInfo);
} public class MethodReturnInstrumentator extends AbstractInstrumentator<ReturnToInstrumentInfo> {

    @Override
    protected boolean mayBeInstrument(@NotNull ReturnToInstrumentInfo info) {
        ReturnInstrumentationAstParent parent
                = info.getParent().accept(new MethodInstrumentationParentFinder(info), null);
        if (parent == null) {
            return false;
        }
        Optional<List<JCTree.JCStatement>> returnCheckOptional = buildReturnCheck(info);
        if (!returnCheckOptional.isPresent()) {
            return false;
        }
        List<JCTree.JCStatement> statements = parent.getStatements();
        for (int i = 0; i < statements.size(); i++) {
            JCTree.JCStatement statement = statements.get(i);
            if (statement == info.getReturnExpression()) {
                List<JCTree.JCStatement> newStatements = returnCheckOptional.get();
                for (int j = i + 1; j < statements.size(); j++) {
                    newStatements = newStatements.append(statements.get(j));
                }
                for (int j = i - 1; j >= 0; j--) {
                    newStatements = newStatements.prepend(statements.get(j));
                }
                parent.setStatements(newStatements);
                mayBeLogInstrumentation(info);
                return true;
            }
        }
        // When control flow reaches this place, that means that the AST parent doesn't contain any statments, so,
        // we just populate it with new instructions.
        parent.setStatements(returnCheckOptional.get());
        mayBeLogInstrumentation(info);
        return true;
    }

    @NotNull
    private static Optional<List<JCTree.JCStatement>> buildReturnCheck(@NotNull ReturnToInstrumentInfo info) {
        CompilationUnitProcessingContext context = info.getContext();
        ExpressionTree returnExpression = info.getReturnExpression().getExpression();
        if (!(returnExpression instanceof JCTree.JCExpression)) {
            context.getLogger().reportDetails(String.format(
                    "find a 'return' expression of type %s but got %s",
                    JCTree.JCExpression.class.getName(), returnExpression.getClass().getName()
            ));
            return Optional.empty();
        }
        JCTree.JCExpression returnJcExpression = (JCTree.JCExpression) returnExpression;

        TreeMaker factory = context.getAstFactory();
        Names symbolsTable = context.getSymbolsTable();
        ExceptionTextGenerator<ReturnToInstrumentInfo> generator =
                context.getExceptionTextGeneratorManager().getGenerator(METHOD_RETURN, context.getPluginSettings());
        String errorMessage = generator.generate(info);

        List<JCTree.JCStatement> result = List.of(
                factory.VarDef(
                        factory.Modifiers(0),
                        symbolsTable.fromString(info.getTmpVariableName()),
                        info.getReturnType(),
                        returnJcExpression
                )
        );
        String exceptionToThrow = info.getContext().getPluginSettings().getExceptionToThrow(METHOD_RETURN);
        result = result.append(InstrumentationUtil.buildVarCheck(factory,
                                                                 symbolsTable,
                                                                 info.getTmpVariableName(),
                                                                 errorMessage,
                                                                 exceptionToThrow));
        result = result.append(
                factory.Return(
                        factory.Ident(symbolsTable.fromString(info.getTmpVariableName()))));
        return Optional.of(result);
    }

    private void mayBeLogInstrumentation(@NotNull ReturnToInstrumentInfo info) {
        CompilationUnitProcessingContext context = info.getContext();
        if (context.getPluginSettings().isVerboseMode()) {
            String methodName = info.getQualifiedMethodName();
            if (methodName != null) {
                context.getLogger().info("added a null-check for 'return' expression in method " + methodName + "()");
            }
        }
    }
} public class ReturnToInstrumentInfo implements InstrumentationInfo {

    @NotNull private final CompilationUnitProcessingContext context;
    @NotNull private final ReturnTree                       returnExpression;
    @NotNull private final JCTree.JCExpression              returnType;
    @NotNull private final String                           tmpVariableName;
    @NotNull private final Tree                             parent;

    @Nullable private final String qualifiedMethodName;

    private final String notNullAnnotation;
    private final String notNullByDefaultAnnotationDescription;

    public ReturnToInstrumentInfo(@NotNull CompilationUnitProcessingContext context,
                                  @Nullable String notNullAnnotation,
                                  @Nullable String notNullByDefaultAnnotationDescription,
                                  @NotNull ReturnTree returnExpression,
                                  @NotNull JCTree.JCExpression returnType,
                                  @NotNull String tmpVariableName,
                                  @NotNull Tree parent,
                                  @Nullable String qualifiedMethodName)
    {
        if (notNullAnnotation == null && notNullByDefaultAnnotationDescription == null) {
            throw new IllegalArgumentException(String.format(
                    "Detected an invalid attempt to instrument a method return - either NotNull annotation or "
                    + "NotNullByDefault annotations are undefined. Method: %s()", qualifiedMethodName));
        }
        this.context = context;
        this.notNullAnnotation = notNullAnnotation;
        this.notNullByDefaultAnnotationDescription = notNullByDefaultAnnotationDescription;
        this.returnExpression = returnExpression;
        this.returnType = returnType;
        this.tmpVariableName = tmpVariableName;
        this.parent = parent;
        this.qualifiedMethodName = qualifiedMethodName;
    }

    @Override
    @NotNull
    public InstrumentationType getType() {
        return InstrumentationType.METHOD_RETURN;
    }

    @Override
    @NotNull
    public CompilationUnitProcessingContext getContext() {
        return context;
    }

    @Override
    public String getNotNullAnnotation() {
        return notNullAnnotation;
    }

    @Override
    public String getNotNullByDefaultAnnotationDescription() {
        return notNullByDefaultAnnotationDescription;
    }

    /**
     * @return  '{@code return}' {@code AST} expression to instrument
     */
    @NotNull
    public ReturnTree getReturnExpression() {
        return returnExpression;
    }

    /**
     * @return  target method's return type
     */
    @NotNull
    public JCTree.JCExpression getReturnType() {
        return returnType;
    }

    /**
     * <p>Allows to get a name of the variable to use for storing {@code return}'s expression value (if necessary).</p>
     * <p>E.g. we don't need a temp variable in code like below:</p>
     * <pre>
     * String s = ...;
     * ...
     * return s;
     * </pre>
     * <p>I.e resulting code might be written as follows:</p>
     * <pre>
     * String s = ...;
     * ...
     * if (s == null) {
     *     throw new NullPointerException("...");
     * }
     * return s;
     *     </pre>
     * <p>However, there is a possible situation that we do need a temp variable:</p>
     * <pre>
     * return compute();
     * </pre>
     * <p>This code should be transformed at the following way:</p>
     * <pre>
     * String tmpVar = compute();
     * if (tmpVar == null) {
     *     throw new NullPointerException("...");
     * }
     * return tmpVar;
     * </pre>
     *
     * @return  name of the variable to use for storing {@code return}'s expression value (if necessary)
     */
    @NotNull
    public String getTmpVariableName() {
        return tmpVariableName;
    }

    /**
     * @return  parent {@code AST} element for the target {@code return} expression to check
     */
    @NotNull
    public Tree getParent() {
        return parent;
    }

    /**
     * @return  qualified method name which {@code NotNull} parameter should be instrumented
     *          (if that information is available)
     */
    @Nullable
    public String getQualifiedMethodName() {
        return qualifiedMethodName;
    }
} public class ParameterInstrumentator extends AbstractInstrumentator<ParameterToInstrumentInfo> {

    @Override
    protected boolean mayBeInstrument(@NotNull ParameterToInstrumentInfo info) {
        String parameterName = info.getMethodParameter().getName().toString();
        CompilationUnitProcessingContext context = info.getContext();
        ExceptionTextGenerator<ParameterToInstrumentInfo> generator =
                context.getExceptionTextGeneratorManager().getGenerator(METHOD_PARAMETER, context.getPluginSettings());
        String errorMessage = generator.generate(info);
        TreeMaker factory = context.getAstFactory();
        Names symbolsTable = context.getSymbolsTable();
        JCTree.JCBlock body = info.getBody();
        String exceptionToThrow = info.getContext().getPluginSettings().getExceptionToThrow(METHOD_PARAMETER);
        JCTree.JCIf varCheck = buildVarCheck(factory, symbolsTable, parameterName, errorMessage, exceptionToThrow);
        if (info.isConstructor() && isFirstStatementThisOrSuperCall(body)) {
            List<JCTree.JCStatement> newStatements = List.of(varCheck);
            List<JCTree.JCStatement> statements = body.getStatements();
            for (int i = 1; i < statements.size(); i++) {
                newStatements = newStatements.append(statements.get(i));
            }
            newStatements = newStatements.prepend(statements.get(0));
            body.stats = newStatements;
        } else {
            body.stats = body.stats.prepend(varCheck);
        }

        if (context.getPluginSettings().isVerboseMode()) {
            String methodName = info.getQualifiedMethodName();
            String methodNotice = methodName == null ? "" : " in the method " + methodName + "()";
            context.getLogger().info(String.format(
                    "added a null-check for argument '%s'%s",
                    parameterName, methodNotice
            ));
        }
        return true;
    }

    private static boolean isFirstStatementThisOrSuperCall(@NotNull JCTree.JCBlock body) {
        List<JCTree.JCStatement> statements = body.getStatements();
        if (statements.isEmpty()) {
            return false;
        }
        JCTree.JCStatement expressionCandidate = statements.get(0);
        if (expressionCandidate instanceof ExpressionStatementTree) {
            ExpressionStatementTree expression = (ExpressionStatementTree) expressionCandidate;
            ExpressionTree methodInvocationCandidate = expression.getExpression();
            if (methodInvocationCandidate instanceof MethodInvocationTree) {
                MethodInvocationTree methodInvocation = (MethodInvocationTree) methodInvocationCandidate;
                ExpressionTree methodSelect = methodInvocation.getMethodSelect();
                if (methodSelect != null) {
                    String select = methodSelect.toString();
                    return "this".equals(select) || "super".equals(select);
                }
            }
        }
        return false;
    }
} public class ParameterToInstrumentInfo implements InstrumentationInfo {

    @NotNull private final CompilationUnitProcessingContext compilationUnitProcessingContext;
    @NotNull private final VariableTree                     methodParameter;
    @NotNull private final JCTree.JCBlock                   body;

    private final String notNullAnnotation;
    private final String notNullByDefaultAnnotationDescription;

    @Nullable private final String qualifiedMethodName;

    private final int     methodParameterIndex;
    private final int     methodParametersNumber;
    private final boolean constructor;

    public ParameterToInstrumentInfo(@NotNull CompilationUnitProcessingContext compilationUnitProcessingContext,
                                     @Nullable String notNullAnnotation,
                                     @Nullable String notNullByDefaultAnnotationDescription,
                                     @NotNull VariableTree methodParameter,
                                     @NotNull JCTree.JCBlock body,
                                     @Nullable String qualifiedMethodName,
                                     int methodParameterIndex,
                                     int methodParametersNumber,
                                     boolean constructor)
    {
        if (notNullAnnotation == null && notNullByDefaultAnnotationDescription == null) {
            throw new IllegalArgumentException(String.format(
                    "Detected an invalid attempt to instrument a method parameter - either NotNull annotation or "
                    + "NotNullByDefault annotations are undefined. Method: %s(), parameter: %s",
                    qualifiedMethodName, methodParameter.getName()));
        }
        this.compilationUnitProcessingContext = compilationUnitProcessingContext;
        this.notNullAnnotation = notNullAnnotation;
        this.notNullByDefaultAnnotationDescription = notNullByDefaultAnnotationDescription;
        this.methodParameter = methodParameter;
        this.body = body;
        this.qualifiedMethodName = qualifiedMethodName;
        this.methodParameterIndex = methodParameterIndex;
        this.methodParametersNumber = methodParametersNumber;
        this.constructor = constructor;
    }

    @Override
    @NotNull
    public InstrumentationType getType() {
        return InstrumentationType.METHOD_PARAMETER;
    }

    @Override
    @NotNull
    public CompilationUnitProcessingContext getContext() {
        return compilationUnitProcessingContext;
    }

    @Override
    public String getNotNullAnnotation() {
        return notNullAnnotation;
    }

    @Override
    public String getNotNullByDefaultAnnotationDescription() {
        return notNullByDefaultAnnotationDescription;
    }

    /**
     * @return {@code AST} element for the method parameter marked by the {@code NotNull} annotation
     */
    @NotNull
    public VariableTree getMethodParameter() {
        return methodParameter;
    }

    /**
     * @return body of the method which target parameter is marked by the {@code NotNull} annotation
     */
    @NotNull
    public JCTree.JCBlock getBody() {
        return body;
    }

    /**
     * @return  qualified method name which {@code NotNull} parameter should be instrumented
     *          (if that information is available)
     */
    @Nullable
    public String getQualifiedMethodName() {
        return qualifiedMethodName;
    }

    /**
     * @return target parameter's index (zero-based)
     */
    public int getMethodParameterIndex() {
        return methodParameterIndex;
    }

    /**
     * @return total number of the target method's parameters
     */
    public int getMethodParametersNumber() {
        return methodParametersNumber;
    }

    /**
     * @return  {@code true} if target method is a constructor
     */
    public boolean isConstructor() {
        return constructor;
    }
} public abstract class AbstractLogger implements TrautePluginLogger {

    /**
     * <p>
     *     There is a possible case that javac implementation is changed not in a backward compatible way in
     *     a new release. This plugin might stop working then (e.g. new approach should be used for fetching
     *     {@code AST} builder).
     * </p>
     * <p>
     *     We don't want to generate numerous errors for the compilation then but report at most once.
     * </p>
     * <p>
     *     This flag allows to check if a problem has already been reported.
     * </p>
     */
    private boolean problemReported;

    /**
     * Delegates to the {@link #report(String)} but adds more generic info to the given problem details
     *
     * @param problemDetails details of the problem to report
     */
    @Override
    public void reportDetails(@NotNull String problemDetails) {
        report(getProblemMessage(problemDetails));
    }

    /**
     * Shows given problem message to the end-user if necessary.
     *
     * @param message   a problem message to show
     */
    @Override
    public void report(@NotNull String message) {
        // Do not report a problem more than once
        if (!problemReported) {
            warn(message);
            problemReported = true;
        }
    }

    protected abstract void warn(@NotNull String message);

    @NotNull
    public abstract Object getKey();

    /**
     * Prepares a problem message to show end-user in case the plugin can't do the job during compilation.
     *
     * @param details   exact problem details, will be appended to the general problem description suffix
     * @return          a problem message to use
     */
    @NotNull
    public static String getProblemMessage(@NotNull String details) {
        return String.format(
                "NotNull-instrumentation failed, it might be that javac implementation has significantly changed "
                + "- '%s' javac plugin expected to %s. %s", TrauteConstants.PLUGIN_NAME, details, getProblemMessageSuffix()
        );
    }

    @NotNull
    public static String getProblemMessageSuffix() {
        return "Please contact the plugin's author via traute.java@gmail.com";
    }
} public class CompilerOutputLogger extends AbstractLogger {

    private static final String NOTICE_PREFIX = String.format("[%s javac plugin]: ", TrauteConstants.PLUGIN_NAME);

    @NotNull private final Log log;

    public CompilerOutputLogger(@NotNull Log log) {
        this.log = log;
    }

    @Override
    @NotNull
    public Log getKey() {
        return log;
    }

    @Override
    public void info(@NotNull String message) {
        log.printRawLines(Log.WriterKind.NOTICE, NOTICE_PREFIX + message);
    }

    @Override
    public void warn(@NotNull String message) {
        log.rawWarning(-1, message);
    }
} public class FileLogger extends AbstractLogger {

    @NotNull private final File file;

    public FileLogger(@NotNull File file) {
        ensureFileExists(file);
        this.file = file;
    }

    private static void ensureFileExists(@NotNull File file) {
        File dir = file.getParentFile();
        if (dir.isDirectory()) {
            return;
        }
        boolean created = dir.mkdirs();
        if (!created) {
            throw new IllegalStateException(String.format("Can't create log file '%s'", file.getAbsolutePath()));
        }
    }

    @Override
    @NotNull
    public Object getKey() {
        return file.getAbsolutePath().intern();
    }

    @Override
    protected void warn(@NotNull String message) {
        println(message, "WARN");
    }

    @Override
    public void info(@NotNull String message) {
        println(message, "INFO");
    }

    private void println(@NotNull String message, @NotNull String logLevel) {
        try {
            try (BufferedWriter writer = new BufferedWriter(new FileWriter(file, true))) {
                writer.write('[');
                writer.write(logLevel);
                writer.write("] ");
                writer.write(message);
                writer.newLine();
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }
} public interface TrautePluginLogger {

    void info(@NotNull String message);

    /**
     * Delegates to the {@link #report(String)} but adds more generic info to the given problem details
     *
     * @param problemDetails details of the problem to report
     */
    void reportDetails(@NotNull String problemDetails);

    /**
     * Asks to report given problem message (the message is not modified and reported as-is).
     *
     * @param message   a message to report
     */
    void report(@NotNull String message);
} public class ExceptionTextGeneratorManager {

    private static final Map<InstrumentationType, ExceptionTextGenerator<?>> DEFAULT_GENERATORS = new HashMap<>();
    static {
        DEFAULT_GENERATORS.put(InstrumentationType.METHOD_PARAMETER, new DefaultParameterExceptionTextGenerator());
        DEFAULT_GENERATORS.put(InstrumentationType.METHOD_RETURN, new DefaultReturnExceptionTextGenerator());
        if (DEFAULT_GENERATORS.size() != InstrumentationType.values().length) {
            throw new RuntimeException(String.format(
                    "Default exception text generators for failed checks are not registered for all "
                    + "instrumentation types. All types: %s, registered types: %s",
                    Arrays.toString(InstrumentationType.values()), DEFAULT_GENERATORS.keySet()
            ));
        }
    }

    private final ConcurrentMap<InstrumentationType, ExceptionTextGenerator<?>> generators = new ConcurrentHashMap<>();

    private final ExceptionTextGeneratorFactory factory = new ExceptionTextGeneratorFactory();

    @NotNull private final TrautePluginLogger logger;

    public ExceptionTextGeneratorManager(@NotNull TrautePluginLogger logger) {
        this.logger = logger;
    }

    @SuppressWarnings("unchecked")
    @NotNull
    public <T> ExceptionTextGenerator<T> getGenerator(@NotNull InstrumentationType type,
                                                      @NotNull TrautePluginSettings settings)
    {
        return (ExceptionTextGenerator<T>) generators.computeIfAbsent(type, k -> {
            String pattern = settings.getExceptionTextPattern(type);
            if (pattern == null) {
                return DEFAULT_GENERATORS.get(type);
            }
            Optional<ExceptionTextGenerator<?>> o = factory.build(type, pattern, logger);
            return o.orElse(DEFAULT_GENERATORS.get(type));
        });
    }
} @NotNull
public static TrautePluginSettingsBuilder settingsBuilder() {
    return new TrautePluginSettingsBuilder();
} public static final String OPTION_LOG_VERBOSE = "traute.log.verbose"; public static final String OPTION_PREFIX_ANNOTATIONS_NOT_NULL_BY_DEFAULT = "traute.annotations.not.null.by.default."; public static final String SEPARATOR = ":"; @NotNull
public static String getProblemMessageSuffix() {
    return "Please contact the plugin's author via traute.java@gmail.com";
}
import com.sun.source.tree.CompilationUnitTree;
import com.sun.source.util.JavacTask;
import com.sun.source.util.Plugin;
import com.sun.source.util.TaskEvent;
import com.sun.source.util.TaskListener;
import com.sun.tools.javac.api.BasicJavacTask;
import com.sun.tools.javac.processing.JavacProcessingEnvironment;
import com.sun.tools.javac.tree.TreeMaker;
import com.sun.tools.javac.util.Context;
import com.sun.tools.javac.util.Log;
import com.sun.tools.javac.util.Names;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import tech.harmonysoft.oss.traute.common.instrumentation.InstrumentationType;
import tech.harmonysoft.oss.traute.common.settings.TrautePluginSettings;
import tech.harmonysoft.oss.traute.common.settings.TrautePluginSettingsBuilder;
import tech.harmonysoft.oss.traute.common.stats.StatsCollector;
import tech.harmonysoft.oss.traute.common.util.TrauteConstants;
import tech.harmonysoft.oss.traute.javac.common.CompilationUnitProcessingContext;
import tech.harmonysoft.oss.traute.javac.common.InstrumentationApplianceFinder;
import tech.harmonysoft.oss.traute.javac.common.PackageInfoManager;
import tech.harmonysoft.oss.traute.javac.instrumentation.Instrumentator;
import tech.harmonysoft.oss.traute.javac.instrumentation.method.MethodReturnInstrumentator;
import tech.harmonysoft.oss.traute.javac.instrumentation.method.ReturnToInstrumentInfo;
import tech.harmonysoft.oss.traute.javac.instrumentation.parameter.ParameterInstrumentator;
import tech.harmonysoft.oss.traute.javac.instrumentation.parameter.ParameterToInstrumentInfo;
import tech.harmonysoft.oss.traute.javac.log.AbstractLogger;
import tech.harmonysoft.oss.traute.javac.log.CompilerOutputLogger;
import tech.harmonysoft.oss.traute.javac.log.FileLogger;
import tech.harmonysoft.oss.traute.javac.log.TrautePluginLogger;
import tech.harmonysoft.oss.traute.javac.text.ExceptionTextGeneratorManager;
import javax.tools.JavaFileObject;
import java.io.File;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.lang.ref.WeakReference;
import java.lang.reflect.Field;
import java.util.*;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import static java.lang.reflect.Modifier.*;
import static java.util.stream.Collectors.joining;
import static java.util.stream.Collectors.toList;
import static tech.harmonysoft.oss.traute.common.settings.TrautePluginSettingsBuilder.settingsBuilder;
import static tech.harmonysoft.oss.traute.common.util.TrauteConstants.OPTION_LOG_VERBOSE;
import static tech.harmonysoft.oss.traute.common.util.TrauteConstants.OPTION_PREFIX_ANNOTATIONS_NOT_NULL_BY_DEFAULT;
import static tech.harmonysoft.oss.traute.common.util.TrauteConstants.SEPARATOR;
import static tech.harmonysoft.oss.traute.javac.log.AbstractLogger.getProblemMessageSuffix;
package tech.harmonysoft.oss.traute.javac;




/**
 * <p>A {@code javac} plugin which inserts {@code null}-checks for target method arguments and returns from method.</p>
 * <p><b><i>Method argument check example</i></b></p>
 * <p>Consider the sources below:</p>
 * <pre>
 * public void service(&#064;NotNull Data data) {
 *     // Method instructions
 * }
 * </pre>
 * <p>
 *     When this code is compiled with the current plugin enabled, resulting binary looks like if it's compiled
 *     from a source below:
 * </p>
 * <pre>
 * public void serve(&#064;NotNull Data data) {
 *     if (data == null) {
 *         throw new NullPointerException("Argument 's' of type Data is declared as &#064;NotNull but got null for it");
 *     }
 *     // Method instructions
 * }
 * </pre>
 * <i>Note: exact message text is slightly different in a way that it provides more details about the problem.</i>
 * <p><b><i>Method return type example</i></b></p>
 * <p>Consider the source below:</p>
 * <pre>
 * &#064;NotNull
 * public Data fetch() {
 *     return dao.fetch();
 * }
 * </pre>
 * <p>
 *     When it's compiled with the current plugin enabled, resulting binary looks like if it's compiled
 *     from a source below:
 * </p>
 * <pre>
 * &#064;NotNull
 * public Data fetch() {
 *     Data tmpVar1 = dao.fetch();
 *     if (tmpVar1 == null) {
 *         throw new NullPointerException("Detected an attempt to return null from a method marked by &#064;NotNull");
 *     }
 *     return tmpVar1;
 * }
 * </pre>
 * <i>Note: exact message text is slightly different in a way that it provides more details about the problem.</i>
 */
public class TrauteJavacPlugin implements Plugin {

    private final AtomicReference<WeakReference<AbstractLogger>> loggerRef             = new AtomicReference<>();
    private final AtomicReference<TrautePluginSettings>          pluginSettingsRef     = new AtomicReference<>();

    private final Instrumentator<ParameterToInstrumentInfo> parameterInstrumentator = new ParameterInstrumentator();
    private final Instrumentator<ReturnToInstrumentInfo>    methodInstrumentator    = new MethodReturnInstrumentator();
    private final Set<String>                               pluginOptionKeys        = new HashSet<>();
    private final PackageInfoManager                        packageInfoManager      = new PackageInfoManager();