An array with distributed storage, in Swift.

SegmentedArray.swift

//
//  SegmentedArray.swift
//  NodeCollections
//
//  Created by Daryle Walker on 4/15/18.
//  Copyright © 2018 Daryle Walker. All rights reserved.
//


// MARK: Globals

/// The maximum number of elements per segment.
let globalMaxElementsPerSegment = 128

// MARK: Primary Definition

/// An array that should have its elements in segments, to gain the memory advantages of linked lists.
public struct SegmentedArray<Element>: ExpressibleByArrayLiteral {

    /// The count of elements typically in a segment.
    public static var maxElementsPerSegment: Int { return globalMaxElementsPerSegment }

    /// A block of contigous elements.
    typealias Block = ContiguousArray<Element>
    /// Storage for multiple blocks.
    typealias SuperBlock = [Block]

    /// The storage: possibly-sparse array of contiguous arrays.
    var blocks: SuperBlock

    public init(arrayLiteral elements: Element...) {
        let meps = SegmentedArray.maxElementsPerSegment
        let elementCount = elements.count
        let (blockCount, stragglerCount) = elementCount.quotientAndRemainder(dividingBy: meps)
        blocks = SuperBlock()
        blocks.reserveCapacity(blockCount + stragglerCount.signum())
        for i in 0..<blockCount {
            var newBlock = Block()
            newBlock.reserveCapacity(meps)
            newBlock.append(contentsOf: elements[(i * meps)..<((i + 1) * meps)])
            blocks.append(newBlock)
        }
        if stragglerCount > 0 {
            var lastBlock = Block()
            lastBlock.reserveCapacity(meps)
            lastBlock.append(contentsOf: elements[(blockCount * meps)...])
            blocks.append(lastBlock)
        }
    }

}

// MARK: - Descriptions

extension SegmentedArray: CustomStringConvertible {

    /// Generate a string representation for the collection, possibly going deep for per-element descriptions.
    func normalOutput(useDebugOutput: Bool) -> String {
        let elementConversion: (Element) -> String = useDebugOutput ? String.init(reflecting:) : String.init(describing:)
        return "[\(blocks.map { $0.map(elementConversion).joined(separator: ", ") }.joined(separator: ", "))]"
    }

    public var description: String {
        return normalOutput(useDebugOutput: true)
    }

}

extension SegmentedArray: CustomDebugStringConvertible {

    /// Generate debugging data for the collection, and possibly its elements.
    func debuggingOutput(goingDeep: Bool) -> String {
        let elementConversion: (Element) -> String = goingDeep ? String.init(reflecting:) : String.init(describing:)
        var blockDescriptions = [String]()
        for b in blocks {
            let elementDescriptions = b.map(elementConversion).joined(separator: ", ")
            let reserved = Swift.max(0, Swift.min(b.capacity, SegmentedArray.maxElementsPerSegment) - b.count)
            let blockDescription: String
            if reserved > 0 {
                let reservedDescription = "(\(reserved) element(s) reserved)"
                blockDescription = [elementDescriptions, reservedDescription].joined(separator: " ")
            } else {
                blockDescription = elementDescriptions
            }
            blockDescriptions.append(blockDescription)
        }

        let blocksReserved = Swift.max(0, blocks.capacity - blocks.count)
        if blocksReserved > 0 {
            let reservedDescription = "(\(blocksReserved) block(s) reserved)"
            return "[\([blockDescriptions.joined(separator: "; "), reservedDescription].joined(separator: " "))]"
        } else {
            return "[\(blockDescriptions.joined(separator: "; "))]"
        }
    }

    public var debugDescription: String {
        return debuggingOutput(goingDeep: false)
    }

}

// MARK: Collection

extension SegmentedArray: BidirectionalCollection, MutableCollection, RangeReplaceableCollection {

    public struct Index: Comparable, Hashable {
        /// Points to the block of the target element
        let major: SuperBlock.Index
        /// Points to the element within the target block
        let minor: Block.Index

        public static func < (lhs: SegmentedArray<Element>.Index, rhs: SegmentedArray<Element>.Index) -> Bool {
            return (lhs.major, lhs.minor) < (rhs.major, rhs.minor)
        }
    }

    public var startIndex: Index {
        guard !blocks.isEmpty else { return Index(major: 0, minor: 0) }

        return Index(major: blocks.startIndex, minor: blocks.first!.startIndex)
    }
    public var endIndex: Index {
        guard !blocks.isEmpty else { return Index(major: 0, minor: 0) }

        return Index(major: blocks.indices.last!, minor: blocks.last!.endIndex)
    }

    public subscript(position: Index) -> Element {
        get { return blocks[position.major][position.minor] }
        set { blocks[position.major][position.minor] = newValue }
    }

