From f2862b476edcef83412c7af8687c9cd8e4097c0f Mon Sep 17 00:00:00 2001
From: Peter Ebden <peter.ebden@gmail.com>
Date: Mon, 24 Jul 2017 01:48:29 +0100
Subject: [PATCH] Memoise calls to fullyQualifiedNameToSwaggerName to speed it
 up for large registries (#421)

Memoise calls to fullyQualifiedNameToSwaggerName for all registries ever seen instead of just the last one.
---
 protoc-gen-swagger/genswagger/template.go     | 19 ++++++++++++++++---
 .../genswagger/template_test.go               |  3 ++-
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/protoc-gen-swagger/genswagger/template.go b/protoc-gen-swagger/genswagger/template.go
index 170e0e4befd..b2353aa5341 100644
--- a/protoc-gen-swagger/genswagger/template.go
+++ b/protoc-gen-swagger/genswagger/template.go
@@ -8,6 +8,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"sync"
 
 	pbdescriptor "github.com/golang/protobuf/protoc-gen-go/descriptor"
 	"github.com/grpc-ecosystem/grpc-gateway/protoc-gen-grpc-gateway/descriptor"
@@ -329,9 +330,21 @@ func renderEnumerationsAsDefinition(enums enumMap, d swaggerDefinitionsObject, r
 
 // Take in a FQMN or FQEN and return a swagger safe version of the FQMN
 func fullyQualifiedNameToSwaggerName(fqn string, reg *descriptor.Registry) string {
-	return resolveFullyQualifiedNameToSwaggerName(fqn, append(reg.GetAllFQMNs(), reg.GetAllFQENs()...))
+	registriesSeenMutex.Lock()
+	defer registriesSeenMutex.Unlock()
+	if mapping, present := registriesSeen[reg]; present {
+		return mapping[fqn]
+	}
+	mapping := resolveFullyQualifiedNameToSwaggerNames(append(reg.GetAllFQMNs(), reg.GetAllFQENs()...))
+	registriesSeen[reg] = mapping
+	return mapping[fqn]
 }
 
+// registriesSeen is used to memoise calls to resolveFullyQualifiedNameToSwaggerNames so
+// we don't repeat it unnecessarily, since it can take some time.
+var registriesSeen = map[*descriptor.Registry]map[string]string{}
+var registriesSeenMutex sync.Mutex
+
 // Take the names of every proto and "uniq-ify" them. The idea is to produce a
 // set of names that meet a couple of conditions. They must be stable, they
 // must be unique, and they must be shorter than the FQN.
@@ -339,7 +352,7 @@ func fullyQualifiedNameToSwaggerName(fqn string, reg *descriptor.Registry) strin
 // This likely could be made better. This will always generate the same names
 // but may not always produce optimal names. This is a reasonably close
 // approximation of what they should look like in most cases.
-func resolveFullyQualifiedNameToSwaggerName(fqn string, messages []string) string {
+func resolveFullyQualifiedNameToSwaggerNames(messages []string) map[string]string {
 	packagesByDepth := make(map[int][][]string)
 	uniqueNames := make(map[string]string)
 
@@ -379,7 +392,7 @@ func resolveFullyQualifiedNameToSwaggerName(fqn string, messages []string) strin
 			}
 		}
 	}
-	return uniqueNames[fqn]
+	return uniqueNames
 }
 
 // Swagger expects paths of the form /path/{string_value} but grpc-gateway paths are expected to be of the form /path/{string_value=strprefix/*}. This should reformat it correctly.
diff --git a/protoc-gen-swagger/genswagger/template_test.go b/protoc-gen-swagger/genswagger/template_test.go
index 1d438936368..4633a0bbac6 100644
--- a/protoc-gen-swagger/genswagger/template_test.go
+++ b/protoc-gen-swagger/genswagger/template_test.go
@@ -758,7 +758,8 @@ func TestResolveFullyQualifiedNameToSwaggerName(t *testing.T) {
 	}
 
 	for _, data := range tests {
-		output := resolveFullyQualifiedNameToSwaggerName(data.input, data.listOfFQMNs)
+		names := resolveFullyQualifiedNameToSwaggerNames(data.listOfFQMNs)
+		output := names[data.input]
 		if output != data.output {
 			t.Errorf("Expected fullyQualifiedNameToSwaggerName(%v) to be %s but got %s",
 				data.input, data.output, output)