A* Generic Implementation in Swift

Question

My first Code Review post.

I'm up to a generic implementation of the A* search algorithm in Swift (for now, it's a single goal implementation).

Here's what's been coded so far:

// State : a protocole for states in the search space
protocol State : Equatable {

    // successors() : returns an array of successors states in the search space
    func successors() -> [Successor<Self>]

    // heuristic(goal) : returns the heuristic value for a given states in relation to a given goal state
    func heuristic(goal:Self) -> Double

    // id : a string identifying a state
    var id : String { get }

}

// States are compared by their id
func ==<T:State>(lhs:T, rhs:T) -> Bool {
    return lhs.id==rhs.id
}

// Successor : represents a successor state and its cost
struct Successor<T:State> {
    var state: T
    var cost: Double
}

// Plan : a plan of states
struct Plan<T:State> {

    // states : an array of states that make a plan
    var states: [T]

    // cost : the total cost of the plan
    var cost: Double

    // isNot(another) : checks if another plan is different from the current one
    func isNot(another: Plan) -> Bool {
        return !(states == another.states && cost == another.cost)
    }

}

// AStar<TState> : finds the A* solution (nil if no solution found) given a start state and goal state
func AStar<TState:State>(start: TState, goal: TState) -> Plan<TState>? {

    var fringe : [Plan<TState>] = [Plan(states: [start], cost: 0)]

    while fringe.count>0 {

        let bestPlan = fringe.minElement({
            a,b
            in
            a.cost + a.states.last!.heuristic(goal) < b.cost + b.states.last!.heuristic(goal)
        })!

        fringe = fringe.filter({
            plan in plan.isNot(bestPlan)
        })

        if bestPlan.states.last! == goal {
            return bestPlan
        }else{
            let successors = bestPlan.states.last!.successors()
            for successor in successors.filter({ s in !bestPlan.states.contains(s.state) }) {
                let newPlan = Plan(states: bestPlan.states+[successor.state], cost: bestPlan.cost+successor.cost)
                fringe.append(newPlan)
            }
        }

    }

    return nil

}

I'm here for some suggestions to make this Swiftier

EDIT:

I tested this implementation using a graph of vertices representing cities with costs representing road distances. I used the straight line distance as a heuristic value :

let adjacencyList : [String:[String:Double]] = [
    "oued"      :   ["biskra":150.0],
    "biskra"    :   ["batna":120.0, "oued":150.0],
    "batna"     :   ["biskra":120.0, "barika":40.0, "setif":100.0, "constantine":110.0],
    "barika"    :   ["batna":40.0, "setif":50.0],
    "setif"     :   ["batna":100.0, "barika":50.0, "constantine":50.0, "bejaia":80.0],
    "constantine":  ["batna":110.0, "setif":50.0, "annaba":80.0],
    "bajaia"    :   ["setif":80.0, "annaba":30.0],
    "annaba"    :   ["constantine":80.0, "bejaia":30.0]
]

let straightLineDistances = [
    "biskra"    :   ["annaba":220.0],
    "batna"     :   ["annaba":140.0],
    "barika"    :   ["annaba":200.0],
    "setif"     :   ["annaba":100.0],
    "constantine":  ["annaba":80.0],
    "bejaia"    :   ["annaba":30.0],
    "oued"      :   ["annaba":320.0],
    "annaba"    :   ["annaba":0.0]
]


final class Vertex : State, CustomDebugStringConvertible {

    let label : String

    init(label:String) {
        self.label = label
    }

    func successors() -> [Successor<Vertex>] {
        return adjacencyList[label]!.map { x in Successor<Vertex>(state:Vertex(label: x.0),cost: x.1) }
    }

    func heuristic(goal:Vertex) -> Double {
        return straightLineDistances[label]![goal.label]!
    }

    var id : String {
        return label
    }

    var debugDescription : String {
        return id
    }

}

let solution = AStar(Vertex(label: "biskra"), goal: Vertex(label: "annaba"))

print(solution)

And the output solutions was the expected A* solution. But I'm more concerned about the elegance of this implementation

Answer 1

This code isn't bad at all. You make great use of the generic features of Swift and you also go for the functional approach over the iterative one whenever it's a good fit. Here are some ways to make your code "Swiftier":