    public func index(after i: Index) -> Index {
        var majorI = i.major
        var minorI = i.minor
        blocks[majorI].formIndex(after: &minorI)
        if minorI == blocks[majorI].endIndex, majorI < blocks.indices.last! {
            blocks.formIndex(after: &majorI)
            minorI = blocks[majorI].startIndex
        }
        return Index(major: majorI, minor: minorI)
    }
    public func index(before i: Index) -> Index {
        var majorI = i.major
        var minorI = i.minor
        if minorI == blocks[majorI].startIndex {
            blocks.formIndex(before: &majorI)
            minorI = blocks[majorI].endIndex
        }
        blocks[majorI].formIndex(before: &minorI)
        return Index(major: majorI, minor: minorI)
    }

    public init() {
        self = []
    }
    public mutating func replaceSubrange<C: Collection>(_ subrange: Range<Index>, with newElements: C) where C.Element == Element {
        let offset = distance(from: startIndex, to: subrange.lowerBound)
        removeSubrange(subrange)
        insert(contentsOf: newElements, at: index(startIndex, offsetBy: offset))
    }

    // MARK: Optimized Members

    public var underestimatedCount: Int { return blocks.first?.count ?? 0 }
    public var count: Int { return blocks.map { $0.count }.reduce(0, +) }

    public func distance(from start: Index, to end: Index) -> Int {
        // Prevent an illegal dereference from the second guard.
        guard start != startIndex || end != endIndex else { return count }
        // The third guard and the main code can't handle this.
        guard start.major != end.major else { return blocks[start.major].distance(from: start.minor, to: end.minor) }
        // Don't bother with reversal calculations in the main code.
        guard start.major < end.major else { return -distance(from: end, to: start) }

        return blocks[start.major].distance(from: start.minor, to: blocks[start.major].endIndex) + blocks[blocks.index(after: start.major) ..< end.major].map({ $0.count }).reduce(0, +) + blocks[end.major].distance(from: blocks[end.major].startIndex, to: end.minor)
    }
    public func index(_ i: Index, offsetBy n: Int, limitedBy limit: Index) -> Index? {
        guard n != 0 else { return i }

        var runningMinor = i.minor
        var runningOffset = n
        if n > 0 {
            guard limit >= i else { return index(i, offsetBy: n) }

            for bi in blocks.indices[i.major...] {
                let nextBlockIndex = blocks.index(after: bi)
                let limitMinor = (bi == limit.major) ? limit.minor : blocks[bi].endIndex
                if let resultMinor = blocks[bi].index(runningMinor, offsetBy: runningOffset, limitedBy: limitMinor) {
                    if resultMinor == blocks[bi].endIndex, nextBlockIndex < blocks.endIndex {
                        return Index(major: nextBlockIndex, minor: blocks[nextBlockIndex].startIndex)
                    } else {
                        return Index(major: bi, minor: resultMinor)
                    }
                } else if bi == limit.major {
                    return nil
                } else {
                    runningOffset -= blocks[bi].distance(from: runningMinor, to: limitMinor)
                    runningMinor = blocks[nextBlockIndex].startIndex
                }
            }
        } else {
            guard limit <= i else { return index(i, offsetBy: n) }

            for bi in blocks.indices[...i.major].reversed() {
                let limitMinor = (bi == limit.major) ? limit.minor : blocks[bi].startIndex
                if let resultMinor = blocks[bi].index(runningMinor, offsetBy: runningOffset, limitedBy: limitMinor) {
                    return Index(major: bi, minor: resultMinor)
                } else if bi == limit.major {
                    return nil
                } else {
                    runningOffset -= blocks[bi].distance(from: runningMinor, to: limitMinor)
                    runningMinor = blocks[blocks.index(before: bi)].endIndex
                }
            }
        }

        preconditionFailure("The block with the limiting index was never encountered")
    }
    public func index(_ i: Index, offsetBy n: Int) -> Index {
        guard n != 0 else { return i }

        if n > 0 {
            return index(i, offsetBy: n, limitedBy: endIndex)!
        } else {
            return index(i, offsetBy: n, limitedBy: startIndex)!
        }
    }

