HTMLStreamer/Sources/HTMLStreamer/InlineArray3.swift

//
//  InlineArray3.swift
//  HTMLStreamer
//
//  Created by Shadowfacts on 11/19/23.
//

import Foundation

/// An array with inline space for up to 3 elements.
///
/// If the array grows beyond 3 elements, it will be stored out-of-line.
/// Once that happens, the array will never return to being stored inline,
/// since the allocation cost has already been paid.
struct InlineArray3<E> {
    private var storage: Storage
    
    init() {
        self.storage = .inline(nil, nil, nil)
    }
}

extension InlineArray3 {
    fileprivate enum Storage {
        case inline(Element?, Element?, Element?)
        case array(ContiguousArray<Element>)
    }
}

extension InlineArray3: ExpressibleByArrayLiteral {
    init(arrayLiteral elements: Element...) {
        switch elements.count {
        case 0:
            self.storage = .inline(nil, nil, nil)
        case 1:
            self.storage = .inline(elements[0], nil, nil)
        case 2:
            self.storage = .inline(elements[0], elements[1], nil)
        case 3:
            self.storage = .inline(elements[0], elements[1], elements[2])
        default:
            self.storage = .array(.init(elements))
        }
    }
}

extension InlineArray3: MutableCollection {
    typealias Element = E
    typealias Index = Int
    typealias Indices = Range<Int>
    
    subscript(position: Int) -> Element {
        _read {
            precondition(position < endIndex)
            switch storage {
            case .inline(let a, let b, let c):
                switch position {
                case 0:
                    yield a.unsafelyUnwrapped
                case 1:
                    yield b.unsafelyUnwrapped
                case 2:
                    yield c.unsafelyUnwrapped
                default:
                    fatalError("unreachable")
                }
            case .array(let arr):
                yield arr[position]
            }
        }
        _modify {
            precondition(position < endIndex)
            switch storage {
            case .inline(let a, let b, let c):
                switch position {
                case 0:
                    var newValue = a.unsafelyUnwrapped
                    yield &newValue
                    storage = .inline(newValue, b, c)
                case 1:
                    var newValue = b.unsafelyUnwrapped
                    yield &newValue
                    storage = .inline(a, newValue, c)
                case 2:
                    var newValue = c.unsafelyUnwrapped
                    yield &newValue
                    storage = .inline(a, b, newValue)
                default:
                    fatalError("unreachable")
                }
            case .array(var arr):
                yield &arr[position]
            }
        }
    }
    
    var startIndex: Int {
        0
    }
    
    var endIndex: Int {
        switch storage {
        case .inline(let a, let b, let c):
            a == nil ? 0 : b == nil ? 1 : c == nil ? 2 : 3
        case .array(let arr):
            arr.endIndex
        }
    }
}

extension InlineArray3: BidirectionalCollection {
}

extension InlineArray3: RandomAccessCollection {
}

extension InlineArray3: RangeReplaceableCollection {
    mutating func replaceSubrange<C>(_ subrange: Range<Int>, with newElements: C) where C: Collection, Element == C.Element {
        switch storage {
        case .array(var arr):
            arr.replaceSubrange(subrange, with: newElements)
            storage = .array(arr)
        case .inline(var a, var b, var c):
            if count - subrange.count + newElements.count <= 3 {
                // remove elements at subrange indices
                if subrange.contains(2) {
                    c = nil
                }
                if subrange.contains(1) {
                    b = c
                    c = nil
                }
                if subrange.contains(0) {
                    a = b
                    b = c
                    c = nil
                }
                
                // insert newElements starting at subrange.lowerBound
                for (offset, el) in newElements.enumerated() {
                    // assert that we have space to insert
                    assert(c == nil)
                    let newIndex = subrange.lowerBound + offset
                    switch newIndex {
                    case 2:
                        c = el
                    case 1:
                        c = b
                        b = el
                    case 0:
                        c = b
                        b = a
                        a = el
                    default:
                        fatalError("unreachable")
                    }
                }
                
                storage = .inline(a, b, c)
            } else {
                var arr: ContiguousArray = if let a {
                    if let b {
                        if let c {
                            [a, b, c]
                        } else {
                            [a, b]
                        }
                    } else {
                        [a]
                    }
                } else {
                    []
                }
                arr.replaceSubrange(subrange, with: newElements)
                storage = .array(arr)
            }
        }
    }
}

private extension Collection {
    func safeIndex(_ index: Index, offsetBy: Int) -> Index? {
        var index = index
        var offsetBy = offsetBy
        while offsetBy > 0 {
            if index < endIndex {
                formIndex(after: &index)
                offsetBy -= 1
            } else {
                return nil
            }
        }
        return index >= endIndex ? nil : index
    }
}

extension InlineArray3.Storage: Equatable where E: Equatable {
}

extension InlineArray3: Equatable where E: Equatable {
}

extension InlineArray3: CustomStringConvertible {
    var description: String {
        switch storage {
        case .inline(nil, nil, nil):
            return "[]"
        case .inline(.some(let a), nil, nil):
            return "[\(a)]"
        case .inline(.some(let a), .some(let b), nil):
            return "[\(a), \(b)]"
        case .inline(.some(let a), .some(let b), .some(let c)):
            return "[\(a), \(b), \(c)]"
        case .inline(_, _, _):
            fatalError("InlineArray3 invariant violated")
        case .array(let arr):
            return arr.description
        }
    }
}