timboudreau · July 20, 2025 05:12
diff --git a/TreiberishStack.swift b/TreiberishStack.swift
 import Atomics

 /// A thread-safe, lockless, concurrent sequence which can be appended to and drained.
 /// This is not quite a Treiber stack in the sense that single elements cannot
 /// be **popped** off of it - Treiber stacks rely on every interaction with a given memory
 /// address being a *single atomic read, write or exchange operation*. The Swift atomics
 /// package is a bit too anemic for that - we would need to be able to exchange a pointer
 /// to the head cell with a pointer to the head cell's child (if any) without reading the
 /// head pointer twice.  That's implementable in C or assembly, but not with the access
 /// Swift atomics give us.
 ///
 /// So in order to make this work, unlike a traditional treiber stack, we append to the
 /// *tail* rather than swapping the head for a new head with the old head as its child
 /// Iteration is in FIFO order.  So, Trieber-ish stack, not Treiber stack.
 ///
 /// Implements `Sequence` - the iterator returned may grow in length if the stack is
 /// pushed to while iterating, but will never get smaller, even if the stack is emptied.
 ///
 public final class TreiberishStack<Element : Sendable> : Sequence, Sendable {
    // IMPORTANT: If we don't use `ordering: .sequentiallyConsistent` here, we will encounter occasional
    // segfaults when setting the head to nil.
    private let _head : ManagedAtomic<TreiberCell<Element>?>
    private var head : TreiberCell<Element>? { _head.load(ordering: .acquiring) }
    public var first : Element? { head?.element }

    public var last : Element? {
        var h = head
        while h != nil {
            if let n = h!.next {
                h = n
            } else {
                return h!.element
            }
        }
        return nil
    }

    public var count : Int {
        var h = head
        var result = 0
        while let hh = h {
            result += 1
            h = hh.next
        }
        return result
    }
    
    /// isEmpty is constant time
    public var isEmpty : Bool { head == nil }
    
    /// Inherited from Sequence; count returned here will be precisely accurate
    public var underestimatedCount: Int { self.count }

    public subscript(_ index : Int) -> Element? { head?[index] }
    
    public init() { _head  = .init(nil) }
    
    public init<S: Sequence<Element>>(_ sequence : S) where S.Iterator.Element == Element {
        var head : TreiberCell<Element>? = nil
        for element in sequence {
            if let h = head {
                h.push(TreiberCell(element))
            } else {
                head = TreiberCell(element)
            }
        }
        _head = .init(head)
    }
    
    /// Append an element
    public func push(_ element : Element) {
        let cell = TreiberCell(element)
        if let existingHead = _head.compareExchange(expected: nil, desired: cell, ordering: .sequentiallyConsistent).1 {
            existingHead.push(cell)
        }
    }
    
    /// Append many elements in a single atomic operation
    public func push<S: Sequence>(contentsOf : S) where S.Element == Element {
        // Build our tail before we append it
        var append : TreiberCell<Element>? = nil
        for item in contentsOf {
            if let a = append {
                a.push(TreiberCell(item))
            } else {
                append = TreiberCell(item)
            }
        }
        if let a = append {
            if let existingHead = _head.compareExchange(expected: nil, desired: a, ordering: .sequentiallyConsistent).1 {
                existingHead.push(a)
            }
        }
    }
    
    /// Returns an iterator which may grow if the stack is concurrently
    /// appended to, but cannot shrink.  Iteration is in FIFO order.
    public func makeIterator() -> TreiberishStackIterator<Element> {
        .init(cell: head)
    }
    
    public func clear() {
        _head.store(nil, ordering: .releasing)
    }
    
    public func drain() -> [Element] {
        var result : [Element] = []
        if let oldHead = _head.exchange(nil, ordering: .sequentiallyConsistent) {
            // Pending: analyze the compiled code and see if we were really getting
            // tail recursion here, in which case, put it back.
 //            oldHead.visit{ result.append($0) }
            var h : TreiberCell<Element>? = oldHead
            while let hh = h {
                result.append(hh.element)
                h = hh.next
            }
        }
        return result
    }
    
    public func exchangeHead(_ newHead : Element) -> [Element] {
        var result : [Element] = []
        if let oldHead = _head.exchange(TreiberCell(newHead), ordering: .sequentiallyConsistent) {
 //            oldHead.visit{ result.append($0) }
            var h : TreiberCell<Element>? = oldHead
            while let hh = h {
                result.append(hh.element)
                h = hh.next
            }
        }
        return result
    }
    
    public func map<T>(_ f : (Element) throws -> T) rethrows -> [T] {
        var result : [T] = []
        if let h = head {
            try h.visit { element in
                result.append(try f(element))
            }
        }
        return result
    }

    public func compactMap<T>(_ f : (Element) throws -> T?) rethrows -> [T] {
        var result : [T] = []
        if let h = head {
            try h.visit { element in
                if let t = try f(element) {
                    result.append(t)
                }
            }
        }
        return result
    }

    public func contains(where f : (Element) -> Bool) -> Bool {
        if let h = head {
            var curr : TreiberCell<Element>? = h
            while let c = curr {
                if f(c.element) {
                    return true
                }
                curr = c.next
            }
        }
        return false
    }
 }

 fileprivate final class TreiberCell<Element : Sendable> : Sendable, AtomicOptionalWrappable, AtomicReference {
    private let _next : ManagedAtomicLazyReference<TreiberCell<Element>> = .init()
    let element : Element
    var next : TreiberCell<Element>? { _next.load() }
    var count : Int { 1 + (next?.count ?? 0) }
    
    var last : Element {
        if let n = next {
            return n.last
        }
        return element
    }
    
    subscript(_ index : Int) -> Element? {
        if index < 0 {
            return nil
        }
        switch index {
        case 0 :
            return element
        default:
            if let n = self.next {
                return n[index - 1]
            }
            return nil
        }
    }
    
    init(_ element: Element) { self.element = element }
    
    func visit(_ f : (Element) throws -> Void) rethrows {
        try f(element)
        if let next = self.next {
            try next.visit(f)
        }
    }
    
    func push(_ cell : TreiberCell<Element>) {
        let next = _next.storeIfNilThenLoad(cell)
        if next !== cell {
            next.push(cell)
        }
    }
 }

 public struct TreiberishStackIterator<E : Sendable> : IteratorProtocol {
    fileprivate var cell : TreiberCell<E>?
    
    public mutating func next() -> E? {
        guard let current = cell else { return nil }
        defer { cell = current.next }
        return current.element
    }
 }
	import Atomics

	/// A thread-safe, lockless, concurrent sequence which can be appended to and drained.
	/// This is not quite a Treiber stack in the sense that single elements cannot
	/// be popped off of it - Treiber stacks rely on every interaction with a given memory
	/// address being a single atomic read, write or exchange operation. The Swift atomics
	/// package is a bit too anemic for that - we would need to be able to exchange a pointer
	/// to the head cell with a pointer to the head cell's child (if any) without reading the
	/// head pointer twice. That's implementable in C or assembly, but not with the access
	/// Swift atomics give us.
	///
	/// So in order to make this work, unlike a traditional treiber stack, we append to the
	/// tail rather than swapping the head for a new head with the old head as its child
	/// Iteration is in FIFO order. So, Trieber-ish stack, not Treiber stack.
	///
	/// Implements `Sequence` - the iterator returned may grow in length if the stack is
	/// pushed to while iterating, but will never get smaller, even if the stack is emptied.
	///
	public final class TreiberishStack<Element : Sendable> : Sequence, Sendable {
	// IMPORTANT: If we don't use `ordering: .sequentiallyConsistent` here, we will encounter occasional
	// segfaults when setting the head to nil.
	private let _head : ManagedAtomic<TreiberCell<Element>?>
	private var head : TreiberCell<Element>? { _head.load(ordering: .acquiring) }
	public var first : Element? { head?.element }

	public var last : Element? {
	var h = head
	while h != nil {
	if let n = h!.next {
	h = n
	} else {
	return h!.element
	}
	}
	return nil
	}

	public var count : Int {
	var h = head
	var result = 0
	while let hh = h {
	result += 1
	h = hh.next
	}
	return result
	}

	/// isEmpty is constant time
	public var isEmpty : Bool { head == nil }

	/// Inherited from Sequence; count returned here will be precisely accurate
	public var underestimatedCount: Int { self.count }

	public subscript(_ index : Int) -> Element? { head?[index] }

	public init() { _head = .init(nil) }

	public init<S: Sequence<Element>>(_ sequence : S) where S.Iterator.Element == Element {
	var head : TreiberCell<Element>? = nil
	for element in sequence {
	if let h = head {
	h.push(TreiberCell(element))
	} else {
	head = TreiberCell(element)
	}
	}
	_head = .init(head)
	}

	/// Append an element
	public func push(_ element : Element) {
	let cell = TreiberCell(element)
	if let existingHead = _head.compareExchange(expected: nil, desired: cell, ordering: .sequentiallyConsistent).1 {
	existingHead.push(cell)
	}
	}

	/// Append many elements in a single atomic operation
	public func push<S: Sequence>(contentsOf : S) where S.Element == Element {
	// Build our tail before we append it
	var append : TreiberCell<Element>? = nil
	for item in contentsOf {
	if let a = append {
	a.push(TreiberCell(item))
	} else {
	append = TreiberCell(item)
	}
	}
	if let a = append {
	if let existingHead = _head.compareExchange(expected: nil, desired: a, ordering: .sequentiallyConsistent).1 {
	existingHead.push(a)
	}
	}
	}

	/// Returns an iterator which may grow if the stack is concurrently
	/// appended to, but cannot shrink. Iteration is in FIFO order.
	public func makeIterator() -> TreiberishStackIterator<Element> {
	.init(cell: head)
	}

	public func clear() {
	_head.store(nil, ordering: .releasing)
	}

	public func drain() -> [Element] {
	var result : [Element] = []
	if let oldHead = _head.exchange(nil, ordering: .sequentiallyConsistent) {
	// Pending: analyze the compiled code and see if we were really getting
	// tail recursion here, in which case, put it back.
	// oldHead.visit{ result.append($0) }
	var h : TreiberCell<Element>? = oldHead
	while let hh = h {
	result.append(hh.element)
	h = hh.next
	}
	}
	return result
	}

	public func exchangeHead(_ newHead : Element) -> [Element] {
	var result : [Element] = []
	if let oldHead = _head.exchange(TreiberCell(newHead), ordering: .sequentiallyConsistent) {
	// oldHead.visit{ result.append($0) }
	var h : TreiberCell<Element>? = oldHead
	while let hh = h {
	result.append(hh.element)
	h = hh.next
	}
	}
	return result
	}

	public func map<T>(_ f : (Element) throws -> T) rethrows -> [T] {
	var result : [T] = []
	if let h = head {
	try h.visit { element in
	result.append(try f(element))
	}
	}
	return result
	}

	public func compactMap<T>(_ f : (Element) throws -> T?) rethrows -> [T] {
	var result : [T] = []
	if let h = head {
	try h.visit { element in
	if let t = try f(element) {
	result.append(t)
	}
	}
	}
	return result
	}

	public func contains(where f : (Element) -> Bool) -> Bool {
	if let h = head {
	var curr : TreiberCell<Element>? = h
	while let c = curr {
	if f(c.element) {
	return true
	}
	curr = c.next
	}
	}
	return false
	}
	}

	fileprivate final class TreiberCell<Element : Sendable> : Sendable, AtomicOptionalWrappable, AtomicReference {
	private let _next : ManagedAtomicLazyReference<TreiberCell<Element>> = .init()
	let element : Element
	var next : TreiberCell<Element>? { _next.load() }
	var count : Int { 1 + (next?.count ?? 0) }

	var last : Element {
	if let n = next {
	return n.last
	}
	return element
	}

	subscript(_ index : Int) -> Element? {
	if index < 0 {
	return nil
	}
	switch index {
	case 0 :
	return element
	default:
	if let n = self.next {
	return n[index - 1]
	}
	return nil
	}
	}

	init(_ element: Element) { self.element = element }

	func visit(_ f : (Element) throws -> Void) rethrows {
	try f(element)
	if let next = self.next {
	try next.visit(f)
	}
	}

	func push(_ cell : TreiberCell<Element>) {
	let next = _next.storeIfNilThenLoad(cell)
	if next !== cell {
	next.push(cell)
	}
	}
	}

	public struct TreiberishStackIterator<E : Sendable> : IteratorProtocol {
	fileprivate var cell : TreiberCell<E>?

	public mutating func next() -> E? {
	guard let current = cell else { return nil }
	defer { cell = current.next }
	return current.element
	}
	}
No results found