swift - Extending Collection with a recursive property/method that depends on the element type

Question

In the context of this question, I though about how one could implement a property or method that counts across all nesting levels in collections.

Intuitively, something this should work:

extension Collection { 
    var flatCount: Int {
        if self.count == 0 {
            return 0
        } else if self.first is Collection { // .Iterator.Element: Collection
            return self.reduce(0) { (res, elem) -> Int in
                res + (elem as! Collection).flatCount // ERROR
            }
        } else {
            return self.reduce(0) { (res,_) in res + 1 }
        }
    }
}

However, we are not allowed to cast values to protocol types that have associated types.

So I thought to make the Element type more explicit, like so:

extension Collection {
    var flatCount: Int {
        return Self.flatCountH(self)
    }

    private static final func 
        flatCountH<C: Collection, D>(_ c: C) -> Int 
            where Iterator.Element == D, D: Collection {
        return c.reduce(0) { (res: Int, elem: D) -> Int in 
            (res + elem.flatCount) as Int // Ambiguous type 
        }
    }

    private static final func flatCountH<C: Collection>(_ c: C) -> Int {
        return c.reduce(0) { $0 + $1.flatCount } // Unable to infer closure type
    }
}

But this is apparently asking too much of the type inferrer.

Now I took a step back and decided to stop trying to wrench everything into one extension:

extension Collection {
    var flatCount: Int {
        // There's no count on Collection, so...
        return self.reduce(0) { (res,_) in res + 1 }
    }
}

extension Collection where Iterator.Element: Collection {
    var flatCount: Int {
        return self.reduce(0) { $0 + $1.flatCount }
    }
}

Now this compiles -- yay! -- but the dispatching is off: $1.flatCount does not bind to the second, recursive version, but always to the first, plain version. That is, flatCount only counts the first nesting level.

Is there a way to juggle types and/or dispatching in a way to express this function? Or am I going about it in a completely roundabout way (or two)?

---

Side note: in the last example and in the first function, I don't use

self.reduce(0) { $0 + 1 }

because that does not compile; here, $0 is the pair of both anonymous parameters! I think that this is needlessly surprising behaviour and posted a change request to the Swift bugtracker.

Answer 1

I don't believe that it's currently possible to write a recursive extension like this, where the base case is determined by the conformance of a static type.

Although note that Collection does have a count property requirement, it's just of type IndexDistance (an associated type), rather than Int. Therefore if this were possible, you could express it as:

extension Collection {
    var flatCount: IndexDistance {
        return count
    }
}

extension Collection where Iterator.Element: Collection {
    var flatCount: IndexDistance {
        // compiler error: unable to infer closure type in the current context
        // (if you expand it out, it will tell you that it's because
        //  $1.flatCount is ambiguous)
        return self.reduce(0) { $0 + $1.flatCount }
    }
}

However, this yields a compiler error (although why it doesn't for when flatCount is an Int, I have no idea – they should both either consistently compile or not compile). The problem is that Swift wants to statically dispatch $1.flatCount – therefore meaning that it can only pick one of the extensions to call (and in this case, the compiler thinks both are equally valid).

The only way static dispatch could work here is if the implementations were specialised for each concrete type of Collection that they're called on. In that case, the ambiguity would be resolved as the compiler would know the concrete type inside the implementation, and thus know whether Iterator.Element.Iterator.Element : Collection, and dispatch accordingly.

However, currently specialisation is only an optimisation (due to the fact that it has the potential to drastically bloat code size without the use of inlining to counteract this additional cost) – therefore it's not possible to guarantee that static dispatch would work for all cases.

Even if $1.flatCount was able to be dynamically dispatched, through for example a protocol witness table (see this great WWDC talk on them), overload resolution based on the type constraints of the extensions would need to take place at runtime (in order to determine which extension to call). However, Swift doesn't resolve overloads at runtime (it would be expensive). Instead, the overload itself is resolved at compile time, and dynamic dispatch then allows the implementation of that overload to be polymorphic with respect to the value that it's called on (i.e it can dispatch to a the value's own implementation of that overload).

---

Unfortunately, I think probably the closest you'll be able to get is to write an extension for Array and use conditional type-casting in order to iterate through nested arrays:

extension Array {
    var flatCount: Int {

        var iterator = makeIterator()

        if let first = iterator.next() as? [Any] {
            // must be an array of arrays – otherwise $1 as! [Any] will crash.
            // feel free to add error handling or adding support for heterogeneous arrays
            // by doing an O(n) walk.
            return iterator.reduce(first.flatCount) { $0 + ($1 as! [Any]).flatCount }
        } else {
            return count
        }
    }
}

let arr = [[[[2, 3, 4]], [3, 4, 5, 6]], [57, 89]]

print(arr.flatCount) // 9

Although note that as @MartinR points out in the comments below, the conversion as(?/!) [Any] will create a new array in most cases (due to the difference in how Swift stores concrete typed and abstract typed values – see this Q&A), making the above implementation not particularly efficient.

One potential solution to this is to use a 'dummy protocol' in order to declare the flatCount property:

// dummy protocol to prevent conversions of arrays with concrete-typed elements to [Any].
protocol _Array {
    var flatCount: Int { get }
}

extension Array : _Array {
    var flatCount: Int {

        var iterator = makeIterator()

        if let first = iterator.next() as? _Array {
            // same comment as above, can crash for heterogeneous arrays.
            return iterator.reduce(first.flatCount) { $0 + ($1 as! _Array).flatCount }
        } else {
            return count
        }
    }
}

This avoids the O(n) conversion from an array of concrete-typed elements to abstract-typed elements (instead, only a single box is created for a given array).

If we do a rough quick benchmark of the two implementations (in a Release build on a MacBook Pro) with the array:

let arr = Array(repeating: Array(repeating: Array(repeating: 1, count: 100), count: 100), count: 1000)

For 10 repeated calls to flatCount, the first extension gives a time of 31.7 seconds. The same benchmark applied to the second implementation yields 0.93 seconds.