Improve how redundantType rule handles Set literals (#2410)

Co-authored-by: calda <1811727+calda@users.noreply.github.com>

Sources/Rules/RedundantType.swift

@@ -129,6 +129,16 @@ public extension FormatRule {
                 let (matches, i, j, wasValue) = formatter.compare(typeStartingAfter: equalsIndex, withTypeStartingAfter: colonIndex, typeEndIndex: typeEndIndex)
                 if matches {
                     removeType(after: equalsIndex, i: i, j: j, wasValue: wasValue)
+                } else if isInferred,
+                          let tokenAfterEquals = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: equalsIndex),
+                          formatter.tokens[tokenAfterEquals] == .startOfScope("["),
+                          let (baseTypeIndex, openAngle, argTypeIndex) = formatter.singleGenericArgType(afterColon: colonIndex, typeEndIndex: typeEndIndex),
+                          formatter.tokens[baseTypeIndex] == .identifier("Set"),
+                          let elementType = formatter.inferredArrayLiteralElementType(at: tokenAfterEquals),
+                          formatter.tokens[argTypeIndex] == elementType
+                {
+                    // The generic argument is redundant (inferred from the array literal)
+                    formatter.removeTokens(in: openAngle ... typeEndIndex)
                 }
             }
         }
@@ -230,6 +240,61 @@ extension Formatter {
         return (true, i, j, wasValue)
     }
 
+    /// For a type annotation of the form `TypeName<SingleArg>`, returns the indices of
+    /// the base type, the opening `<`, and the generic argument token.
+    /// Returns nil if the type has multiple generic arguments, a complex argument type,
+    /// or no generic argument at all.
+    func singleGenericArgType(afterColon colonIndex: Int, typeEndIndex: Int)
+        -> (baseTypeIndex: Int, openAngle: Int, argTypeIndex: Int)?
+    {
+        guard let baseTypeIndex = index(of: .nonSpaceOrCommentOrLinebreak, after: colonIndex),
+              case .identifier = tokens[baseTypeIndex],
+              let openAngle = index(of: .nonSpaceOrCommentOrLinebreak, after: baseTypeIndex),
+              tokens[openAngle] == .startOfScope("<"),
+              let argTypeIndex = index(of: .nonSpaceOrCommentOrLinebreak, after: openAngle),
+              case .identifier = tokens[argTypeIndex],
+              let closeAngle = index(of: .nonSpaceOrCommentOrLinebreak, after: argTypeIndex),
+              closeAngle == typeEndIndex,
+              tokens[closeAngle] == .endOfScope(">")
+        else { return nil }
+        return (baseTypeIndex, openAngle, argTypeIndex)
+    }
+
+    /// Returns the inferred element type for a homogeneous array literal, or nil if the
+    /// array is empty, contains non-literal elements, or has mixed element types.
+    func inferredArrayLiteralElementType(at index: Int) -> Token? {
+        guard tokens[index] == .startOfScope("["),
+              let endIndex = endOfScope(at: index)
+        else { return nil }
+
+        var elementType: Token? = nil
+        var i = index
+
+        while let nextIndex = self.index(of: .nonSpaceOrCommentOrLinebreak, after: i),
+              nextIndex < endIndex
+        {
+            let token = tokens[nextIndex]
+
+            if token == .delimiter(",") {
+                i = nextIndex
+                continue
+            }
+
+            let inferred = typeToken(forValueToken: token)
+            // typeToken returns the token unchanged for non-literals; skip those
+            guard inferred != token else { return nil }
+
+            if let existing = elementType {
+                if existing != inferred { return nil }
+            } else {
+                elementType = inferred
+            }
+            i = token.isStringDelimiter ? (endOfScope(at: nextIndex) ?? nextIndex) : nextIndex
+        }
+
+        return elementType
+    }
+
     /// Returns the equivalent type token for a given value token
     func typeToken(forValueToken token: Token) -> Token {
         switch token {

Tests/Rules/RedundantTypeTests.swift

@@ -826,4 +826,82 @@ final class RedundantTypeTests: XCTestCase {
         let options = FormatOptions(propertyTypes: .inferLocalsOnly)
         testFormatting(for: input, rule: .redundantType, options: options, exclude: [.simplifyGenericConstraints])
     }
+
+    func testRedundantGenericArgRemovedForSetLiteral() {
+        let input = """
+        let set: Set<String> = ["a", "b", "c"]
+        """
+        let output = """
+        let set: Set = ["a", "b", "c"]
+        """
+        let options = FormatOptions(propertyTypes: .inferred)
+        testFormatting(for: input, output, rule: .redundantType, options: options)
+    }
+
+    func testNoRedundantGenericArgRemovedForSetLiteralExplicit() {
+        let input = """
+        let set: Set<String> = ["a", "b", "c"]
+        """
+        let options = FormatOptions(propertyTypes: .explicit)
+        testFormatting(for: input, rule: .redundantType, options: options)
+    }
+
+    func testNoRedundantGenericArgRemovedForArrayTypeLiteral() {
+        let input = """
+        let array: Array<String> = ["a", "b", "c"]
+        """
+        let options = FormatOptions(propertyTypes: .inferred)
+        testFormatting(for: input, rule: .redundantType, options: options, exclude: [.typeSugar])
+    }
+
+    func testNoRedundantGenericArgRemovedForArrayTypeLiteralExplicit() {
+        let input = """
+        let array: Array<String> = ["a", "b", "c"]
+        """
+        let options = FormatOptions(propertyTypes: .explicit)
+        testFormatting(for: input, rule: .redundantType, options: options, exclude: [.typeSugar])
+    }
+
+    func testNoRedundantGenericArgRemovedForCustomArrayLiteralType() {
+        let input = """
+        let custom: MyCustomArrayLiteralType<String> = ["a", "b", "c"]
+        """
+        let options = FormatOptions(propertyTypes: .inferred)
+        testFormatting(for: input, rule: .redundantType, options: options)
+    }
+
+    func testRedundantGenericArgRemovedForSetIntLiteral() {
+        let input = """
+        let set: Set<Int> = [1, 2, 3]
+        """
+        let output = """
+        let set: Set = [1, 2, 3]
+        """
+        let options = FormatOptions(propertyTypes: .inferred)
+        testFormatting(for: input, output, rule: .redundantType, options: options)
+    }
+
+    func testNoRedundantGenericArgRemovedForMismatchedType() {
+        let input = """
+        let set: Set<Double> = [1, 2, 3]
+        """
+        let options = FormatOptions(propertyTypes: .inferred)
+        testFormatting(for: input, rule: .redundantType, options: options)
+    }
+
+    func testNoRedundantGenericArgRemovedForDictionaryLiteral() {
+        let input = """
+        let dict: MyType<String> = ["key": "value"]
+        """
+        let options = FormatOptions(propertyTypes: .inferred)
+        testFormatting(for: input, rule: .redundantType, options: options)
+    }
+
+    func testNoRedundantGenericArgRemovedForMultipleGenericArgs() {
+        let input = """
+        let pair: MyPair<String, Int> = ["a", 1]
+        """
+        let options = FormatOptions(propertyTypes: .inferred)
+        testFormatting(for: input, rule: .redundantType, options: options)
+    }
 }