Matrix && Comonad

struct MatrixIndex: Comparable, CustomStringConvertible {
    let row: Int
    let col: Int
    
    var description: String {
        return "(\(row), \(col))"
    }
}

func ==(lhs: MatrixIndex, rhs: MatrixIndex) -> Bool {
    return (lhs.row, lhs.col) == (rhs.row, rhs.col)
}

func <(lhs: MatrixIndex, rhs: MatrixIndex) -> Bool {
    return lhs.row < rhs.row || lhs.col < rhs.col
}

struct MatrixRow<Element> {
    private let m: Matrix<Element>
    private let row: Int
    
    init(matrix: Matrix<Element>, row: Int) {
        self.m = matrix
        self.row = row
    }
    
    subscript(col: Int) -> Element {
        return m[(row, col)]
    }
}

struct Matrix<Element> {
    private var array: [Element]
    let size: (rows: Int, cols: Int)
    
    let startIndex = MatrixIndex(row: 0, col: 0)
    var endIndex: MatrixIndex {
        return MatrixIndex(row: size.rows, col: 0)
    }
    
    init(repeating: Element, size: (Int, Int)) {
        self.array = Array<Element>(repeating: repeating, count: size.0 * size.1)
        self.size = size
    }
    
    fileprivate init(array: Array<Element>, size: (Int, Int)) {
        self.array = array
        self.size = size
    }
    
    subscript(index: MatrixIndex) -> Element {
        get {
            return array[index.row * size.cols + index.col]
        }
        mutating set {
            array[index.row * size.cols + index.col] = newValue
        }
    }
    
    subscript(index: (Int, Int)) -> Element {
        get {
            return self[MatrixIndex(row: index.0, col: index.1)]
        }
        mutating set {
            self[MatrixIndex(row: index.0, col: index.1)] = newValue
        }
    }
    
    subscript(row: Int) -> MatrixRow<Element> {
        return MatrixRow(matrix: self, row: row)
    }
    
    func map<U>(_ transform: (Element) throws -> U) rethrows -> Matrix<U> {
        return try Matrix<U>(array: array.map(transform), size: size)
    }
    
    var debugDescription: String {
        return array.debugDescription
    }
    
    func index(after i: MatrixIndex) -> MatrixIndex {
        guard i != endIndex else {
            return i
        }
        
        let nextRow = i.col + 1 < size.cols ? i.row : i.row + 1
        let nextCol = nextRow == i.row ? i.col + 1 : 0
        return MatrixIndex(row: nextRow, col: nextCol)
    }
    
    func index(before i: MatrixIndex) -> MatrixIndex {
        guard i != startIndex else {
            return i
        }
        
        let nextRow = i.col > 0 ? i.row : i.row - 1
        let nextCol = (nextRow == i.row ? i.col : size.cols) - 1
        return MatrixIndex(row: nextRow, col: nextCol)
    }
}

extension Matrix: CustomDebugStringConvertible {}

struct ComonadMatrix<Element> {
    let matrix: Matrix<Element>
    let index: MatrixIndex
    
    init(matrix: Matrix<Element>) {
        self.init(matrix: matrix, index: matrix.startIndex)
    }
    
    init(matrix: Matrix<Element>, index: MatrixIndex) {
        self.matrix = matrix
        self.index = index
    }
    
    func comap<U>(_ transform: (ComonadMatrix<Element>) throws -> U) rethrows -> ComonadMatrix<U> {
        var array = Array<U>()
        var idx = matrix.startIndex
        while idx < matrix.endIndex {
            try array.append(transform(ComonadMatrix(matrix: matrix, index: idx)))
            idx = matrix.index(after: idx)
        }
        
        let m = Matrix(array: array, size: matrix.size)
        return ComonadMatrix<U>(matrix: m, index: index)
    }
}

var m = Matrix<Int>(repeating: 0, size: (4, 4))
do {
    var idx = m.startIndex
    var i = 1
    while idx < m.endIndex {
        m[idx] = i
        i += 1
        idx = m.index(after: idx)
    }
}

let m2 = ComonadMatrix(matrix: m).comap { comonad -> Int in
    let idx = comonad.index
    let newIdx = MatrixIndex(row: idx.col, col: comonad.matrix.size.rows - 1 - idx.row)
    return comonad.matrix[newIdx]
}