Create a ShapedArray from nested arrays in Swift

I'm trying to create a shaped array structure for handling N-dimensional (multi-dimensional) numerical data in Swift. I would like to create the structure using nested arrays (array of arrays) like this:

// This would create a double-precision 2D array with a shape of 2x3
// ⎛ 1, 2, 3 ⎞
// ⎝ 4, 5, 6 ⎠
let arr = ShapedArray<Double>([[1, 2, 3], [4, 5, 6]])

// This would create a single-precision 3D array with a shape of 2x2x3
// ⎛ ⎛ 1, 2, 3 ⎞ ⎞
// ⎜ ⎝ 4, 5, 6 ⎠ ⎟
// ⎜ ⎛ 7, 8, 9 ⎞ ⎟
// ⎝ ⎝ 0, 1, 2 ⎠ ⎠
let arr = ShapedArray<Float>([[[1, 2, 3],
                               [4, 5, 6]],
                              [[7, 8, 9],
                               [0, 1, 2]]])

Below is my attempt to implement this. The ShapedArray structure stores the data in a flat array. The getShape function returns the shape of the nested arrays. The flatten function returns a single array (flattened array) from nested arrays.

func getShape(_ arr: some Collection) -> [Int] {
    if let first = arr.first as? any Collection {
        return [arr.count] + getShape(first)
    } else {
        return [arr.count]
    }
}

func flatten(_ arrays: [Any]) -> [Any] {
    var result = [Any]()

    for val in arrays {
        if let arr = val as? [Any] {
            result.append(contentsOf: flatten(arr))
        } else {
            result.append(val)
        }
    }

    return result
}

struct ShapedArray<T> {
    let shape: [Int]
    let data: [T]

    init(arrays: [Any]) {
        self.shape = getShape(arrays)
        self.data = flatten(arrays) as! [T]
    }
}

Some examples of using the ShapedArray struct are shown below.

// This works for 1D array
let sa = ShapedArray<Int>(arrays: [1, 2, 3, 4])

// This doesn't work
// let sa = ShapedArray<Float>(arrays: [1, 2, 3, 4])

// This works for 2D array
let saa = ShapedArray<Int>(arrays: [[0, 1, 2], [3, 4, 5]])

// This doesn't work
// let saa = ShapedArray<Float>(arrays: [[0, 1, 2], [3, 4, 5]])

// This works for 3D array
let saaa = ShapedArray<Int>(arrays: [[[0, 1, 2],
                                      [3, 4, 5]],
                                     [[6, 7, 8],
                                      [9, 0, 1]]])

This code only works for arrays of integers. If I try to create a ShapedArray of floats or doubles I get an error like this:

Could not cast value of type 'Swift.Int' (0x7ff84ad1b2a0) to 'Swift.Float' (0x7ff84ad1ae28).

So apparently the type information of the ShapedArray is not being passed to the result of the flattened array. How can I flatten nested arrays and define the element type? Or is there a different approach I should use?

Solution

(I will assume that the array is always well-formed, i.e. nothing like [1, [2, 3]] or [[1,2,3], [4,5]], since your code seems to assume this as well.)

The problem is that you are taking a [Any], so Swift has no idea that the array literals represent nested arrays of Float/Double. Since you are using integer literals as the array elements, that's what Swift thinks you want to create. You then cast these Ints to Float/Double, hence the crash.

This can be simply fixed annotating the type of the array literal:

let shapedArray = ShapedArray<Float>(arrays: 
     [[[0, 1, 2],
       [3, 4, 5]],
      [[6, 7, 8],              
       [9, 0, 1]]] as [[[Float]]]
)

That said, taking [Any] is not very type safe. A more type-safe solution would be to have the array of a sum type like this:

enum ArrayOrElement<T>: ExpressibleByArrayLiteral {
    case array([ArrayOrElement<T>])
    case element(T)
    
    init(arrayLiteral elements: ArrayOrElement<T>...) {
        self = .array(elements)
    }
}

// ExpressibleByXXXLiteral conformances
extension ArrayOrElement: ExpressibleByIntegerLiteral where T: ExpressibleByIntegerLiteral {
    init(integerLiteral value: T.IntegerLiteralType) {
        self = .element(T(integerLiteral: value))
    }
}

extension ArrayOrElement: ExpressibleByFloatLiteral where T: ExpressibleByFloatLiteral {
    init(floatLiteral value: T.FloatLiteralType) {
        self = .element(T(floatLiteral: value))
    }
}
// you can also do this for boolean and string literals...

Then you can write getShape and flatten that works with this type:

func getShape<T>(_ arr: [ArrayOrElement<T>]) -> [Int] {
    switch arr.first {
    case .array(let inner):
        [arr.count] + getShape(inner)
    default:
        [arr.count]
    }
}

func flatten<T>(_ arrays: [ArrayOrElement<T>]) -> [T] {
    var result = [T]()
    for val in arrays {
        switch val {
        case .array(let arr):
            result.append(contentsOf: flatten(arr))
        case .element(let e):
            result.append(e)
        }
    }

    return result
}

Then ShapedArray.init can take a [ArrayOrElement<T>].

struct ShapedArray<T>: ExpressibleByArrayLiteral {
    let shape: [Int]
    let data: [T]
    
    init(arrayLiteral elements: ArrayOrElement<T>...) {
        self.shape = getShape(elements)
        self.data = flatten(elements)
    }
}

Usage:

let shapedArray: ShapedArray<Double> = [
    [
        [0, 1, 2],
        [3, 4, 5],
    ],
    [
        [6, 7, 8],
        [9, 0, 1],
    ],
]

Note that the elements of the innermost array has to be literals for the implicit conversion to ArrayOrElement to work. If they are not literals (e.g. if you use a variable), you need to wrap it with .element(...).

You might also consider using MLShapedArray instead, but this only supports a small number of element types.