External Library Models Integration (#922)

The newly added `library-model` module consists of a CLI process that
takes an input directory with annotated java source files as a command
line parameter and uses `com.github.javaparser` APIS to generate
`libmodels.astubx` file containing method stubs for methods that return
@Nullable. This can be run using the existing `JarInferEnabled` and
`JarInferUseReturnAnnotations` flags.

This allows us to be able catch issues as shown in the below example
from externally annotated source code:

```java
@NullMarked
public class AnnotationExample {
    @Nullable
    public String makeUpperCase(String inputString) {
        if (inputString == null || inputString.isEmpty()) {
            return null;
        } else {
            return inputString.toUpperCase();
        }
    }
}

```

```java
class Test {
    static AnnotationExample annotationExample = new AnnotationExample();
    static void test(String value){}
    static void testPositive() {
        // BUG: Diagnostic contains: passing @Nullable parameter 'annotationExample.makeUpperCase(\"nullaway\")'
        test(annotationExample.makeUpperCase(\"nullaway\"));
    }    
}
```

---------

Co-authored-by: Manu Sridharan <msridhar@gmail.com>
Co-authored-by: Lázaro Clapp <lazaro.clapp@gmail.com>

21 files changed, 743 insertions, 39 deletions

diff --git a/build.gradle b/build.gradle
index 43ca0ba..5bf4aeb 100644
--- a/build.gradle
+++ b/build.gradle
@@ -95,9 +95,9 @@ subprojects { project ->
         google()
     }
 
-    // For some reason, spotless complains when applied to the jar-infer folder itself, even
-    // though there is no top-level :jar-infer project
-    if (project.name != "jar-infer") {
+    // Spotless complains when applied to the folders containing projects
+    // when they do not have a build.gradle file
+    if (project.name != "jar-infer" && project.name != "library-model") {
         project.apply plugin: "com.diffplug.spotless"
         spotless {
             java {
diff --git a/code-coverage-report/build.gradle b/code-coverage-report/build.gradle
index 3e5420f..939178e 100644
--- a/code-coverage-report/build.gradle
+++ b/code-coverage-report/build.gradle
@@ -80,4 +80,6 @@ dependencies {
     implementation project(':jar-infer:nullaway-integration-test')
     implementation project(':guava-recent-unit-tests')
     implementation project(':jdk-recent-unit-tests')
+    implementation project(':library-model:library-model-generator')
+    implementation project(':library-model:library-model-generator-integration-test')
 }
diff --git a/gradle/dependencies.gradle b/gradle/dependencies.gradle
index bf767fc..596096a 100755
--- a/gradle/dependencies.gradle
+++ b/gradle/dependencies.gradle
@@ -75,6 +75,7 @@ def build = [
     errorProneTestHelpersOld: "com.google.errorprone:error_prone_test_helpers:${oldestErrorProneVersion}",
     checkerDataflow         : "org.checkerframework:dataflow-nullaway:${versions.checkerFramework}",
     guava                   : "com.google.guava:guava:30.1-jre",
+    javaparser              : "com.github.javaparser:javaparser-core:3.25.8",
     javaxValidation         : "javax.validation:validation-api:2.0.1.Final",
     jspecify                : "org.jspecify:jspecify:0.3.0",
     jsr305Annotations       : "com.google.code.findbugs:jsr305:3.0.2",
diff --git a/jar-infer/jar-infer-cli/build.gradle b/jar-infer/jar-infer-cli/build.gradle
index 3245d66..9dafc4d 100644
--- a/jar-infer/jar-infer-cli/build.gradle
+++ b/jar-infer/jar-infer-cli/build.gradle
@@ -18,6 +18,7 @@ dependencies {
     implementation deps.build.commonscli
     implementation deps.build.guava
     implementation project(":jar-infer:jar-infer-lib")
+    implementation project(":library-model:library-model-generator")
 
     testImplementation deps.test.junit4
     testImplementation(deps.build.errorProneTestHelpers) {
diff --git a/jar-infer/jar-infer-lib/build.gradle b/jar-infer/jar-infer-lib/build.gradle
index 8ea4f9b..c0f6c34 100644
--- a/jar-infer/jar-infer-lib/build.gradle
+++ b/jar-infer/jar-infer-lib/build.gradle
@@ -37,6 +37,7 @@ dependencies {
     api deps.build.guava
     api deps.build.commonsIO
     compileOnly deps.build.errorProneCheckApi
+    implementation project(":library-model:library-model-generator")
 
     testImplementation deps.test.junit4
     testImplementation(deps.build.errorProneTestHelpers) {
diff --git a/jar-infer/jar-infer-lib/src/main/java/com/uber/nullaway/jarinfer/DefinitelyDerefedParamsDriver.java b/jar-infer/jar-infer-lib/src/main/java/com/uber/nullaway/jarinfer/DefinitelyDerefedParamsDriver.java
index fc85241..6616785 100644
--- a/jar-infer/jar-infer-lib/src/main/java/com/uber/nullaway/jarinfer/DefinitelyDerefedParamsDriver.java
+++ b/jar-infer/jar-infer-lib/src/main/java/com/uber/nullaway/jarinfer/DefinitelyDerefedParamsDriver.java
@@ -43,6 +43,8 @@ import com.ibm.wala.types.ClassLoaderReference;
 import com.ibm.wala.types.TypeReference;
 import com.ibm.wala.util.collections.Iterator2Iterable;
 import com.ibm.wala.util.config.FileOfClasses;
+import com.uber.nullaway.libmodel.MethodAnnotationsRecord;
+import com.uber.nullaway.libmodel.StubxWriter;
 import java.io.ByteArrayInputStream;
 import java.io.DataOutputStream;
 import java.io.File;
@@ -437,7 +439,7 @@ public class DefinitelyDerefedParamsDriver {
       }
       methodRecords.put(
           sign,
-          new MethodAnnotationsRecord(
+          MethodAnnotationsRecord.create(
               nullableReturns.contains(sign) ? ImmutableSet.of("Nullable") : ImmutableSet.of(),
               ImmutableMap.copyOf(argAnnotation)));
       nullableReturns.remove(sign);
@@ -445,7 +447,7 @@ public class DefinitelyDerefedParamsDriver {
     for (String nullableReturnMethodSign : Iterator2Iterable.make(nullableReturns.iterator())) {
       methodRecords.put(
           nullableReturnMethodSign,
-          new MethodAnnotationsRecord(ImmutableSet.of("Nullable"), ImmutableMap.of()));
+          MethodAnnotationsRecord.create(ImmutableSet.of("Nullable"), ImmutableMap.of()));
     }
     StubxWriter.write(out, importedAnnotations, packageAnnotations, typeAnnotations, methodRecords);
   }
diff --git a/jar-infer/jar-infer-lib/src/main/java/com/uber/nullaway/jarinfer/MethodAnnotationsRecord.java b/jar-infer/jar-infer-lib/src/main/java/com/uber/nullaway/jarinfer/MethodAnnotationsRecord.java
deleted file mode 100644
index 5f94b27..0000000
--- a/jar-infer/jar-infer-lib/src/main/java/com/uber/nullaway/jarinfer/MethodAnnotationsRecord.java
+++ /dev/null
@@ -1,26 +0,0 @@
-package com.uber.nullaway.jarinfer;
-
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableSet;
-
-/** A record describing the annotations associated with a java method and its arguments. */
-final class MethodAnnotationsRecord {
-  private final ImmutableSet<String> methodAnnotations;
-  // 0 means receiver
-  private final ImmutableMap<Integer, ImmutableSet<String>> argumentAnnotations;
-
-  MethodAnnotationsRecord(
-      ImmutableSet<String> methodAnnotations,
-      ImmutableMap<Integer, ImmutableSet<String>> argumentAnnotations) {
-    this.methodAnnotations = methodAnnotations;
-    this.argumentAnnotations = argumentAnnotations;
-  }
-
-  ImmutableSet<String> getMethodAnnotations() {
-    return methodAnnotations;
-  }
-
-  ImmutableMap<Integer, ImmutableSet<String>> getArgumentAnnotations() {
-    return argumentAnnotations;
-  }
-}
diff --git a/library-model/library-model-generator-cli/build.gradle b/library-model/library-model-generator-cli/build.gradle
new file mode 100644
index 0000000..94fd8a3
--- /dev/null
+++ b/library-model/library-model-generator-cli/build.gradle
@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2024. Uber Technologies
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+plugins {
+    id "java-library"
+    id "com.github.johnrengelman.shadow"
+}
+
+jar{
+    manifest {
+        attributes('Main-Class':'com.uber.nullaway.libmodel.LibraryModelGeneratorCLI')
+    }
+    // add this classifier so that the output file for the jar task differs from
+    // the output file for the shadowJar task (otherwise they overwrite each other's
+    // outputs, forcing the tasks to always re-run)
+    archiveClassifier = "nonshadow"
+}
+
+shadowJar {
+    mergeServiceFiles()
+    configurations = [
+        project.configurations.runtimeClasspath
+    ]
+    archiveClassifier = ""
+}
+shadowJar.dependsOn jar
+assemble.dependsOn shadowJar
+
+dependencies {
+    implementation project(":library-model:library-model-generator")
+}
diff --git a/library-model/library-model-generator-cli/src/main/java/com/uber/nullaway/libmodel/LibraryModelGeneratorCLI.java b/library-model/library-model-generator-cli/src/main/java/com/uber/nullaway/libmodel/LibraryModelGeneratorCLI.java
new file mode 100644
index 0000000..daea085
--- /dev/null
+++ b/library-model/library-model-generator-cli/src/main/java/com/uber/nullaway/libmodel/LibraryModelGeneratorCLI.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) 2024 Uber Technologies, Inc.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package com.uber.nullaway.libmodel;
+
+/**
+ * A CLI tool for invoking the process for {@link LibraryModelGenerator} which generates astubx
+ * file(s) from a directory containing annotated source code to be used as external library models.
+ */
+public class LibraryModelGeneratorCLI {
+  /**
+   * This is the main method of the cli tool. It parses the source files within a specified
+   * directory, obtains meaningful Nullability annotation information and writes it into an astubx
+   * file.
+   *
+   * @param args Command line arguments for the directory containing source files and the output
+   *     directory.
+   */
+  public static void main(String[] args) {
+    if (args.length != 2) {
+      System.out.println(
+          "Incorrect number of command line arguments. Required arguments:  <inputSourceDirectory> <outputDirectory>");
+      return;
+    }
+    LibraryModelGenerator libraryModelGenerator = new LibraryModelGenerator();
+    libraryModelGenerator.generateAstubxForLibraryModels(args[0], args[1]);
+  }
+}
diff --git a/library-model/library-model-generator-integration-test/build.gradle b/library-model/library-model-generator-integration-test/build.gradle
new file mode 100644
index 0000000..664e2bc
--- /dev/null
+++ b/library-model/library-model-generator-integration-test/build.gradle
@@ -0,0 +1,32 @@
+/*
+ * Copyright (C) 2024. Uber Technologies
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+plugins {
+    id "java-library"
+    id "nullaway.java-test-conventions"
+}
+
+dependencies {
+    testImplementation project(":nullaway")
+    testImplementation project(":library-model:test-library-model-generator")
+    testImplementation deps.test.junit4
+    testImplementation(deps.build.errorProneTestHelpers) {
+        exclude group: "junit", module: "junit"
+    }
+    implementation deps.build.guava
+    implementation deps.build.javaparser
+    compileOnly deps.apt.autoValueAnnot
+    annotationProcessor deps.apt.autoValue
+}
diff --git a/library-model/library-model-generator-integration-test/src/test/java/com/uber/nullaway/libmodel/LibraryModelIntegrationTest.java b/library-model/library-model-generator-integration-test/src/test/java/com/uber/nullaway/libmodel/LibraryModelIntegrationTest.java
new file mode 100644
index 0000000..c0149fd
--- /dev/null
+++ b/library-model/library-model-generator-integration-test/src/test/java/com/uber/nullaway/libmodel/LibraryModelIntegrationTest.java
@@ -0,0 +1,126 @@
+package com.uber.nullaway.libmodel;
+
+import com.google.errorprone.CompilationTestHelper;
+import com.uber.nullaway.NullAway;
+import java.util.Arrays;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+public class LibraryModelIntegrationTest {
+
+  @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+  private CompilationTestHelper compilationHelper;
+
+  @Before
+  public void setup() {
+    compilationHelper = CompilationTestHelper.newInstance(NullAway.class, getClass());
+  }
+
+  @Test
+  public void libraryModelNullableReturnsTest() {
+    compilationHelper
+        .setArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber",
+                "-XepOpt:NullAway:JarInferEnabled=true",
+                "-XepOpt:NullAway:JarInferUseReturnAnnotations=true"))
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import com.uber.nullaway.libmodel.AnnotationExample;",
+            "class Test {",
+            "  static AnnotationExample annotationExample = new AnnotationExample();",
+            "  static void test(String value){",
+            "  }",
+            "  static void testPositive() {",
+            "    // BUG: Diagnostic contains: passing @Nullable parameter 'annotationExample.makeUpperCase(\"nullaway\")'",
+            "    test(annotationExample.makeUpperCase(\"nullaway\"));",
+            "  }",
+            "  static void testNegative() {",
+            "    test(annotationExample.nullReturn());",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void libraryModelNullableReturnsArrayTest() {
+    compilationHelper
+        .setArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber",
+                "-XepOpt:NullAway:JarInferEnabled=true",
+                "-XepOpt:NullAway:JarInferUseReturnAnnotations=true"))
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import com.uber.nullaway.libmodel.AnnotationExample;",
+            "class Test {",
+            "  static AnnotationExample annotationExample = new AnnotationExample();",
+            "  static void test(Integer[] value){",
+            "  }",
+            "  static void testPositive() {",
+            "    // BUG: Diagnostic contains: passing @Nullable parameter 'annotationExample.generateIntArray(7)'",
+            "    test(annotationExample.generateIntArray(7));",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void libraryModelWithoutJarInferEnabledTest() {
+    compilationHelper
+        .setArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber"))
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import com.uber.nullaway.libmodel.AnnotationExample;",
+            "class Test {",
+            "  static AnnotationExample annotationExample = new AnnotationExample();",
+            "  static void test(String value){",
+            "  }",
+            "  static void testNegative() {",
+            "    // Since the JarInferEnabled and JarInferUseReturnAnnotations flags are not set, we don't get an error here",
+            "    test(annotationExample.makeUpperCase(\"nullaway\"));",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void libraryModelInnerClassNullableReturnsTest() {
+    compilationHelper
+        .setArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber",
+                "-XepOpt:NullAway:JarInferEnabled=true",
+                "-XepOpt:NullAway:JarInferUseReturnAnnotations=true"))
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import com.uber.nullaway.libmodel.AnnotationExample;",
+            "class Test {",
+            "  static AnnotationExample.InnerExample innerExample = new AnnotationExample.InnerExample();",
+            "  static void test(String value){",
+            "  }",
+            "  static void testPositive() {",
+            "    // BUG: Diagnostic contains: passing @Nullable parameter 'innerExample.returnNull()'",
+            "    test(innerExample.returnNull());",
+            "  }",
+            "}")
+        .doTest();
+  }
+}
diff --git a/library-model/library-model-generator/build.gradle b/library-model/library-model-generator/build.gradle
new file mode 100644
index 0000000..1d497fc
--- /dev/null
+++ b/library-model/library-model-generator/build.gradle
@@ -0,0 +1,26 @@
+/*
+ * Copyright (C) 2024. Uber Technologies
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+plugins {
+    id 'java-library'
+    id 'nullaway.java-test-conventions'
+}
+
+dependencies {
+    implementation deps.build.guava
+    implementation deps.build.javaparser
+    compileOnly deps.apt.autoValueAnnot
+    annotationProcessor deps.apt.autoValue
+}
diff --git a/library-model/library-model-generator/src/main/java/com/uber/nullaway/libmodel/LibraryModelGenerator.java b/library-model/library-model-generator/src/main/java/com/uber/nullaway/libmodel/LibraryModelGenerator.java
new file mode 100644
index 0000000..306a366
--- /dev/null
+++ b/library-model/library-model-generator/src/main/java/com/uber/nullaway/libmodel/LibraryModelGenerator.java
@@ -0,0 +1,265 @@
+/*
+ * Copyright (c) 2024 Uber Technologies, Inc.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package com.uber.nullaway.libmodel;
+
+import com.github.javaparser.ParseResult;
+import com.github.javaparser.ParserConfiguration.LanguageLevel;
+import com.github.javaparser.ast.CompilationUnit;
+import com.github.javaparser.ast.ImportDeclaration;
+import com.github.javaparser.ast.PackageDeclaration;
+import com.github.javaparser.ast.body.ClassOrInterfaceDeclaration;
+import com.github.javaparser.ast.body.MethodDeclaration;
+import com.github.javaparser.ast.expr.AnnotationExpr;
+import com.github.javaparser.ast.type.ArrayType;
+import com.github.javaparser.ast.type.ClassOrInterfaceType;
+import com.github.javaparser.ast.visitor.VoidVisitorAdapter;
+import com.github.javaparser.utils.CollectionStrategy;
+import com.github.javaparser.utils.ParserCollectionStrategy;
+import com.github.javaparser.utils.ProjectRoot;
+import com.github.javaparser.utils.SourceRoot;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import java.io.DataOutputStream;
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.Optional;
+
+/**
+ * Utilized for generating an astubx file from a directory containing annotated Java source code.
+ *
+ * <p>This class utilizes com.github.javaparser APIs to analyze Java source files within a specified
+ * directory. It processes the annotated Java source code to generate an astubx file that contains
+ * the required annotation information to be able to generate library models.
+ */
+public class LibraryModelGenerator {
+
+  public void generateAstubxForLibraryModels(String inputSourceDirectory, String outputDirectory) {
+    Map<String, MethodAnnotationsRecord> methodRecords = processDirectory(inputSourceDirectory);
+    writeToAstubx(outputDirectory, methodRecords);
+  }
+
+  /**
+   * Parses all the source files within the directory using javaparser.
+   *
+   * @param sourceDirectoryRoot Directory containing annotated java source files.
+   * @return a Map containing the Nullability annotation information from the source files.
+   */
+  private Map<String, MethodAnnotationsRecord> processDirectory(String sourceDirectoryRoot) {
+    Map<String, MethodAnnotationsRecord> methodRecords = new LinkedHashMap<>();
+    Path root = dirnameToPath(sourceDirectoryRoot);
+    AnnotationCollectorCallback ac = new AnnotationCollectorCallback(methodRecords);
+    CollectionStrategy strategy = new ParserCollectionStrategy();
+    // Required to include directories that contain a module-info.java, which don't parse by
+    // default.
+    strategy.getParserConfiguration().setLanguageLevel(LanguageLevel.JAVA_17);
+    ProjectRoot projectRoot = strategy.collect(root);
+
+    projectRoot
+        .getSourceRoots()
+        .forEach(
+            sourceRoot -> {
+              try {
+                sourceRoot.parse("", ac);
+              } catch (IOException e) {
+                throw new RuntimeException(e);
+              }
+            });
+    return methodRecords;
+  }
+
+  /**
+   * Writes the Nullability annotation information into the output directory as an astubx file.
+   *
+   * @param outputPath Output Directory.
+   * @param methodRecords Map containing the collected Nullability annotation information.
+   */
+  private void writeToAstubx(
+      String outputPath, Map<String, MethodAnnotationsRecord> methodRecords) {
+    if (methodRecords.isEmpty()) {
+      return;
+    }
+    Map<String, String> importedAnnotations =
+        ImmutableMap.of(
+            "NonNull", "org.jspecify.annotations.NonNull",
+            "Nullable", "org.jspecify.annotations.Nullable");
+    Path outputPathInstance = Paths.get(outputPath);
+    try {
+      Files.createDirectories(outputPathInstance.getParent());
+      try (DataOutputStream dos = new DataOutputStream(Files.newOutputStream(outputPathInstance))) {
+        StubxWriter.write(
+            dos,
+            importedAnnotations,
+            Collections.emptyMap(),
+            Collections.emptyMap(),
+            methodRecords);
+      }
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  public Path dirnameToPath(String dir) {
+    File f = new File(dir);
+    String absoluteDir = f.getAbsolutePath();
+    if (absoluteDir.endsWith("/.")) {
+      absoluteDir = absoluteDir.substring(0, absoluteDir.length() - 2);
+    }
+    return Paths.get(absoluteDir);
+  }
+
+  private static class AnnotationCollectorCallback implements SourceRoot.Callback {
+
+    private final AnnotationCollectionVisitor annotationCollectionVisitor;
+
+    public AnnotationCollectorCallback(Map<String, MethodAnnotationsRecord> methodRecords) {
+      this.annotationCollectionVisitor = new AnnotationCollectionVisitor(methodRecords);
+    }
+
+    @Override
+    public Result process(Path localPath, Path absolutePath, ParseResult<CompilationUnit> result) {
+      Result res = Result.SAVE;
+      Optional<CompilationUnit> opt = result.getResult();
+      if (opt.isPresent()) {
+        CompilationUnit cu = opt.get();
+        cu.accept(annotationCollectionVisitor, null);
+      }
+      return res;
+    }
+  }
+
+  private static class AnnotationCollectionVisitor extends VoidVisitorAdapter<Void> {
+
+    private String parentName = "";
+    private boolean isJspecifyNullableImportPresent = false;
+    private boolean isNullMarked = false;
+    private Map<String, MethodAnnotationsRecord> methodRecords;
+    private static final String ARRAY_RETURN_TYPE_STRING = "Array";
+    private static final String NULL_MARKED = "NullMarked";
+    private static final String NULLABLE = "Nullable";
+    private static final String JSPECIFY_NULLABLE_IMPORT = "org.jspecify.annotations.Nullable";
+
+    public AnnotationCollectionVisitor(Map<String, MethodAnnotationsRecord> methodRecords) {
+      this.methodRecords = methodRecords;
+    }
+
+    @Override
+    public void visit(PackageDeclaration pd, Void arg) {
+      this.parentName = pd.getNameAsString();
+      super.visit(pd, null);
+    }
+
+    @Override
+    public void visit(ImportDeclaration id, Void arg) {
+      if (id.getName().toString().contains(JSPECIFY_NULLABLE_IMPORT)) {
+        this.isJspecifyNullableImportPresent = true;
+      }
+      super.visit(id, null);
+    }
+
+    @Override
+    public void visit(ClassOrInterfaceDeclaration cid, Void arg) {
+      /*This logic assumes an explicit @NullMarked annotation on the top-level class within a
+      source file, and it's expected that each source file contains only one top-level class. The
+      logic does not currently handle cases where @NullMarked annotations appear on some nested
+      classes but not others. It also does not consider annotations within package-info.java or
+      module-info.java files.*/
+      parentName += "." + cid.getNameAsString();
+      cid.getAnnotations()
+          .forEach(
+              a -> {
+                if (a.getNameAsString().equalsIgnoreCase(NULL_MARKED)) {
+                  this.isNullMarked = true;
+                }
+              });
+      super.visit(cid, null);
+      // We reset the variable that constructs the parent name after visiting all the children.
+      parentName = parentName.substring(0, parentName.lastIndexOf("." + cid.getNameAsString()));
+    }
+
+    @Override
+    public void visit(MethodDeclaration md, Void arg) {
+      if (this.isNullMarked && hasNullableReturn(md)) {
+        methodRecords.put(
+            parentName + ":" + getMethodReturnTypeString(md) + " " + md.getSignature().toString(),
+            MethodAnnotationsRecord.create(ImmutableSet.of("Nullable"), ImmutableMap.of()));
+      }
+      super.visit(md, null);
+    }
+
+    /**
+     * Determines if a MethodDeclaration can return null.
+     *
+     * @param md The MethodDeclaration instance.
+     * @return {@code true} if the method can return null, {@code false} otherwise.
+     */
+    private boolean hasNullableReturn(MethodDeclaration md) {
+      if (md.getType() instanceof ArrayType) {
+        /* For an Array return type the annotation is on the type when the Array instance is
+        Nullable(Object @Nullable []) and on the node when the elements inside are
+        Nullable(@Nullable Object []) */
+        for (AnnotationExpr annotation : md.getType().getAnnotations()) {
+          if (isAnnotationNullable(annotation)) {
+            return true;
+          }
+        }
+      } else {
+        for (AnnotationExpr annotation : md.getAnnotations()) {
+          if (isAnnotationNullable(annotation)) {
+            return true;
+          }
+        }
+      }
+      return false;
+    }
+
+    /**
+     * Takes a MethodDeclaration and returns the String value for the return type that will be
+     * written into the astubx file.
+     *
+     * @param md The MethodDeclaration instance.
+     * @return The return type string value to be written into the astubx file.
+     */
+    private String getMethodReturnTypeString(MethodDeclaration md) {
+      if (md.getType() instanceof ClassOrInterfaceType) {
+        return md.getType().getChildNodes().get(0).toString();
+      } else if (md.getType() instanceof ArrayType) {
+        return ARRAY_RETURN_TYPE_STRING;
+      } else {
+        return md.getType().toString();
+      }
+    }
+
+    private boolean isAnnotationNullable(AnnotationExpr annotation) {
+      // We only consider jspecify Nullable annotations(star imports are not supported).
+      return (annotation.getNameAsString().equalsIgnoreCase(NULLABLE)
+              && this.isJspecifyNullableImportPresent)
+          || annotation.getNameAsString().equalsIgnoreCase(JSPECIFY_NULLABLE_IMPORT);
+    }
+  }
+}
diff --git a/library-model/library-model-generator/src/main/java/com/uber/nullaway/libmodel/MethodAnnotationsRecord.java b/library-model/library-model-generator/src/main/java/com/uber/nullaway/libmodel/MethodAnnotationsRecord.java
new file mode 100644
index 0000000..7651b8b
--- /dev/null
+++ b/library-model/library-model-generator/src/main/java/com/uber/nullaway/libmodel/MethodAnnotationsRecord.java
@@ -0,0 +1,20 @@
+package com.uber.nullaway.libmodel;
+
+import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+
+/** A record describing the annotations associated with a java method and its arguments. */
+@AutoValue
+public abstract class MethodAnnotationsRecord {
+
+  public static MethodAnnotationsRecord create(
+      ImmutableSet<String> methodAnnotations,
+      ImmutableMap<Integer, ImmutableSet<String>> argumentAnnotations) {
+    return new AutoValue_MethodAnnotationsRecord(methodAnnotations, argumentAnnotations);
+  }
+
+  abstract ImmutableSet<String> methodAnnotations();
+
+  abstract ImmutableMap<Integer, ImmutableSet<String>> argumentAnnotations();
+}
diff --git a/jar-infer/jar-infer-lib/src/main/java/com/uber/nullaway/jarinfer/StubxWriter.java b/library-model/library-model-generator/src/main/java/com/uber/nullaway/libmodel/StubxWriter.java
index 326097f..9dc6052 100644
--- a/jar-infer/jar-infer-lib/src/main/java/com/uber/nullaway/jarinfer/StubxWriter.java
+++ b/library-model/library-model-generator/src/main/java/com/uber/nullaway/libmodel/StubxWriter.java
@@ -1,4 +1,4 @@
-package com.uber.nullaway.jarinfer;
+package com.uber.nullaway.libmodel;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
@@ -12,7 +12,7 @@ import java.util.Map;
 import java.util.Set;
 
 /** Simple writer for the astubx format. */
-final class StubxWriter {
+public final class StubxWriter {
   /**
    * The file magic number for version 0 .astubx files. It should be the first four bytes of any
    * compatible .astubx file.
@@ -31,7 +31,7 @@ final class StubxWriter {
    *     MethodAnnotationsRecord}
    * @exception IOException On output error.
    */
-  static void write(
+  public static void write(
       DataOutputStream out,
       Map<String, String> importedAnnotations,
       Map<String, Set<String>> packageAnnotations,
@@ -93,13 +93,13 @@ final class StubxWriter {
     int methodAnnotationSize = 0;
     int methodArgumentRecordsSize = 0;
     for (Map.Entry<String, MethodAnnotationsRecord> entry : methodRecords.entrySet()) {
-      methodAnnotationSize += entry.getValue().getMethodAnnotations().size();
-      methodArgumentRecordsSize += entry.getValue().getArgumentAnnotations().size();
+      methodAnnotationSize += entry.getValue().methodAnnotations().size();
+      methodArgumentRecordsSize += entry.getValue().argumentAnnotations().size();
     }
     out.writeInt(methodAnnotationSize);
     // Followed by those records as pairs of ints pointing into the dictionary
     for (Map.Entry<String, MethodAnnotationsRecord> entry : methodRecords.entrySet()) {
-      for (String annot : entry.getValue().getMethodAnnotations()) {
+      for (String annot : entry.getValue().methodAnnotations()) {
         out.writeInt(encodingDictionary.get(entry.getKey()));
         out.writeInt(encodingDictionary.get(importedAnnotations.get(annot)));
       }
@@ -110,7 +110,7 @@ final class StubxWriter {
     //  argument position)
     for (Map.Entry<String, MethodAnnotationsRecord> entry : methodRecords.entrySet()) {
       for (Map.Entry<Integer, ImmutableSet<String>> argEntry :
-          entry.getValue().getArgumentAnnotations().entrySet()) {
+          entry.getValue().argumentAnnotations().entrySet()) {
         for (String annot : argEntry.getValue()) {
           out.writeInt(encodingDictionary.get(entry.getKey()));
           out.writeInt(argEntry.getKey());
diff --git a/library-model/test-library-model-generator/build.gradle b/library-model/test-library-model-generator/build.gradle
new file mode 100644
index 0000000..676abca
--- /dev/null
+++ b/library-model/test-library-model-generator/build.gradle
@@ -0,0 +1,55 @@
+/*
+ * Copyright (C) 2024. Uber Technologies
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+plugins {
+    id "java-library"
+}
+
+def testInputsPath = "${rootProject.projectDir}/library-model/test-library-model-generator/src/main/resources/sample_annotated/src"
+def astubxPath = "com/uber/nullaway/libmodel/provider/libmodels.astubx"
+
+jar {
+    manifest {
+        attributes(
+                'Created-By'     : "Gradle ${gradle.gradleVersion}",
+                'Build-Jdk'      : "${System.properties['java.version']} (${System.properties['java.vendor']} ${System.properties['java.vm.version']})",
+                'Build-OS'       : "${System.properties['os.name']} ${System.properties['os.arch']} ${System.properties['os.version']}"
+                )
+    }
+}
+
+jar.doLast {
+    javaexec {
+        classpath = files("${rootProject.projectDir}/library-model/library-model-generator-cli/build/libs/library-model-generator-cli.jar")
+        args = [
+            testInputsPath,
+            "${jar.destinationDirectory.get()}/${astubxPath}"
+        ]
+    }
+    exec {
+        workingDir "./build/libs"
+        commandLine "jar", "uf", "test-library-model-generator.jar", astubxPath
+    }
+}
+
+dependencies {
+    compileOnly deps.apt.autoService
+    annotationProcessor deps.apt.autoService
+    compileOnly project(":nullaway")
+    implementation deps.build.jsr305Annotations
+}
+
+jar.dependsOn ":library-model:library-model-generator-cli:assemble"
diff --git a/library-model/test-library-model-generator/src/main/java/com/uber/nullaway/libmodel/AnnotationExample.java b/library-model/test-library-model-generator/src/main/java/com/uber/nullaway/libmodel/AnnotationExample.java
new file mode 100644
index 0000000..ce8405f
--- /dev/null
+++ b/library-model/test-library-model-generator/src/main/java/com/uber/nullaway/libmodel/AnnotationExample.java
@@ -0,0 +1,42 @@
+package com.uber.nullaway.libmodel;
+
+import java.util.Locale;
+
+/**
+ * This class has the same name as the class under
+ * resources/sample_annotated/src/com/uber/nullaway/libmodel/AnnotationExample.java because we use
+ * this as the unannotated version for our test cases to see if we are appropriately processing the
+ * annotations as an external library model.
+ */
+public class AnnotationExample {
+  public String makeUpperCase(String inputString) {
+    if (inputString == null || inputString.isEmpty()) {
+      return null;
+    } else {
+      return inputString.toUpperCase(Locale.ROOT);
+    }
+  }
+
+  public Integer[] generateIntArray(int size) {
+    if (size <= 0) {
+      return null;
+    } else {
+      Integer[] result = new Integer[size];
+      for (int i = 0; i < size; i++) {
+        result[i] = i + 1;
+      }
+      return result;
+    }
+  }
+
+  public String nullReturn() {
+    return null;
+  }
+
+  public static class InnerExample {
+
+    public String returnNull() {
+      return null;
+    }
+  }
+}
diff --git a/library-model/test-library-model-generator/src/main/java/com/uber/nullaway/libmodel/provider/TestProvider.java b/library-model/test-library-model-generator/src/main/java/com/uber/nullaway/libmodel/provider/TestProvider.java
new file mode 100644
index 0000000..9b4c5e7
--- /dev/null
+++ b/library-model/test-library-model-generator/src/main/java/com/uber/nullaway/libmodel/provider/TestProvider.java
@@ -0,0 +1,14 @@
+package com.uber.nullaway.libmodel.provider;
+
+import com.google.auto.service.AutoService;
+import com.uber.nullaway.jarinfer.JarInferStubxProvider;
+import java.util.Collections;
+import java.util.List;
+
+@AutoService(JarInferStubxProvider.class)
+public class TestProvider implements JarInferStubxProvider {
+  @Override
+  public List<String> pathsToStubxFiles() {
+    return Collections.singletonList("libmodels.astubx");
+  }
+}
diff --git a/library-model/test-library-model-generator/src/main/resources/sample_annotated/src/com/uber/nullaway/libmodel/AnnotationExample.java b/library-model/test-library-model-generator/src/main/resources/sample_annotated/src/com/uber/nullaway/libmodel/AnnotationExample.java
new file mode 100644
index 0000000..f95c540
--- /dev/null
+++ b/library-model/test-library-model-generator/src/main/resources/sample_annotated/src/com/uber/nullaway/libmodel/AnnotationExample.java
@@ -0,0 +1,48 @@
+package com.uber.nullaway.libmodel;
+
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
+@NullMarked
+public class AnnotationExample {
+
+    @Nullable
+    public String makeUpperCase(String inputString) {
+        if (inputString == null || inputString.isEmpty()) {
+            return null;
+        } else {
+            return inputString.toUpperCase();
+        }
+    }
+
+    public Integer @Nullable [] generateIntArray(int size) {
+        if (size <= 0) {
+            return null;
+        } else {
+            Integer[] result = new Integer[size];
+            for (int i = 0; i < size; i++) {
+                result[i] = i + 1;
+            }
+            return result;
+        }
+    }
+
+    /**
+     * This method exists to test that
+     * we do not process this annotation.
+     * Since for the purposes of this tool,
+     * we are only considering the jspecify annotation.
+     */
+    @javax.annotation.Nullable
+    public String nullReturn() {
+        return null;
+    }
+
+    public static class InnerExample {
+
+        @Nullable
+        public String returnNull() {
+            return null;
+        }
+    }
+}
diff --git a/nullaway/src/main/java/com/uber/nullaway/handlers/InferredJARModelsHandler.java b/nullaway/src/main/java/com/uber/nullaway/handlers/InferredJARModelsHandler.java
index bd70059..495ac46 100644
--- a/nullaway/src/main/java/com/uber/nullaway/handlers/InferredJARModelsHandler.java
+++ b/nullaway/src/main/java/com/uber/nullaway/handlers/InferredJARModelsHandler.java
@@ -218,7 +218,8 @@ public class InferredJARModelsHandler extends BaseNoOpHandler {
         if (methodArgAnnotations != null) {
           Set<String> methodAnnotations = methodArgAnnotations.get(RETURN);
           if (methodAnnotations != null) {
-            if (methodAnnotations.contains("javax.annotation.Nullable")) {
+            if (methodAnnotations.contains("javax.annotation.Nullable")
+                || methodAnnotations.contains("org.jspecify.annotations.Nullable")) {
               LOG(DEBUG, "DEBUG", "Nullable return for method: " + methodSign);
               return true;
             }
diff --git a/settings.gradle b/settings.gradle
index 575e1cf..95c4141 100644
--- a/settings.gradle
+++ b/settings.gradle
@@ -32,3 +32,7 @@ include ':jdk-recent-unit-tests'
 include ':code-coverage-report'
 include ':sample-app'
 include ':jar-infer:test-android-lib-jarinfer'
+include ':library-model:library-model-generator'
+include ':library-model:library-model-generator-integration-test'
+include ':library-model:library-model-generator-cli'
+include ':library-model:test-library-model-generator'