    public mutating func append(_ newElement: Element) {
        let meps = SegmentedArray.maxElementsPerSegment
        if (blocks.last?.count ?? Int.max) >= meps {
            var newBlock = Block()
            newBlock.reserveCapacity(meps)
            blocks.append(newBlock)
        }

        blocks[blocks.index(before: blocks.endIndex)].append(newElement)
    }
    public mutating func insert(_ newElement: Element, at i: Index) {
        insert(contentsOf: CollectionOfOne(newElement), at: i)
    }
    public mutating func insert<C: Collection>(contentsOf newElements: C, at i: Index) where C.Element == Element {
        guard !newElements.isEmpty else { return }
        guard i != endIndex else { return append(contentsOf: newElements) }

        // Fill in spaces around the insertion point with new elements before using new blocks.
        var nextBlockIndex: SuperBlock.Index
        let meps = SegmentedArray.maxElementsPerSegment
        var newIterator = newElements.makeIterator()
        var firstBlockSuffix = Block()
        firstBlockSuffix.reserveCapacity(meps)
        if i == startIndex {
            // Like the else block very below, but without risking dereferencing startIndex while empty.
            nextBlockIndex = blocks.startIndex
        } else if i.minor > blocks[i.major].startIndex {
            // Break off elements past the insertion point, turning insertions...
            firstBlockSuffix[...] = blocks[i.major][i.minor...]
            blocks[i.major].removeSubrange(i.minor...)

            // ...into local appends.
            while blocks[i.major].count < meps, let next = newIterator.next() {
                blocks[i.major].append(next)
            }

            nextBlockIndex = blocks.index(after: i.major)
        } else if let previousBlockIndex = blocks.index(i.major, offsetBy: -1, limitedBy: blocks.startIndex) {
            // Can't fill in forward, but let's try locally appending backwards.
            while blocks[previousBlockIndex].count < meps, let next = newIterator.next() {
                blocks[previousBlockIndex].append(next)
            }

            nextBlockIndex = i.major
        } else {
            // No wiggle room.
            nextBlockIndex = i.major
        }

        // Put the remaining new elements into new blocks.
        var newestBlock: Block? = nil
        while let next = newIterator.next() {
            if case .none = newestBlock {
                newestBlock = Block()
                newestBlock!.reserveCapacity(meps)
            }
            newestBlock!.append(next)

            if newestBlock!.count >= meps {
                blocks.insert(newestBlock!, at: nextBlockIndex)
                blocks.formIndex(after: &nextBlockIndex)
                newestBlock = nil
            }
        }

        // Insert the incomplete block of new elements.
        if let lastNewBlock = newestBlock, !lastNewBlock.isEmpty {
            blocks.insert(lastNewBlock, at: nextBlockIndex)
            blocks.formIndex(after: &nextBlockIndex)
            newestBlock = nil
        }

        // Re-insert the broken-off elements.
        while blocks[blocks.index(before: nextBlockIndex)].count < meps, !firstBlockSuffix.isEmpty {
            blocks[blocks.index(before: nextBlockIndex)].append(firstBlockSuffix.removeFirst())
        }

        if !firstBlockSuffix.isEmpty {
            if nextBlockIndex == blocks.endIndex {
                blocks.append(Block())
                blocks[nextBlockIndex].reserveCapacity(meps)
            }

            let transferCount = Swift.max(0, Swift.min(firstBlockSuffix.count, meps - blocks[nextBlockIndex].count))
            blocks[nextBlockIndex].insert(contentsOf: firstBlockSuffix.suffix(transferCount), at: blocks[nextBlockIndex].startIndex)
            firstBlockSuffix.removeLast(transferCount)
            if !firstBlockSuffix.isEmpty {
                blocks.insert(firstBlockSuffix, at: nextBlockIndex)
            }
        }
    }

    public mutating func remove(at position: Index) -> Element {
        defer { removeSubrange(position ..< index(after: position)) }
        return self[position]
    }
    public mutating func removeSubrange(_ bounds: Range<Index>) {
        switch (bounds.lowerBound == startIndex, bounds.upperBound == endIndex) {
        case (true, true):
            // Remove everything.
            blocks.removeAll(keepingCapacity: true)
        case (true, false):
            // Remove prefix.
            blocks[bounds.upperBound.major].removeSubrange(..<bounds.upperBound.minor)
            blocks.removeSubrange(..<bounds.upperBound.major)
        case (false, true):
            // Remove suffix.
            if bounds.lowerBound.minor == blocks[bounds.lowerBound.major].startIndex {
                blocks.removeSubrange(bounds.lowerBound.major...)
            } else {
                blocks[bounds.lowerBound.major].removeSubrange(bounds.lowerBound.minor...)
                blocks.removeSubrange(blocks.index(after: bounds.lowerBound.major)...)
            }
        case (false, false):
            // Remove in the middle.
            if bounds.lowerBound.major == bounds.upperBound.major {
                // Remove in a single block.
                blocks[bounds.lowerBound.major].removeSubrange(bounds.lowerBound.minor ..< bounds.upperBound.minor)
            } else {
                // Remove the suffix of the last prefix block and the prefix of the first suffix block.
                blocks[bounds.lowerBound.major].removeSubrange(bounds.lowerBound.minor...)
                blocks[bounds.upperBound.major].removeSubrange(..<bounds.upperBound.minor)

                // Remove the intermediate blocks.
                switch (blocks[bounds.lowerBound.major].isEmpty, blocks[bounds.upperBound.major].isEmpty) {
                case (false, false):
                    blocks.removeSubrange(blocks.index(after: bounds.lowerBound.major) ..< bounds.upperBound.major)
                case (false, true):
                    blocks.removeSubrange(blocks.index(after: bounds.lowerBound.major) ... bounds.upperBound.major)
                case (true, false):
                    blocks.removeSubrange(bounds.lowerBound.major ..< bounds.upperBound.major)
                case (true, true):
                    blocks.removeSubrange(bounds.lowerBound.major ... bounds.upperBound.major)
                }
            }
        }
    }

}

// MARK: Equality

extension SegmentedArray: Equatable where Element: Equatable {

    public static func == (lhs: SegmentedArray, rhs: SegmentedArray) -> Bool {
        return lhs.elementsEqual(rhs)
    }

}