From a18f8e7187a4e2c4ab6b314b29c8d575df72b497 Mon Sep 17 00:00:00 2001
From: Tristan Swadell <tswadell@google.com>
Date: Thu, 14 Dec 2023 15:06:32 -0800
Subject: [PATCH] Introduce 'wither' method for changing type traits (#871)

---
 checker/types.go           |  6 +++++-
 common/types/types.go      | 17 +++++++++++++++++
 common/types/types_test.go | 11 +++++++++++
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/checker/types.go b/checker/types.go
index e2373d1b..ed51beed 100644
--- a/checker/types.go
+++ b/checker/types.go
@@ -137,7 +137,11 @@ func internalIsAssignable(m *mapping, t1, t2 *types.Type) bool {
 	case types.BoolKind, types.BytesKind, types.DoubleKind, types.IntKind, types.StringKind, types.UintKind,
 		types.AnyKind, types.DurationKind, types.TimestampKind,
 		types.StructKind:
-		return t1.IsAssignableType(t2)
+		// Test whether t2 is assignable from t1. The order of this check won't usually matter;
+		// however, there may be cases where type capabilities are expanded beyond what is supported
+		// in the current common/types package. For example, an interface designation for a group of
+		// Struct types.
+		return t2.IsAssignableType(t1)
 	case types.TypeKind:
 		return kind2 == types.TypeKind
 	case types.OpaqueKind, types.ListKind, types.MapKind:
diff --git a/common/types/types.go b/common/types/types.go
index 76624eef..aef5c693 100644
--- a/common/types/types.go
+++ b/common/types/types.go
@@ -373,6 +373,23 @@ func (t *Type) TypeName() string {
 	return t.runtimeTypeName
 }
 
+// WithTraits creates a copy of the current Type and sets the trait mask to the traits parameter.
+//
+// This method should be used with Opaque types where the type acts like a container, e.g. vector.
+func (t *Type) WithTraits(traits int) *Type {
+	if t == nil {
+		return nil
+	}
+	return &Type{
+		kind:                    t.kind,
+		parameters:              t.parameters,
+		runtimeTypeName:         t.runtimeTypeName,
+		isAssignableType:        t.isAssignableType,
+		isAssignableRuntimeType: t.isAssignableRuntimeType,
+		traitMask:               traits,
+	}
+}
+
 // String returns a human-readable definition of the type name.
 func (t *Type) String() string {
 	if len(t.Parameters()) == 0 {
diff --git a/common/types/types_test.go b/common/types/types_test.go
index 68e3aefe..4a0e3709 100644
--- a/common/types/types_test.go
+++ b/common/types/types_test.go
@@ -787,6 +787,17 @@ func TestTypeHasTrait(t *testing.T) {
 	}
 }
 
+func TestTypeWithTraits(t *testing.T) {
+	vec := NewOpaqueType("vector", NewTypeParamType("T"))
+	if vec.HasTrait(traits.SizerType) {
+		t.Error("vec.HasTrait(SizerType) returned true")
+	}
+	vec = vec.WithTraits(traits.SizerType | traits.ContainerType | traits.IndexerType)
+	if !vec.HasTrait(traits.SizerType) {
+		t.Errorf("vec.HasTrait(SizerType) returned false after WithTraits() call")
+	}
+}
+
 func TestTypeConvertToType(t *testing.T) {
 	if BoolType.ConvertToType(TypeType) != TypeType {
 		t.Error("ConvertToType(TypeType) did not produce type value")