diff --git a/java/hiddenapi_modular.go b/java/hiddenapi_modular.go
index 8e17d56..86466d6 100644
--- a/java/hiddenapi_modular.go
+++ b/java/hiddenapi_modular.go
@@ -417,12 +417,12 @@
 // an entry for every single member in the dex implementation jars of the individual modules. Every
 // signature in any of the other files MUST be included in this file.
 //
-// moduleSpecificFlagsPaths are the paths to the flags files generated by each module using
-// information from the baseFlagsPath as well as from annotations within the source.
+// annotationFlags is the path to the annotation flags file generated from annotation information
+// in each module.
 //
-// augmentationInfo is a struct containing paths to files that augment the information provided by
-// the moduleSpecificFlagsPaths.
-func buildRuleToGenerateHiddenApiFlags(ctx android.BuilderContext, name, desc string, outputPath android.WritablePath, baseFlagsPath android.Path, moduleSpecificFlagsPaths android.Paths, flagFileInfo *hiddenAPIFlagFileInfo) {
+// flagFileInfo is a struct containing paths to files that augment the information provided by
+// the annotationFlags.
+func buildRuleToGenerateHiddenApiFlags(ctx android.BuilderContext, name, desc string, outputPath android.WritablePath, baseFlagsPath android.Path, annotationFlags android.Path, flagFileInfo *hiddenAPIFlagFileInfo) {
 
 	// The file which is used to record that the flags file is valid.
 	var validFile android.WritablePath
@@ -452,7 +452,7 @@
 	command := rule.Command().
 		BuiltTool("generate_hiddenapi_lists").
 		FlagWithInput("--csv ", baseFlagsPath).
-		Inputs(moduleSpecificFlagsPaths).
+		Input(annotationFlags).
 		FlagWithOutput("--output ", tempPath)
 
 	// Add the options for the different categories of flag files.
@@ -490,31 +490,18 @@
 // * metadata.csv
 // * index.csv
 // * all-flags.csv
-func hiddenAPIGenerateAllFlagsForBootclasspathFragment(ctx android.ModuleContext, contents []android.Module, stubJarsByKind map[android.SdkKind]android.Paths, flagFileInfo *hiddenAPIFlagFileInfo) {
-
+func hiddenAPIGenerateAllFlagsForBootclasspathFragment(ctx android.ModuleContext, contents []hiddenAPIModule, stubJarsByKind map[android.SdkKind]android.Paths, flagFileInfo *hiddenAPIFlagFileInfo) {
 	hiddenApiSubDir := "modular-hiddenapi"
 
-	bootDexJars := android.Paths{}
-	classesJars := android.Paths{}
-	for _, module := range contents {
-		if hiddenAPI, ok := module.(hiddenAPIIntf); ok {
-			classesJars = append(classesJars, hiddenAPI.classesJars()...)
-			bootDexJar := hiddenAPI.bootDexJar()
-			if bootDexJar == nil {
-				ctx.ModuleErrorf("module %s does not provide a dex jar", module)
-			} else {
-				bootDexJars = append(bootDexJars, bootDexJar)
-			}
-		} else {
-			ctx.ModuleErrorf("module %s does not implement hiddenAPIIntf", module)
-		}
-	}
-
 	// Generate the stub-flags.csv.
+	bootDexJars := extractBootDexJarsFromHiddenAPIModules(ctx, contents)
 	stubFlagsCSV := android.PathForModuleOut(ctx, hiddenApiSubDir, "stub-flags.csv")
 	rule := ruleToGenerateHiddenAPIStubFlagsFile(ctx, stubFlagsCSV, bootDexJars, stubJarsByKind)
 	rule.Build("modularHiddenAPIStubFlagsFile", "modular hiddenapi stub flags")
 
+	// Extract the classes jars from the contents.
+	classesJars := extractClassJarsFromHiddenAPIModules(ctx, contents)
+
 	// Generate the set of flags from the annotations in the source code.
 	annotationFlagsCSV := android.PathForModuleOut(ctx, hiddenApiSubDir, "annotation-flags.csv")
 	buildRuleToGenerateAnnotationFlags(ctx, "modular hiddenapi annotation flags", classesJars, stubFlagsCSV, annotationFlagsCSV)
@@ -523,7 +510,7 @@
 	metadataCSV := android.PathForModuleOut(ctx, hiddenApiSubDir, "metadata.csv")
 	buildRuleToGenerateMetadata(ctx, "modular hiddenapi metadata", classesJars, stubFlagsCSV, metadataCSV)
 
-	// Generate the index file from the annotations in the source code.
+	// Generate the index file from the CSV files in the classes jars.
 	indexCSV := android.PathForModuleOut(ctx, hiddenApiSubDir, "index.csv")
 	buildRuleToGenerateIndex(ctx, "modular hiddenapi index", classesJars, indexCSV)
 
@@ -537,7 +524,7 @@
 	// Generate the all-flags.csv which are the flags that will, in future, be encoded into the dex
 	// files.
 	outputPath := android.PathForModuleOut(ctx, hiddenApiSubDir, "all-flags.csv")
-	buildRuleToGenerateHiddenApiFlags(ctx, "modularHiddenApiAllFlags", "modular hiddenapi all flags", outputPath, stubFlagsCSV, android.Paths{annotationFlagsCSV}, flagFileInfo)
+	buildRuleToGenerateHiddenApiFlags(ctx, "modularHiddenApiAllFlags", "modular hiddenapi all flags", outputPath, stubFlagsCSV, annotationFlagsCSV, flagFileInfo)
 
 	// Store the paths in the info for use by other modules and sdk snapshot generation.
 	flagFileInfo.StubFlagsPaths = android.Paths{stubFlagsCSV}
@@ -546,3 +533,41 @@
 	flagFileInfo.IndexPaths = android.Paths{indexCSV}
 	flagFileInfo.AllFlagsPaths = android.Paths{outputPath}
 }
+
+// gatherHiddenAPIModuleFromContents gathers the hiddenAPIModule from the supplied contents.
+func gatherHiddenAPIModuleFromContents(ctx android.ModuleContext, contents []android.Module) []hiddenAPIModule {
+	hiddenAPIModules := []hiddenAPIModule{}
+	for _, module := range contents {
+		if hiddenAPI, ok := module.(hiddenAPIModule); ok {
+			hiddenAPIModules = append(hiddenAPIModules, hiddenAPI)
+		} else if _, ok := module.(*DexImport); ok {
+			// Ignore this for the purposes of hidden API processing
+		} else {
+			ctx.ModuleErrorf("module %s does not implement hiddenAPIModule", module)
+		}
+	}
+	return hiddenAPIModules
+}
+
+// extractBootDexJarsFromHiddenAPIModules extracts the boot dex jars from the supplied modules.
+func extractBootDexJarsFromHiddenAPIModules(ctx android.ModuleContext, contents []hiddenAPIModule) android.Paths {
+	bootDexJars := android.Paths{}
+	for _, module := range contents {
+		bootDexJar := module.bootDexJar()
+		if bootDexJar == nil {
+			ctx.ModuleErrorf("module %s does not provide a dex jar", module)
+		} else {
+			bootDexJars = append(bootDexJars, bootDexJar)
+		}
+	}
+	return bootDexJars
+}
+
+// extractClassJarsFromHiddenAPIModules extracts the class jars from the supplied modules.
+func extractClassJarsFromHiddenAPIModules(ctx android.ModuleContext, contents []hiddenAPIModule) android.Paths {
+	classesJars := android.Paths{}
+	for _, module := range contents {
+		classesJars = append(classesJars, module.classesJars()...)
+	}
+	return classesJars
+}