From 945ef1528ee9082dff721874ee142e0b15b7235e Mon Sep 17 00:00:00 2001
From: Nick Lockwood <nick@charcoaldesign.co.uk>
Date: Sun, 12 Jul 2020 17:05:36 +0100
Subject: [PATCH] Add inference for noSpaceOperators option

---
 Sources/Inference.swift     | 36 ++++++++++++++++++-----
 Sources/Rules.swift         |  6 ++--
 Tests/InferenceTests.swift  | 58 ++++++++++++++++++++++++++++---------
 Tests/RulesTests.swift      | 20 ++++++++-----
 Tests/XCTestManifests.swift |  8 ++++-
 5 files changed, 96 insertions(+), 32 deletions(-)

diff --git a/Sources/Inference.swift b/Sources/Inference.swift
index f33547901..6ee1592ea 100644
--- a/Sources/Inference.swift
+++ b/Sources/Inference.swift
@@ -139,20 +139,42 @@ private struct Inference {
         options.allowInlineSemicolons = allow
     }
 
-    let spaceAroundRangeOperators = OptionInferrer { formatter, options in
-        var spaced = 0, unspaced = 0
-        formatter.forEach(.rangeOperator) { i, _ in
-            guard let nextToken = formatter.next(.nonSpaceOrCommentOrLinebreak, after: i),
+    let noSpaceOperators = OptionInferrer { formatter, options in
+        var spaced = [String: Int](), unspaced = [String: Int]()
+        formatter.forEach(.operator) { i, token in
+            guard case let .operator(name, .infix) = token, name != ".",
+                let nextToken = formatter.next(.nonSpaceOrCommentOrLinebreak, after: i),
                 nextToken.string != ")", nextToken.string != "," else {
                 return
             }
             if formatter.token(at: i + 1)?.isSpaceOrLinebreak == true {
-                spaced += 1
+                spaced[name, default: 0] += 1
             } else {
-                unspaced += 1
+                unspaced[name, default: 0] += 1
+            }
+        }
+        var noSpaceOperators = Set<String>()
+        let operators = Set(spaced.keys).union(unspaced.keys)
+        for name in operators where unspaced[name, default: 0] > spaced[name, default: 0] + 1 {
+            noSpaceOperators.insert(name)
+        }
+        // Related pairs
+        let relatedPairs = [
+            ("...", "..<"), ("*", "/"), ("*=", "/="), ("+", "-"), ("+=", "-="),
+            ("==", "!="), ("<", ">"), ("<=", ">="), ("<<", ">>"),
+        ]
+        for pair in relatedPairs {
+            if noSpaceOperators.contains(pair.0),
+                !noSpaceOperators.contains(pair.1),
+                !operators.contains(pair.1) {
+                noSpaceOperators.insert(pair.1)
+            } else if noSpaceOperators.contains(pair.1),
+                !noSpaceOperators.contains(pair.0),
+                !operators.contains(pair.0) {
+                noSpaceOperators.insert(pair.0)
             }
         }
-        options.spaceAroundRangeOperators = (spaced >= unspaced)
+        options.noSpaceOperators = noSpaceOperators
     }
 
     let useVoid = OptionInferrer { formatter, options in
diff --git a/Sources/Rules.swift b/Sources/Rules.swift
index 6e708885d..cb5192c41 100644
--- a/Sources/Rules.swift
+++ b/Sources/Rules.swift
@@ -608,8 +608,7 @@ public struct _FormatRules {
     ) { formatter in
         formatter.forEach(.rangeOperator) { i, token in
             guard case let .operator(name, .infix) = token else { return }
-            if !formatter.options.spaceAroundRangeOperators ||
-                formatter.options.noSpaceOperators.contains(name) {
+            if !formatter.options.spaceAroundRangeOperators {
                 if formatter.token(at: i + 1)?.isSpace == true,
                     formatter.token(at: i - 1)?.isSpace == true,
                     let nextToken = formatter.next(.nonSpace, after: i),
@@ -619,7 +618,8 @@ public struct _FormatRules {
                     formatter.removeToken(at: i + 1)
                     formatter.removeToken(at: i - 1)
                 }
-            } else {
+            } else if formatter.options.spaceAroundRangeOperators,
+                !formatter.options.noSpaceOperators.contains(name) {
                 if formatter.token(at: i + 1)?.isSpaceOrLinebreak == false {
                     formatter.insertToken(.space(" "), at: i + 1)
                 }
diff --git a/Tests/InferenceTests.swift b/Tests/InferenceTests.swift
index ded3c50e8..7bc5ad6f7 100644
--- a/Tests/InferenceTests.swift
+++ b/Tests/InferenceTests.swift
@@ -138,20 +138,6 @@ class InferenceTests: XCTestCase {
         XCTAssertFalse(options.allowInlineSemicolons)
     }
 
-    // MARK: spaceAroundRangeOperators
-
-    func testInferSpaceAroundRangeOperators() {
-        let input = "let foo = 0 ..< bar\n;let baz = 1...quux"
-        let options = inferFormatOptions(from: tokenize(input))
-        XCTAssertTrue(options.spaceAroundRangeOperators)
-    }
-
-    func testInferNoSpaceAroundRangeOperators() {
-        let input = "let foo = 0..<bar\n;let baz = 1...quux"
-        let options = inferFormatOptions(from: tokenize(input))
-        XCTAssertFalse(options.spaceAroundRangeOperators)
-    }
-
     // MARK: useVoid
 
     func testInferUseVoid() {
@@ -631,4 +617,48 @@ class InferenceTests: XCTestCase {
         let options = inferFormatOptions(from: tokenize(input))
         XCTAssertFalse(options.indentCase)
     }
+
+    // MARK: nospaceoperators
+
+    func testInferNoSpaceAroundTimesOperator() {
+        let input = "let foo = a*b + c / d + e*f"
+        let options = inferFormatOptions(from: tokenize(input))
+        XCTAssertEqual(options.noSpaceOperators, ["*"])
+    }
+
+    func testInferSpaceAroundTimesOperator() {
+        let input = "let foo = a*b + c * d + e*f"
+        let options = inferFormatOptions(from: tokenize(input))
+        XCTAssertEqual(options.noSpaceOperators, [])
+    }
+
+    func testInferNoSpaceAroundTimesAndDivideOperator() {
+        let input = "let foo = a*b + c*d"
+        let options = inferFormatOptions(from: tokenize(input))
+        XCTAssertEqual(options.noSpaceOperators, ["/", "*"])
+    }
+
+    func testInferSpaceAroundRangeOperators() {
+        let input = "let foo = a ... b"
+        let options = inferFormatOptions(from: tokenize(input))
+        XCTAssertEqual(options.noSpaceOperators, [])
+    }
+
+    func testInferSpaceAroundRangeOperators2() {
+        let input = "let foo = a ..< b"
+        let options = inferFormatOptions(from: tokenize(input))
+        XCTAssertEqual(options.noSpaceOperators, [])
+    }
+
+    func testInferSpaceAroundRangeOperators3() {
+        let input = "let foo = a...b; let bar = a...b; let baz = a ..< b"
+        let options = inferFormatOptions(from: tokenize(input))
+        XCTAssertEqual(options.noSpaceOperators, ["..."])
+    }
+
+    func testInferNoSpaceAroundRangeOperators3() {
+        let input = "let foo = a...b; let bar = a...b"
+        let options = inferFormatOptions(from: tokenize(input))
+        XCTAssertEqual(options.noSpaceOperators, ["...", "..<"])
+    }
 }
diff --git a/Tests/RulesTests.swift b/Tests/RulesTests.swift
index 9b8f81ce5..4749aa94f 100644
--- a/Tests/RulesTests.swift
+++ b/Tests/RulesTests.swift
@@ -3987,15 +3987,14 @@ class RulesTests: XCTestCase {
         testFormatting(for: input, output, rule: FormatRules.ranges)
     }
 
-    // spaceAroundRangeOperators = true
-
-    func testNoSpaceAroundRangeOperatorsWithCustomOptions() {
-        let input = "foo ..< bar"
-        let output = "foo..<bar"
-        let options = FormatOptions(spaceAroundRangeOperators: false)
-        testFormatting(for: input, output, rule: FormatRules.ranges, options: options)
+    func testNoSpaceAroundRangeOperatorsWithExplicitNoSpace() {
+        let input = "foo..<bar"
+        let options = FormatOptions(noSpaceOperators: ["..<"])
+        testFormatting(for: input, rule: FormatRules.ranges, options: options)
     }
 
+    // spaceAroundRangeOperators = true
+
     func testNoSpaceAddedAroundVariadic() {
         let input = "foo(bar: Int...)"
         let options = FormatOptions(spaceAroundRangeOperators: true)
@@ -4041,6 +4040,13 @@ class RulesTests: XCTestCase {
 
     // spaceAroundRangeOperators = false
 
+    func testNoSpaceAroundRangeOperatorsWithCustomOptions() {
+        let input = "foo ..< bar"
+        let output = "foo..<bar"
+        let options = FormatOptions(spaceAroundRangeOperators: false)
+        testFormatting(for: input, output, rule: FormatRules.ranges, options: options)
+    }
+
     func testSpaceNotRemovedBeforeLeadingRangeOperatorWithSpaceAroundRangeOperatorsFalse() {
         let input = "let range = ..<foo.endIndex"
         let options = FormatOptions(spaceAroundRangeOperators: false)
diff --git a/Tests/XCTestManifests.swift b/Tests/XCTestManifests.swift
index f45fac732..a6ac84e0c 100644
--- a/Tests/XCTestManifests.swift
+++ b/Tests/XCTestManifests.swift
@@ -228,7 +228,9 @@ extension InferenceTests {
         ("testInferNoIndentCase", testInferNoIndentCase),
         ("testInferNoOctalGrouping", testInferNoOctalGrouping),
         ("testInferNoSpaceAfterOperatorFunc", testInferNoSpaceAfterOperatorFunc),
-        ("testInferNoSpaceAroundRangeOperators", testInferNoSpaceAroundRangeOperators),
+        ("testInferNoSpaceAroundRangeOperators3", testInferNoSpaceAroundRangeOperators3),
+        ("testInferNoSpaceAroundTimesAndDivideOperator", testInferNoSpaceAroundTimesAndDivideOperator),
+        ("testInferNoSpaceAroundTimesOperator", testInferNoSpaceAroundTimesOperator),
         ("testInferNoTrailingCommas", testInferNoTrailingCommas),
         ("testInferNoTruncateBlanklines", testInferNoTruncateBlanklines),
         ("testInferOctetHexGrouping", testInferOctetHexGrouping),
@@ -243,6 +245,9 @@ extension InferenceTests {
         ("testInferSelfInInitOnly", testInferSelfInInitOnly),
         ("testInferSpaceAfterOperatorFunc", testInferSpaceAfterOperatorFunc),
         ("testInferSpaceAroundRangeOperators", testInferSpaceAroundRangeOperators),
+        ("testInferSpaceAroundRangeOperators2", testInferSpaceAroundRangeOperators2),
+        ("testInferSpaceAroundRangeOperators3", testInferSpaceAroundRangeOperators3),
+        ("testInferSpaceAroundTimesOperator", testInferSpaceAroundTimesOperator),
         ("testInferThousands", testInferThousands),
         ("testInferTrailingCommas", testInferTrailingCommas),
         ("testInferUnhoisted", testInferUnhoisted),
@@ -1176,6 +1181,7 @@ extension RulesTests {
         ("testNoSpaceAroundMultipleOptionalChaining", testNoSpaceAroundMultipleOptionalChaining),
         ("testNoSpaceAroundPrefixMinus", testNoSpaceAroundPrefixMinus),
         ("testNoSpaceAroundRangeOperatorsWithCustomOptions", testNoSpaceAroundRangeOperatorsWithCustomOptions),
+        ("testNoSpaceAroundRangeOperatorsWithExplicitNoSpace", testNoSpaceAroundRangeOperatorsWithExplicitNoSpace),
         ("testNoSpaceBeforeColon", testNoSpaceBeforeColon),
         ("testNoSpaceBeforeComma", testNoSpaceBeforeComma),
         ("testNoSpaceBetweenArrayLiteralAndParen", testNoSpaceBetweenArrayLiteralAndParen),