Matrix.swift

import Foundation

public struct Matrix<Scalar: CustomStringConvertible> {
  public let width: Int
  public let height: Int
  public var grid: [Scalar]
}

// MARK: - Creating matrices

extension Matrix {
  public init(width: Int, height: Int, repeatedValue: Scalar) {
    self.width = width
    self.height = height
    self.grid = .init(repeating: repeatedValue, count: width * height)
  }

  public init(size: (Int, Int), repeatedValue: Scalar) {
    self.init(width: size.0, height: size.1, repeatedValue: repeatedValue)
  }

  public init() {
    (width, height) = (0, 0)
    grid = []
  }
}

extension Matrix where Scalar: ExpressibleByIntegerLiteral {
  /* Creates a matrix where each element is 0. */
  public static func zeros(_ width: Int, height: Int) -> Matrix<Scalar> {
    return Matrix<Scalar>(width: width, height: height, repeatedValue: 0)
  }

  public static func zeros(_ size: (Int, Int)) -> Matrix<Scalar> {
    return Matrix<Scalar>(size: size, repeatedValue: 0)
  }
}

// MARK: - Querying the matrix

extension Matrix {
  public var size: (Int, Int) {
    return (width, height)
  }

  public var length: Int {
    return Swift.max(width, height)
  }

  public var numCells: Int {
    return grid.count
  }

  public subscript(x: Int, y: Int) -> Scalar {
    get { return grid[(y * width) + x] }
    set { grid[(y * width) + x] = newValue }
  }
}

// MARK: - Printable

extension Matrix: CustomStringConvertible {
  public var description: String {
    var description = ""

    for i in 0..<height {
      let contents = (0..<width).map{ self[$0, i].description }.joined(separator: " ")

      switch (i, height) {
      case (0, 1):
        description += "( \(contents) )\n"
      case (0, _):
        description += "⎛ \(contents) ⎞\n"
      case (height - 1, _):
        description += "⎝ \(contents) ⎠\n"
      default:
        description += "⎜ \(contents) ⎥\n"
      }
    }
    return description
  }
}

// MARK: - SequenceType

/* Lets you iterate through the rows of the matrix. */
extension Matrix: Sequence {
  public func makeIterator() -> AnyIterator<ArraySlice<Scalar>> {
    let endIndex = width * height
    var nextRowStartIndex = 0
    return AnyIterator {
      if nextRowStartIndex == endIndex {
        return nil
      } else {
        let currentRowStartIndex = nextRowStartIndex
        nextRowStartIndex += self.width
        return self.grid[currentRowStartIndex..<nextRowStartIndex]
      }
    }
  }
}

// MARK: Helpers
extension Matrix {
  public func tryGet(_ x: Int, _ y: Int) -> Scalar? {
    return (x >= 0 && x < width && y >= 0 && y < height)
        ? self[x, y]
        : nil
  }
}