Aoc2018Day17.swift

extension ClosedRange where Bound == Int {
    init(_ string: Substring) {
        if let range = string.range(of: "..") {
            let x = Int(string.prefix(through: range.lowerBound).dropLast())!
            let y = Int(string.suffix(from: range.upperBound))!
            
            self = x...y
        } else {
            let x = Int(string)!
            self = x...x
        }
    }
}

struct Vein {
    let x, y: ClosedRange<Int>
    
    init(_ string: Substring) {
        let c = string.split(separator: ",")
        let ranges = (ClosedRange(c[0].dropFirst(2)), ClosedRange(c[1].dropFirst(3)))
        
        if string.first == "x" {
            (x, y) = ranges
        } else {
            (y, x) = ranges
        }
    }
}

struct Grid<Element> {
    let xRange: ClosedRange<Int>
    let yRange: ClosedRange<Int>
    private(set) var elements: [Element]
    
    init(repeating element: Element, x: ClosedRange<Int>, y: ClosedRange<Int>) {
        self.xRange = x
        self.yRange = y
        self.elements = Array(repeating: element, count: xRange.count * yRange.count)
    }
    
    private func index(x: Int, y: Int) -> Int {
        assert(xRange.contains(x) && yRange.contains(y))
        return xRange.count * (y - yRange.lowerBound) + x - xRange.lowerBound
    }
    
    subscript(x x: Int, y y: Int) -> Element {
        get { return elements[index(x: x, y: y)] }
        set { elements[index(x: x, y: y)] = newValue }
    }
}

enum State {
    case empty, clay, flowing, settled
    
    var isFlowing: Bool {
        guard case .flowing = self else { return false }
        return true
    }
    
    var isSettled: Bool {
        guard case .settled = self else { return false }
        return true
    }
    
    var isWater: Bool {
        return isFlowing || isSettled
    }
}

enum Direction {
    case down, left, right
}

func day17(input: String) -> (part1: Int, part2: Int) {
    let veins = input.split(separator: "\n").map(Vein.init)
    let xRange = veins.map { $0.x.lowerBound }.min()! - 1...veins.map { $0.x.upperBound }.max()! + 1
    let yRange = veins.map { $0.y.lowerBound }.min()!...veins.map { $0.y.upperBound }.max()!
    
    var grid = Grid(repeating: State.empty, x: xRange, y: yRange)
    
    for vein in veins {
        for x in vein.x {
            for y in vein.y {
                grid[x: x, y: y] = .clay
            }
        }
    }
    
    func fill(x: Int, y: Int, direction: Direction) -> Bool {
        guard yRange.contains(y) else { return false }
        
        let state = grid[x: x, y: y]
        guard state == .empty else { return !state.isFlowing }
        
        grid[x: x, y: y] = .flowing
        guard fill(x: x, y: y + 1, direction: .down) else { return false }
        
        let l = direction == .right || fill(x: x - 1, y: y, direction: .left)
        let r = direction == .left  || fill(x: x + 1, y: y, direction: .right)
        
        guard l && r else { return false }
        
        if direction == .down {
            func isFlowing(_ x: Int) -> Bool {
                return grid[x: x, y: y].isFlowing
            }
            
            func settle(_ x: Int) {
                grid[x: x, y: y] = .settled
            }
            
            settle(x)
            sequence(first: x - 1, next: { $0 - 1 }).prefix(while: isFlowing).forEach(settle)
            sequence(first: x + 1, next: { $0 + 1 }).prefix(while: isFlowing).forEach(settle)
        }
        
        return true
    }
    
    _ = fill(x: 500, y: yRange.lowerBound, direction: .down)
    
    return (
        grid.elements.count(where: { $0.isWater }),
        grid.elements.count(where: { $0.isSettled })
    )
}