diff --git a/java/dexpreopt_bootjars.go b/java/dexpreopt_bootjars.go
index c6aa7fe0d..87f6d5e33 100644
--- a/java/dexpreopt_bootjars.go
+++ b/java/dexpreopt_bootjars.go
@@ -88,6 +88,9 @@ type bootImageConfig struct {
 	images     map[android.ArchType]android.OutputPath  // first image file
 	imagesDeps map[android.ArchType]android.OutputPaths // all files
 
+	// Only for extensions, paths to the primary boot images (grouped by target).
+	primaryImages map[android.ArchType]android.OutputPath
+
 	// File path to a zip archive with all image files (or nil, if not needed).
 	zip android.WritablePath
 }
@@ -355,7 +358,7 @@ func buildBootImageRuleForArch(ctx android.SingletonContext, image *bootImage,
 	}
 
 	if image.extension {
-		artImage := artBootImageConfig(ctx).images[arch]
+		artImage := image.primaryImages[arch]
 		cmd.
 			Flag("--runtime-arg").FlagWithInputList("-Xbootclasspath:", image.dexPathsDeps.Paths(), ":").
 			Flag("--runtime-arg").FlagWithList("-Xbootclasspath-locations:", image.dexLocationsDeps, ":").
diff --git a/java/dexpreopt_config.go b/java/dexpreopt_config.go
index 31bec93eb..637a32f0a 100644
--- a/java/dexpreopt_config.go
+++ b/java/dexpreopt_config.go
@@ -246,10 +246,12 @@ func genBootImageConfigs(ctx android.PathContext) map[string]*bootImageConfig {
 
 		// specific to the framework config
 		frameworkCfg.dexPathsDeps = append(artCfg.dexPathsDeps, frameworkCfg.dexPathsDeps...)
+		frameworkCfg.primaryImages = artCfg.images
 		frameworkCfg.imageLocations = append(artCfg.imageLocations, frameworkCfg.imageLocations...)
 
 		// specific to the jitzygote-framework config
 		frameworkJZCfg.dexPathsDeps = append(artJZCfg.dexPathsDeps, frameworkJZCfg.dexPathsDeps...)
+		frameworkJZCfg.primaryImages = artJZCfg.images
 		frameworkJZCfg.imageLocations = append(artJZCfg.imageLocations, frameworkJZCfg.imageLocations...)
 
 		return configs