Extend HooksRule on extension (#101)

* Extend HooksRule on extension

* Simplify HooksRule

* Remove return type

* Add FIXME

* Apply swiftformat

Sources/TokamakLint/Rules/HooksRule.swift

@@ -17,12 +17,12 @@ struct HooksRule: Rule {
   public static func validate(visitor: TokenVisitor) -> [StyleViolation] {
     var violations: [StyleViolation] = []
 
-    // search for render function
+    // search render function
     let structs = visitor.root.components(hookedComponentProtocols)
 
     structs.forEach { structDecl in
       for render in structDecl.children(with: "render") {
-        // search for Hooks argument name in the render argument list
+        // search Hooks argument name in the render argument list
         let renderFuncDecl = render.firstParent(of: .functionDecl)
         let funcSign = renderFuncDecl?.firstChild(of: .functionSignature)
         let funcParameterList = funcSign?
@@ -35,7 +35,7 @@ struct HooksRule: Rule {
           .firstChild(of: .codeBlockItemList)
         else { return }
 
-        // check that Hooks.state is on the first layer of the render function
+        // check that the Hooks.state is on the first layer of the render function
         hooks.forEach { hook in
           guard let hookMemberAccessExpr = hook
             .firstParent(of: .memberAccessExpr),
@@ -55,6 +55,46 @@ struct HooksRule: Rule {
       }
     }
 
+    // search Hooks extension
+    let extensions = visitor.root.children(with: .extensionDecl)
+      .filter { ext in
+        let simpleTypeIdentifier = ext.firstChild(of: .simpleTypeIdentifier)
+        return simpleTypeIdentifier?.children[0].text == "Hooks"
+      }
+
+    extensions.forEach { ext in
+      // search all memberDeclListItem
+      guard let memberDeclList = ext.firstChild(of: .memberDeclList) else {
+        return
+      }
+      memberDeclList.children.forEach { memberDeclListItem in
+        // search `state` use in memberDeclListItem
+        // FIXME: this should be extended on all possible hooks, not only `state`
+        let states = memberDeclListItem.children(with: "state")
+
+        // search codeBlockItemList in memberDeclListItem
+        guard let memberCodeBlockItemList = memberDeclList.firstChild(of: .codeBlockItemList) else { return }
+
+        states.forEach { state in
+          let stateCodeBlock = state.firstParent(of: .codeBlockItem)
+          guard let safeState = stateCodeBlock else { return }
+          guard memberCodeBlockItemList.children.contains(safeState) else {
+            violations.append(
+              StyleViolation(
+                ruleDescription: OneRenderFunctionRule.description,
+                location: Location(
+                  file: visitor.path,
+                  line: state.range.startRow,
+                  character: state.range.startColumn
+                )
+              )
+            )
+            return
+          }
+        }
+      }
+    }
+
     return violations
   }
 }

Tests/TokamakCLITests/TokamakLintTests.swift

@@ -47,43 +47,43 @@ final class TokamakLintTests: XCTestCase {
 
   func testTwoComponentsCorrectBroken() throws {
     let path = "\(try srcRoot())/TwoComponentsCorrectBroken.swift"
-    let oneRenderFunctionRuleResult = try OneRenderFunctionRule.validate(path: path)
-    XCTAssertEqual(oneRenderFunctionRuleResult.count, 2)
-    XCTAssertEqual(oneRenderFunctionRuleResult[0].location.line, 41)
-    XCTAssertEqual(oneRenderFunctionRuleResult[1].location.line, 60)
+    let result = try OneRenderFunctionRule.validate(path: path)
+    XCTAssertEqual(result.count, 2)
+    XCTAssertEqual(result[0].location.line, 41)
+    XCTAssertEqual(result[1].location.line, 60)
   }
 
   func testTwoComponentsBrokenBroken() throws {
     let path = "\(try srcRoot())/TwoComponentsBrokenBroken.swift"
-    let oneRenderFunctionRuleResult = try OneRenderFunctionRule.validate(path: path)
-    XCTAssertEqual(oneRenderFunctionRuleResult.count, 4)
-    XCTAssertEqual(oneRenderFunctionRuleResult[0].location.line, 10)
-    XCTAssertEqual(oneRenderFunctionRuleResult[1].location.line, 29)
-    XCTAssertEqual(oneRenderFunctionRuleResult[2].location.line, 60)
-    XCTAssertEqual(oneRenderFunctionRuleResult[3].location.line, 79)
+    let result = try OneRenderFunctionRule.validate(path: path)
+    XCTAssertEqual(result.count, 4)
+    XCTAssertEqual(result[0].location.line, 10)
+    XCTAssertEqual(result[1].location.line, 29)
+    XCTAssertEqual(result[2].location.line, 60)
+    XCTAssertEqual(result[3].location.line, 79)
   }
 
   func testTwoComponentsCorrectCorrect() throws {
     let path = "\(try srcRoot())/TwoComponentsCorrectCorrect.swift"
-    let oneRenderFunctionRuleResult = try OneRenderFunctionRule.validate(path: path)
-    XCTAssertEqual(oneRenderFunctionRuleResult.count, 0)
+    let result = try OneRenderFunctionRule.validate(path: path)
+    XCTAssertEqual(result.count, 0)
   }
 
   func testHooksRulePositive() throws {
     let path = "\(try srcRoot())/HooksRulePositive.swift"
-    let oneRenderFunctionRuleResult = try HooksRule.validate(path: path)
-    XCTAssertEqual(oneRenderFunctionRuleResult.count, 0)
+    let result = try HooksRule.validate(path: path)
+    XCTAssertEqual(result.count, 0)
   }
 
   func testHooksRuleNegative() throws {
     let path = "\(try srcRoot())/HooksRuleNegative.swift"
-    let oneRenderFunctionRuleResult = try HooksRule.validate(path: path)
-    XCTAssertEqual(oneRenderFunctionRuleResult.count, 6)
-    XCTAssertEqual(oneRenderFunctionRuleResult[0].location.line, 7)
-    XCTAssertEqual(oneRenderFunctionRuleResult[1].location.line, 15)
-    XCTAssertEqual(oneRenderFunctionRuleResult[2].location.line, 20)
-    XCTAssertEqual(oneRenderFunctionRuleResult[3].location.line, 29)
-    XCTAssertEqual(oneRenderFunctionRuleResult[4].location.line, 34)
-    XCTAssertEqual(oneRenderFunctionRuleResult[5].location.line, 42)
+    let result = try HooksRule.validate(path: path)
+    XCTAssertEqual(result.count, 12)
+    let violationsLines = [
+      7, 15, 20, 40, 45, 53, 79, 81, 29, 31, 69, 71,
+    ]
+    for (i, line) in violationsLines.enumerated() {
+      XCTAssertEqual(result[i].location.line, line)
+    }
   }
 }

ValidationTests/HooksRuleNegative.swift

@@ -22,6 +22,17 @@ struct HookedLeafComponent: LeafComponent {
   }
 }
 
+// don't use Hooks in extesion on non first level
+extension Hooks {
+  var blah: State<Int> {
+    if true {
+      return state(0)
+    } else {
+      return state(42)
+    }
+  }
+}
+
 struct AnotherHookedLeafComponent: LeafComponent {
   static func render(props: Props, hooks: Hooks) -> AnyNode {
     // do not use hooks in the loop
@@ -43,3 +54,41 @@ struct AnotherHookedLeafComponent: LeafComponent {
     }
   }
 }
+
+// good extension among broken
+extension Hooks {
+  var theAnswerToLifeTheUniverseAndEverything: State<Int> {
+    return state(42)
+  }
+}
+
+// don't use state in conditions
+extension Hooks {
+  func test(_ condition: Bool) -> State<Int> {
+    if condition {
+      return state(0)
+    } else {
+      return state(42)
+    }
+  }
+}
+
+struct ConditionHookedLeafComponent: LeafComponent {
+  static func render(props: Props, hooks: Hooks) {
+    if props.condition {
+      return hooks.blah
+    } else {
+      return hooks.blah
+    }
+  }
+}
+
+extension Hooks {
+  var whatDoesTheDogSay: State<String> {
+    return state("Woof-Woof")
+  }
+
+  func sayHiTo(name: String) -> String {
+    return "Hi \(name)!"
+  }
+}

ValidationTests/HooksRulePositive.swift

@@ -6,14 +6,34 @@ struct HookedLeafComponent: LeafComponent {
   }
 }
 
+extension Hooks {
+  var theAnswerToLifeTheUniverseAndEverything: State<Int> {
+    return state(42)
+  }
+}
+
 struct AnotherHookedLeafComponent: LeafComponent {
   static func render(props: Props, hooks: Hooks) -> AnyNode {
     let hookMadeWithLove = hooks.state("")
   }
 }
 
+extension Hooks {
+  var whatDoesTheDogSay: State<String> {
+    return state("Woof-Woof")
+  }
+
+  func sayHiTo(name: String) -> String {
+    return "Hi \(name)!"
+  }
+}
+
 struct OneMoreHookedLeafComponent: LeafComponent {
   static func render(props: Props, hooks: Hooks) -> AnyNode {
     let hookMadeWithLove = hooks.state("")
   }
 }
+
+extension String {
+  let str = "number"
+}