Creating value-type generic Stack in swift with pointers and copy-on-write feature


Swift’s standard library contains all standard data structures you’d expect but it differs from other standard libraries that they’re value type and not reference type. Even their counter part like NSArray, NSMutableArray, NSDictionary etc in Objective-C’s Foundation are all reference type.

Lets try creating a very simple data structure: Stack with just two functionalities push and pop. push will append any element on top of stack and pop will remove the top most element.

The Protocol

protocol StackType {
    typealias Element
    mutating func push(element: Element)
    mutating func pop() -> Element?
}

The protocol is a simple protocol with just two methods marked mutating because the methods will be mutating the stack itself and since it’ll be a homogenous stack we associated the protocol with type Element which will be pushed or popped.

The Storage

To store the elements we’ll need some kind of storage, we can use swift’s Array to store the elements but that’ll be too simple. So lets create a generic class which can store our data.

final class BufferStorage<Element> {

    private var ptr: UnsafeMutablePointer<Element>
    private let capacity: Int

    init(capacity: Int) {
        ptr = UnsafeMutablePointer<Element>.alloc(capacity)
        self.capacity = capacity
    }

    static func copy(buffer: BufferStorage<Element>, count: Int) -> BufferStorage<Element> {
        let storage = BufferStorage<Element>(capacity: buffer.capacity)
        storage.ptr.initializeFrom(buffer.ptr, count: count)
        return storage
    }

    func add(element: Element, at position: Int) {
        (ptr + position).initialize(element)
    }

    func removeAt(position: Int) -> Element {
        let item = (ptr + position).memory
        (ptr + position).destroy()
        return item
    }

    func itemAt(position: Int) -> Element {
        return (ptr + position).memory
    }

    deinit {
        print("de inited")
        ptr.destroy(capacity)
        ptr.dealloc(capacity)
    }
}

This class contains a pointer which will be holding the type Element, which a generic so it can hold anything. For simplicity the pointer is allocated with a certain capacity so we don’t have to deal with expanding memory etc right now.

###Methods

init(capacity: Int)

The init method will allocate memory with capacity which will be using to store our Elements.

func add(element: Element, at position: Int)
func removeAt(position: Int) -> Element

These method will let us add and remove an element at a give position in the buffer

func itemAt(position: Int) -> Element

itemAt method will let us access any Element at the given position

deinit

The deinit method is the main reason we needed the BufferStorage to be a class. Swift doesn’t provide any option to do a cleanup when a value type is removed from a memory. deinit is called when a reference type is removed (A reference type is removed when there are zero references to it), we have to get rid of all the memory we asked for while creating this storage when our stack goes out of scope.

static func copy(buffer: BufferStorage<Element>, count: Int) -> BufferStorage<Element>

We’ll need to copy the storage when our stack is passed around and is mutated.

Note:

(ptr + position).memory This works because of + is overloaded in swift’s stdlib which gives the pointer incremented by the number provided.

public func +<Memory>(lhs: UnsafeMutablePointer<Memory>, rhs: Int) -> UnsafeMutablePointer<Memory>

The Stack

struct Stack<Element>: StackType {

    private var buffer: BufferStorage<Element>
    private var top = 0

    init(capacity: Int = 50) {
        buffer = BufferStorage(capacity: capacity)
    }

    mutating func push(element: Element) {
        if !isUniquelyReferencedNonObjC(&buffer) {
            buffer = BufferStorage.copy(buffer, count: top)
        }
        buffer.add(element, at: top)
        top = top + 1
    }

    mutating func pop() -> Element? {
        guard top > 0 else {
            return nil
        }
        if !isUniquelyReferencedNonObjC(&buffer) {
            buffer = BufferStorage.copy(buffer, count: top)
        }
        let item = buffer.removeAt(top-1)
        top = top - 1
        return item
    }
}

The stack is implemented as a struct containing two properies, buffer which is BufferStorage and top which will track the top of the stack.

The init method will initialize BufferStorage with the provided capacity.

The code is simple and self explanatory, so let just focus on isUniquelyReferencedNonObjC

isUniquelyReferencedNonObjC is a method in stdlib which tells you if a non-objc object ie pure swift object has reference count equal to one or not.

public func isUniquelyReferencedNonObjC<T : AnyObject>(inout object: T) -> Bool

so when one stack instance is assigned into another variable it will share the storage instance until push or pop is called on the new instance (or the old one) and it will detach the old storage and create a new copy for itself to use by calling the copy method on the storage.

if !isUniquelyReferencedNonObjC(&buffer) {
    buffer = BufferStorage.copy(buffer, count: top)
}

This is how copy-on-write is achieved.

It will also be useful to conform the stack to CustomStringConvertible so that its printable and SequenceType so we can use constructs like for..in

extension Stack: CustomStringConvertible {
    var description: String {
        guard top > 0 else {
            return "[]"
        }
        var str = "["
        for x in 0..<top-1 {
            str = str + "\((buffer.itemAt(x))), "
        }
        str = str + "\(buffer.itemAt(top-1))" + "]"
        return str
    }
}

struct StackGenerator<Element>: GeneratorType {
    var stack: Stack<Element>
    mutating func next() -> Element? {
        return stack.pop()
    }
}

extension Stack: SequenceType {
    func generate() -> StackGenerator<Element> {
        return StackGenerator<Element>(stack: self)
    }
}

Okay now let try our stack with Int.

var intStack = Stack<Int>()
intStack.push(1)
intStack.push(2)
intStack.push(3)
intStack.push(4)
intStack.push(5)

var newIntStack = intStack
print("Stacks before popping from new stack original: \(intStack) new: \(newIntStack)")
newIntStack.pop()
print("Stacks after popping from new stack original: \(intStack) new: \(newIntStack)")

This will output

Stacks before popping from new stack original: [1, 2, 3, 4, 5] new: [1, 2, 3, 4, 5]
Stacks after popping from new stack original: [1, 2, 3, 4, 5] new: [1, 2, 3, 4]

The values and memory is shared until pop is called on newIntStack and after pop is called newIntStack’s buffer is replaced with a new buffer.

We can also call SequenceType’s methods to do operations on the stack.

for x in newIntStack {
    print(x)
}

let stackValuesMultipliedByTwo = newIntStack.map { $0 * 2 }

Storing reference type in Stack

Now lets try the stack with reference types, suppose we had a simple class

class DemoClass: CustomStringConvertible {
    let tag: Int
    init(_ tag: Int) {
        self.tag = tag
    }
    deinit {
        print("removing...\(tag)")
    }
    var description: String {
        return "#\(tag)"
    }
}

And added few instances into stack

var classStack = Stack<DemoClass>()
classStack.push(DemoClass(1))
classStack.push(DemoClass(2))

var newClassStack = classStack
newClassStack.pop()

As you might expect popping from newClassStack will not print removing...2 as it is still being strongly referenced by the first stack classStack, however if you call classStack.pop() at the end of the above code removing...2 will be printed as theres no one left to hold that instance of DemoClass. This shows that reference types are not copied when assigned to a new stack and only the pointer holding reference to those objects are copied.

The entire playground friendly code is available here https://gist.github.com/aciidb0mb3r/fd147c8c6b1728664fdb