diff --git a/java/droidstubs.go b/java/droidstubs.go
index 0157185..a8e0a22 100644
--- a/java/droidstubs.go
+++ b/java/droidstubs.go
@@ -283,66 +283,6 @@
 	return module
 }
 
-func getStubsTypeAndTag(tag string) (StubsType, string, error) {
-	if len(tag) == 0 {
-		return Everything, "", nil
-	}
-	if tag[0] != '.' {
-		return Unavailable, "", fmt.Errorf("tag must begin with \".\"")
-	}
-
-	stubsType := Everything
-	// Check if the tag has a stubs type prefix (e.g. ".exportable")
-	for st := Everything; st <= Exportable; st++ {
-		if strings.HasPrefix(tag, "."+st.String()) {
-			stubsType = st
-		}
-	}
-
-	return stubsType, strings.TrimPrefix(tag, "."+stubsType.String()), nil
-}
-
-// Droidstubs' tag supports specifying with the stubs type.
-// While supporting the pre-existing tags, it also supports tags with
-// the stubs type prefix. Some examples are shown below:
-// {.annotations.zip} - pre-existing behavior. Returns the path to the
-// annotation zip.
-// {.exportable} - Returns the path to the exportable stubs src jar.
-// {.exportable.annotations.zip} - Returns the path to the exportable
-// annotations zip file.
-// {.runtime.api_versions.xml} - Runtime stubs does not generate api versions
-// xml file. For unsupported combinations, the default everything output file
-// is returned.
-func (d *Droidstubs) OutputFiles(tag string) (android.Paths, error) {
-	stubsType, prefixRemovedTag, err := getStubsTypeAndTag(tag)
-	if err != nil {
-		return nil, err
-	}
-	switch prefixRemovedTag {
-	case "":
-		stubsSrcJar, err := d.StubsSrcJar(stubsType)
-		return android.Paths{stubsSrcJar}, err
-	case ".docs.zip":
-		docZip, err := d.DocZip(stubsType)
-		return android.Paths{docZip}, err
-	case ".api.txt", android.DefaultDistTag:
-		// This is the default dist path for dist properties that have no tag property.
-		apiFilePath, err := d.ApiFilePath(stubsType)
-		return android.Paths{apiFilePath}, err
-	case ".removed-api.txt":
-		removedApiFilePath, err := d.RemovedApiFilePath(stubsType)
-		return android.Paths{removedApiFilePath}, err
-	case ".annotations.zip":
-		annotationsZip, err := d.AnnotationsZip(stubsType)
-		return android.Paths{annotationsZip}, err
-	case ".api_versions.xml":
-		apiVersionsXmlFilePath, err := d.ApiVersionsXmlFilePath(stubsType)
-		return android.Paths{apiVersionsXmlFilePath}, err
-	default:
-		return nil, fmt.Errorf("unsupported module reference tag %q", tag)
-	}
-}
-
 func (d *Droidstubs) AnnotationsZip(stubsType StubsType) (ret android.Path, err error) {
 	switch stubsType {
 	case Everything:
@@ -1363,6 +1303,46 @@
 
 		rule.Build("nullabilityWarningsCheck", "nullability warnings check")
 	}
+
+	d.setOutputFiles(ctx)
+}
+
+// This method sets the outputFiles property, which is used to set the
+// OutputFilesProvider later.
+// Droidstubs' tag supports specifying with the stubs type.
+// While supporting the pre-existing tags, it also supports tags with
+// the stubs type prefix. Some examples are shown below:
+// {.annotations.zip} - pre-existing behavior. Returns the path to the
+// annotation zip.
+// {.exportable} - Returns the path to the exportable stubs src jar.
+// {.exportable.annotations.zip} - Returns the path to the exportable
+// annotations zip file.
+// {.runtime.api_versions.xml} - Runtime stubs does not generate api versions
+// xml file. For unsupported combinations, the default everything output file
+// is returned.
+func (d *Droidstubs) setOutputFiles(ctx android.ModuleContext) {
+	tagToOutputFileFunc := map[string]func(StubsType) (android.Path, error){
+		"":                     d.StubsSrcJar,
+		".docs.zip":            d.DocZip,
+		".api.txt":             d.ApiFilePath,
+		android.DefaultDistTag: d.ApiFilePath,
+		".removed-api.txt":     d.RemovedApiFilePath,
+		".annotations.zip":     d.AnnotationsZip,
+		".api_versions.xml":    d.ApiVersionsXmlFilePath,
+	}
+	stubsTypeToPrefix := map[StubsType]string{
+		Everything: "",
+		Exportable: ".exportable",
+	}
+	for _, tag := range android.SortedKeys(tagToOutputFileFunc) {
+		for _, stubType := range android.SortedKeys(stubsTypeToPrefix) {
+			tagWithPrefix := stubsTypeToPrefix[stubType] + tag
+			outputFile, err := tagToOutputFileFunc[tag](stubType)
+			if err == nil {
+				ctx.SetOutputFiles(android.Paths{outputFile}, tagWithPrefix)
+			}
+		}
+	}
 }
 
 func (d *Droidstubs) createApiContribution(ctx android.DefaultableHookContext) {