Go: Mutex protect the two global maps

Author: jketema

go/extractor/extractor.go

@@ -38,6 +38,9 @@ type typeParamParentEntry struct {
 	isFromReceiver bool
 }
 
+// typeParamMutex protects typeParamParent and typeParamOrigin.
+var typeParamMutex sync.RWMutex
+
 var typeParamParent map[*types.TypeParam]typeParamParentEntry = make(map[*types.TypeParam]typeParamParentEntry)
 
 var typeParamOrigin map[*types.TypeParam]*types.TypeParam = make(map[*types.TypeParam]*types.TypeParam)
@@ -2043,7 +2046,10 @@ func getObjectBeingUsed(tw *trap.Writer, ident *ast.Ident) types.Object {
 }
 
 func getTypeParamParentLabel(tw *trap.Writer, tp *types.TypeParam) (trap.Label, bool) {
+	typeParamMutex.RLock()
 	entry, exists := typeParamParent[tp]
+	typeParamMutex.RUnlock()
+
 	if !exists {
 		log.Fatalf("Parent of type parameter does not exist: %s %s", tp.String(), tp.Constraint().String())
 	}
@@ -2055,6 +2061,13 @@ func getTypeParamParentLabel(tw *trap.Writer, tp *types.TypeParam) (trap.Label,
 }
 
 func setTypeParamParent(tp *types.TypeParam, parent types.Object, isFromReceiver bool) {
+	typeParamMutex.Lock()
+	defer typeParamMutex.Unlock()
+	setTypeParamParentLocked(tp, parent, isFromReceiver)
+}
+
+// setTypeParamParentLocked requires typeParamMutex to be held for writing.
+func setTypeParamParentLocked(tp *types.TypeParam, parent types.Object, isFromReceiver bool) {
 	entry, exists := typeParamParent[tp]
 	newEntry := typeParamParentEntry{parent, isFromReceiver}
 	if !exists {
@@ -2106,9 +2119,12 @@ func checkObjectNotSpecialized(obj types.Object) {
 // getTypeParamOrigin returns the origin type parameter for a type parameter
 // from an instantiated method.
 func getTypeParamOrigin(tp *types.TypeParam) *types.TypeParam {
+	typeParamMutex.RLock()
 	if origin, exists := typeParamOrigin[tp]; exists {
+		typeParamMutex.RUnlock()
 		return origin
 	}
+	typeParamMutex.RUnlock()
 	return tp
 }
 
@@ -2144,12 +2160,29 @@ func populateTypeParamParentAndOrigin(tp *types.Named, i int, meth *types.Func)
 	originParams := meth.Type().(*types.Signature).TypeParams()
 
 	if instantiatedParams.Len() != originParams.Len() {
-		log.Fatalf("Method instantiation %s has %d type parameters, origin has %d", tp.Method(i), instantiatedParams.Len(), originParams.Len())
+		log.Fatalf("Method instantiation %s has %d type parameters, origin has %d",
+			tp.Method(i), instantiatedParams.Len(), originParams.Len())
 	}
 
 	for j := 0; j < instantiatedParams.Len(); j++ {
-		instantiatedParam := instantiatedParams.At(j)
-		setTypeParamParent(instantiatedParam, meth, false)
-		typeParamOrigin[instantiatedParam] = originParams.At(j)
+		setTypeParamParentAndOrigin(instantiatedParams.At(j), meth, originParams.At(j))
 	}
 }
+
+// setTypeParamParentAndOrigin atomically records the parent and origin of an
+// instantiated method type parameter.
+func setTypeParamParentAndOrigin(instantiatedParam *types.TypeParam, parent *types.Func, originParam *types.TypeParam) {
+	typeParamMutex.Lock()
+	defer typeParamMutex.Unlock()
+
+	setTypeParamParentLocked(instantiatedParam, parent, false)
+
+	if existing, exists := typeParamOrigin[instantiatedParam]; exists {
+		if existing != originParam {
+			log.Fatalf("Origin of type parameter '%s %s' being set to a different value: '%s' vs '%s'",
+				instantiatedParam.String(), instantiatedParam.Constraint().String(), existing.String(), originParam.String())
+		}
+	}
+
+	typeParamOrigin[instantiatedParam] = originParam
+}