From 08f7cf06a2f30a3e036f279bcf955fe0153fc76d Mon Sep 17 00:00:00 2001
From: Bevan Arps <bevan.arps@microsoft.com>
Date: Fri, 23 Feb 2024 09:35:19 +1300
Subject: [PATCH] Reduce the number of types loaded by the generator (#3803)

* Pass configuration into loadAllSchemas()

* Reduce the number of loaded types
---
 .../internal/codegen/pipeline/load_types.go   | 21 +++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/v2/tools/generator/internal/codegen/pipeline/load_types.go b/v2/tools/generator/internal/codegen/pipeline/load_types.go
index 6e8acd1293e..09b33cb3416 100644
--- a/v2/tools/generator/internal/codegen/pipeline/load_types.go
+++ b/v2/tools/generator/internal/codegen/pipeline/load_types.go
@@ -313,10 +313,8 @@ func loadSwaggerData(
 ) (jsonast.SwaggerTypes, error) {
 	schemas, err := loadAllSchemas(
 		ctx,
-		config.SchemaRoot,
-		config.LocalPathPrefix(),
 		idFactory,
-		config.Status.Overrides,
+		config,
 		log)
 	if err != nil {
 		return jsonast.SwaggerTypes{}, err
@@ -684,12 +682,14 @@ func shouldSkipDir(filePath string) bool {
 // shouldSkipDir), and returns those files in a map of path→swagger spec.
 func loadAllSchemas(
 	ctx context.Context,
-	rootPath string,
-	localPathPrefix string,
 	idFactory astmodel.IdentifierFactory,
-	overrides []config.SchemaOverride,
+	config *config.Configuration,
 	log logr.Logger,
 ) (map[string]jsonast.PackageAndSwagger, error) {
+	rootPath := config.SchemaRoot
+	localPathPrefix := config.LocalPathPrefix()
+	overrides := config.Status.Overrides
+
 	var mutex sync.Mutex
 	schemas := make(map[string]jsonast.PackageAndSwagger)
 
@@ -726,6 +726,15 @@ func loadAllSchemas(
 				astmodel.GeneratorVersion,
 				version)
 
+			// We need the file if the version is short (e.g. "v1") because those are often shared between
+			// resource providers.
+			// Alternatively, we need the file if we have configuration for the group
+			fileNeeded := len(version) < 10 ||
+				config.ObjectModelConfiguration.IsGroupConfigured(pkg)
+			if !fileNeeded {
+				return nil
+			}
+
 			// all files are loaded in parallel to speed this up
 			logInfoSparse(
 				log,