Refactoring Swift

By Vagmi on 16 Feb 2016

We are building a mobile app for enterprise asset management. One of the methods in our SyncManager class had a very complex merge method that merged the items in core data and the items received from the server. To give you a perspective this is what the method looked like.

This method takes in sorted business objects (Work Orders) from core data as the first argument and sorted JSON array on work order id as the second argument and does swift's equivalent of pointer increments to merge these lists together. The code has two vars and three while loops. While the code works per se, reasoning about a bug by say forgetting to increment a pointer is very hard.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
private func mergeWorkOrders(workOrders: [WorkOrder], jsonArray: [JSON]) {
  var jsonIndex = 0; var woIndex = 0

    while jsonIndex < jsonArray.count && woIndex < workOrders.count {


      let workOrder = workOrders[woIndex]
        let json = jsonArray[jsonIndex]

        let jsonWoId = json["WORKORDER", 0, "WORKORDERID"].int
        let woID = workOrder.workOrderID

        if jsonWoId == nil || woID == nil {
          print("ERROR: WorkOrder ID is nil")
            return
        }

      if jsonWoId == woID!.integerValue { //Update
        print("Sync: Update \(woID)")
          workOrder.buildFromDictinoary(json)
          woIndex++; jsonIndex++;

      } else if jsonWoId < woID!.integerValue { //Insert
        print("Sync: Insert \(woID)")
          if let aWork = NSEntityDescription.insertNewObjectForEntityForName("WorkOrder", 
              inManagedObjectContext: self.moc) as? WorkOrder {
            aWork.buildFromDictinoary(json)
          } else {
            print("ERROR: Cannot insert workorder to Core Data")
          }
        jsonIndex++

      } else {  //Delete
        print("Sync: Delete \(woID)")
          self.moc.deleteObject(workOrder)
          woIndex++
      }
    }

  /* Now process all remaining records */
  while jsonIndex < jsonArray.count {
    print("Sync: Insert")
      let json = jsonArray[jsonIndex]
      if let aWork = NSEntityDescription.insertNewObjectForEntityForName("WorkOrder", 
          inManagedObjectContext: self.moc) as? WorkOrder {
        aWork.buildFromDictinoary(json)
      } else {
        print("ERROR: Cannot insert workorder to Core Data")
      }
    jsonIndex++
  }

  while woIndex < workOrders.count {
    print("Sync: Delete")
      let workOrder = workOrders[woIndex]
      self.moc.deleteObject(workOrder)
      woIndex++
  }
}

The goal for our refactor is zero vars and a functional way to subtract the lists based on object keys. This is what the transformed functional code looks like.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
private func mergeWorkOrders(workOrders: [WorkOrder], jsonArray: [JSON]) {
  let jsonIdFn : JSON -> Int    = { $0["WORKORDER",0,"WORKORDERID"].intValue }
  let woIdFn : WorkOrder -> Int = { $0.workOrderID!.integerValue }

  let deleteableWOs  = workOrders.subtract(jsonArray, otherKey: jsonIdFn, on: woIdFn)
  let insertableWOs  = jsonArray.subtract(workOrders, otherKey: woIdFn, on: jsonIdFn)
  let modifiableWOs  = jsonArray.subtract(insertableWOs, on: jsonIdFn)
  let workOrdersById = workOrders.indexUniqueOn(woIdFn)

  modifiableWOs.forEach { saveWorkOrder($0, existingWorkOrder: workOrdersById[jsonIdFn($0)]) }
  insertableWOs.forEach { saveWorkOrder($0, existingWorkOrder: nil) }
  deleteableWOs.forEach { self.moc.deleteObject($0) }
}

private func saveWorkOrder(json: JSON, existingWorkOrder: WorkOrder?) {
  let wo = existingWorkOrder == nil ? NSEntityDescription.insertNewObjectForEntityForName("WorkOrder", inManagedObjectContext: self.moc) as? WorkOrder : existingWorkOrder
  wo?.buildFromDictinoary(json)
}

By most measures, the refactored code looks a bit better than the original. We have removed a lot of complexity from various branching and looping and have made the interface consistent and readable.

Swift does not provide the subtract and indexUniqueOn methods. I wrote a bunch of functions for extending the SequenceType.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
extension SequenceType {