1. Spacing. Readability is very important in Swift and your code currently doesn't follow all the best practices. Operators are usually surrounded by spaces (e.g. lhs.id == rhs.id) and colons are usually only followed by a space (e.g. protocol State: Equatable and var id: String).
2. Don't force unwrap. You used states.last! twice. Even though states is never empty in this case, the compiler doesn't actually enforce it. No matter how sure you are that a value isn't nil, it's generally frowned upon to use !. Instead, you could add var lastState: T to the Plan struct which keeps track of the last state.
3. Use trailing closures. If the last parameter of a function is a closure, you can put it outside the brackets: myArray.filter { $0 > 3 }.
4. The isNot(_:) function is not very Swifty. Using != would make much more sense, so the logical thing to do is to extend Plan to conform to Equatable. Also, I don't think you have to check whether cost == another.cost because there's no way the two would be different if the states are equal.
5. Don't use fringe.count > 0, but instead use !fringe.isEmpty. It's both more readable and for some data structures it's also more efficient (in case there's no constant time random access).
6. Instead of force unwrapping minElement, use while let bestPlan = fringe.minElement(...) instead.
7. Don't use short closure arguments such as a, b, x. Either give them more meaningful names, or use anonymous closure arguments ($0, $1, etc.).
8. Instead of using filter to remove bestPlan from fringe, use removeAtIndex instead. You can get this index using fringe.indexOf(bestPlan) (which you can unwrap simultaneously with fringe).
9. Use a guard to return bestPlan in case the goal is reached, both to increase readability and to prevent nesting.
10. Use a where-clause instead of filter to loop through all successors that are not contained in bestPlan.states. Using filter is good, but this is what where is for.
11. The way you construct newPlan isn't very elegant. You could consider adding an append function to Plan for adding a successor.

Here's what my version looks like after the changes:

// State : a protocole for states in the search space
protocol State: Equatable {

    // successors() : returns an array of successors states in the search space
    func successors() -> [Successor<Self>]

    // heuristic(goal) : returns the heuristic value in relation to a given goal state
    func heuristic(goal: Self) -> Double

    // id : a string identifying a state
    var id: String { get }

}

// States are compared by their id
func == <T: State> (lhs: T, rhs: T) -> Bool {
    return lhs.id == rhs.id
}

// Successor : represents a successor state and its cost
struct Successor<T: State> {
    var state: T
    var cost: Double
}

// Plan : a plan of states
struct Plan<T: State> {

    // states : an array of states that make a plan
    var states: [T]

    // lastState : the last state of the plan
    var lastState: T

    // cost : the total cost of the plan
    var cost: Double

    // initialise a plan with a single state
    init(state: T) {
        states = [state]
        lastState = state
        cost = 0
    }

    // append a successor to this plan
    mutating func append(successor: Successor<T>) {
        states.append(successor.state)
        lastState = successor.state
        cost += successor.cost
    }

    // the non-mutating version of append(_:)
    func appending(successor: Successor<T>) -> Plan {
        var new = self
        new.append(successor)
        return new
    }

}

extension Plan: Equatable {}

func == <T: State> (lhs: Plan<T>, rhs: Plan<T>) -> Bool {
    return lhs.states == rhs.states
}

// AStar<TState> : finds the A* solution (nil if no solution found) given a start state and goal state
func AStar <TState: State> (start: TState, goal: TState) -> Plan<TState>? {

    var fringe = [Plan(state: start)]

    // computes the best plan from the fringe array
    // I made this its own function to make the `while let` statement more readable
    func bestPlan() -> Plan<TState>? {
        return fringe.minElement {
            $0.cost + $0.lastState.heuristic(goal) < $1.cost + $1.lastState.heuristic(goal)
        }
    }

    while let bestPlan = bestPlan(), index = fringe.indexOf(bestPlan) {
        fringe.removeAtIndex(index)

        guard bestPlan.lastState != goal else { return bestPlan }

        let successors = bestPlan.lastState.successors()

        for successor in successors where !bestPlan.states.contains(successor.state) {
            fringe.append(bestPlan.appending(successor))
        }
    }

    return nil

}