I recently had occasion to give my old Sudoku talk again. For those who haven’t seen the talk, I live-code my way through a Sudoku problem and together we write a solver that can solve any valid grid. It’s a very fun talk to give and I hope it’s enjoyable to watch as well.

While preparing for the talk, I took the chance to update and modernize some of the code.

I was able to use multi-line strings to represent the grid in a graphical way and get rid of an old helper that is now part of the standard library as allSatisfy(_:). More important than those changes, though, I was able to incorporate the new “primary associated types” for protocols, which had a way bigger impact on the code than I expected.

Let dig in. Here’s how the grid is structured:

public struct Grid {
     private var rows: [[Cell]]
}

There are a few ways to represent the data, but I chose an array of arrays, which represent the rows of the grid.

However, the underlying data structure is not as important, because the primary mode of interaction with this object is through the “units” of the grid, like rows, columns, and boxes. Here are some helpers to extract those:

extension Grid { 
     public var cells: [Cell] {
         return Array(rows.joined())
     }

     public func row(forIndex index: Int) -> [Cell] {
         let rowIndex = index / 9
         return rows[rowIndex]
     }

     public func column(forIndex index: Int) -> [Cell] {
         let columnIndex = index % 9
         return self.rows.map({ row in
             return row[columnIndex]
         })
     }

     public func box(forIndex index: Int) -> [Cell] {
         let rowIndex = index / 9
         let columnIndex = index % 9
         let boxColumnIndex = columnIndex / 3
         let boxRowIndex = rowIndex / 3
         return (0..<3).flatMap({ rowOffset in
             return self.rows[boxRowIndex*3+rowOffset][boxColumnIndex*3..<boxColumnIndex*3+3]
         })
     }
}

The box function is definitely the most daunting here, but what I want to look at here is the return types of all these functions. They all return arrays. Every time you request any of these, the relevant data is copied into a new array. This happens so many times during the solve that it starts to add up. For example, isSolved looks like this:

 public var isSolved: Bool {
     return cells.allSatisfy({ $0.isSettled })
 }

It asks for all the cells in the grid, which creates an 81-cell array, copies all the data to it, and then iterates that array exactly one time, looking for a cell that isn’t settled.

Before Swift 5.7, for the cells property, you could return some Collection. This enables you get rid of the of conversion to an Array, which saves you the copying.

 public var cells: some Collection {
     return rows.joined()
 }

Sadly, this pre-Swift-5.7 code is nearly useless. Even though Swift can infer the full return type of this function, the contract with external callers doesn’t include the type of the element. They get to know it’s a collection and that it has some element type, but not what the element is. They have to use it as an opaque value. It was a pretty big shortcoming when working with the some keyword.

Swift 5.7 changes all that. Protocols can now provide some “primary” associated types, which callers can rely on using the angle-bracket generics syntax that we all know and love:

 public var cells: some Collection<Cell> {
     return rows.joined()
 }

This is huge. Now, since .joined() produces a “lazy” collection, the collection stays lazy, and doesn’t have to copy its items to a new array. Less work, big upside, and the contract with callers includes the element type.

Now, you could say this isn’t anything new. After all, I could have written the concrete type of rows.joined() for the return type of the cells property, and seen the same benefit:

 public var cells: JoinedSequence<[Cell]> {
     return rows.joined()
 }

Not too bad, right? However, here’s what the return type for box(forIndex:) looks like when you add a .lazy to get the intended performance:

public func box(forIndex index: Int) -> LazySequence<
    FlattenSequence<
        LazyMapSequence<
            LazySequence<(Range<Int>)>.Elements,
            ArraySlice<Cell>
        >
    >
> {
    let rowIndex = index / 9
    let columnIndex = index % 9
    let boxColumnIndex = columnIndex / 3
    let boxRowIndex = rowIndex / 3
    return (0..<3).lazy.flatMap({ rowOffset in
        return self.rows[boxRowIndex*3+rowOffset][boxColumnIndex*3..<boxColumnIndex*3+3]
    })
}

This isn’t usable. Not only is it hideous, these are implementation details, fully on display for the world to see. If you ever change the implementation of the function or the internal structure of the Grid type, you’ll have to change the return type as well. If this was a library, callers could even rely on details of the return type that you didn’t intend to expose.

The typical pre-5.7 way around the problem of messy return types was to use AnyCollection. Using AnyCollection has performance trade offs though, since it boxes its value, slowing down accesses and other operations.

How bad is that trade off? My expectation was that not returning an Array and returning an AnyCollection instead (which helps you avoid the expensive copy) would get you almost all of the way there. However, it turns out that you get about a 15% improvement in the solver’s runtime with AnyCollection. Switching to some Collection<Cell> gets you an additional 15% improvement. This means that boxing up the collection is about half as bad as fully copying every item to a new buffer. These are ultimately pretty short arrays (under 100 items each) so algorithms that operate on bigger data will benefit even more from these tweaks.

Another lesson (that I always feel like I have to learn afresh) is that you shouldn’t make assumptions about performance. Always profile! Your intuitions are fallible.

So here’s all of the helpers, returning some Collection<Cell> with the new syntax:

public var cells: some Collection<Cell> {
    return rows.joined()
}

public func row(forIndex index: Int) -> some Collection<Cell> {
    let rowIndex = index / 9
    return rows[rowIndex]
}

public func column(forIndex index: Int) -> some Collection<Cell> {
    let columnIndex = index % 9
    return self.rows.lazy.map({ row in
        return row[columnIndex]
    })
}

public func box(forIndex index: Int) -> some Collection<Cell> {
    let rowIndex = index / 9
    let columnIndex = index % 9
    let boxColumnIndex = columnIndex / 3
    let boxRowIndex = rowIndex / 3
    return (0..<3).lazy.flatMap({ rowOffset in
        return self.rows[boxRowIndex*3+rowOffset][boxColumnIndex*3..<boxColumnIndex*3+3]
    })
}

The code is nicer, better reflects its intent, and hides implementation details from the caller. These are all worthwhile ends in themselves, but we also see substantial performance improvements. This feature is a huge boon for collection-heavy code. By making the simple thing the fast thing and the safe thing for libraries you end up with a win on all fronts.