  func groupBy<T: Hashable>(keyFunc: Generator.Element -> T) -> [T:[Generator.Element]] {
    var dict: [T:[Generator.Element]] = [:]
    for elem in self {
      let key = keyFunc(elem)
      if case nil = dict[key]?.append(elem) { dict[key] = [elem] }
    }
    return dict
  }

  func indexUniqueOn<T: Hashable>(keyFunc: Generator.Element -> T) -> [T:Generator.Element] {
    var dict: [T:Generator.Element] = [:]
    for elem in self {
      let key = keyFunc(elem)
      if case nil = dict[key] {
        dict[key] = elem
      }
    }
    return dict
  }
  
  func subtract<T: Hashable>(otherCollection: Self, on: Self.Generator.Element -> T) -> [Self.Generator.Element] {
    let otherCollectionByKey = otherCollection.indexUniqueOn(on)
    return filter { nil == otherCollectionByKey[on($0)] }
  }
  
  func subtract<T: Hashable,X where X: SequenceType>(otherCollection: X, otherKey: X.Generator.Element -> T, 
      on: Self.Generator.Element -> T) -> [Self.Generator.Element] {
    let otherCollectionByKey = otherCollection.indexUniqueOn(otherKey)
    return filter { nil == otherCollectionByKey[on($0)] }
  }
}

The subtract method takes in another collection and two lambdas that lets it arrive on the common key. This lets us subtract items of one sequence based on the one of the derived keys of the other sequence. The Swift's type system lets us set constraints on generics that help us express this beautifully. Swift's type system is actually quite nice. And this is what the tests for the extension look like.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
import XCTest
@testable import MainModule
class Person  : NSObject {
  var name:String
  var id:Int
  var age:Int
  init(name:String, id: Int, age: Int) {
    self.name = name
    self.id = id
    self.age = age
  }
}
class Employee : NSObject {
  var badgeId: Int
  var email: String
  init(badgeId: Int, email: String) {
    self.badgeId = badgeId
    self.email = email
  }
}
class TestSequenceTypeCategorization: XCTestCase {
  let person1 = Person(name: "Person 1", id: 1, age: 20)
  let person2 = Person(name: "Person 2", id: 22, age: 30)
  let person3 = Person(name: "Person 3", id: 33, age: 20)
  let person4 = Person(name: "Person 4", id: 22, age: 40)
  var persons:[Person] = []
  var persons2:[Person] = []
  
    override func setUp() {
        super.setUp()
      persons = [person1, person2, person3, person4]
      persons2 = [person1, person2, person4]
    }
    
    override func tearDown() {
        super.tearDown()
    }

  func testGroupBy() {
    let personsByAge = persons.groupBy { $0.age }
    assert(personsByAge[20]!.count==2, "The group by count is broken")
    assert(personsByAge[20]!.contains { $0 == person1 }, "Should contain person 1")
    assert(personsByAge[20]!.contains { $0 == person3 }, "Should contain person 3")
  }
  func testIndexOn() {
    let personsIndexedById = persons.indexUniqueOn { $0.id }
    assert(personsIndexedById.keys.count == 3, "should have three keys")
    assert(personsIndexedById[1]==person1, "Id 1 should get person 1")
    assert(personsIndexedById[22]==person2, "Id 22 should get person 2")
    assert(personsIndexedById[33]==person3, "Id 33 should get person 4")
  }
  func testSubtract() {
    let filteredPersons = persons.subtract(persons2, on: {$0.id})
    assert(filteredPersons.count==1,"there should be only one person")
    assert(filteredPersons[0] == person3,"and that should be person3")
  }
  func testSubtractOnOtherTypes() {
    let employees = [
      Employee(badgeId: 1, email: "test1@example.com"),
      Employee(badgeId: 22, email: "test22@example.com"),
      Employee(badgeId: 42, email: "test42@example.com")
    ]
    let filteredPersons = persons.subtract(employees, otherKey: {$0.badgeId} , on: {$0.id})
    let filteredEmployees = employees.subtract(persons, otherKey: {$0.id}, on: {$0.badgeId})
    assert(filteredPersons.count==1, "there must be only one fitered person")
    assert(filteredPersons[0] == person3, "the filtered person should be person 3")
    assert(filteredEmployees.count == 1, "there should be only one filtered employee")
    assert(filteredEmployees[0]==employees[2], "and that should be employee 42")
  }
}
comments powered by Disqus