ToyGC.swift

//
//  ToyGC.swift
//  ToyGC
//
//  Created by satoshia on 2015/08/31.
//  Copyright © 2015 satoshia. All rights reserved.
//

import Darwin
import Foundation


var GC = ToyGC()

func gcInit(maxHeapSize heapSize: Int, stackBase: UnsafePointer<UInt>) {
    GC = ToyGC(maxHeapSize: heapSize, stackBase: stackBase)
}

struct Cell {
    typealias Pointer = UnsafeMutablePointer<Cell>
    var value: Int
    var nextPtr: Pointer

    init() {
        value = 0
        nextPtr = nil
    }

    static func alloc() -> Pointer {
        return GC.alloc()
    }
}

struct ToyGC {
//    var bitmaps: [UInt32]
    let bitmaps: UnsafeMutablePointer<UInt32> // to avoid swift_bridgeObjectRetain()
    let bitmapsSize: Int
    var currentIndex: Int

    var _availableCells: Int

    var size: Int
    var heap: Cell.Pointer
    var heapEnd: Cell.Pointer

    var stackBase: UnsafePointer<UInt>  // not real stack base

    init(maxHeapSize heapSize: Int = 1 << 16, stackBase: UnsafePointer<UInt> = nil) {
        if stackBase == nil {
            let addr = pthread_get_stackaddr_np(pthread_self())
            self.stackBase = UnsafePointer<UInt>(addr)
        } else {
            self.stackBase = stackBase
        }

        if heapSize <= 0 {
            size = 0
            heap = nil
            heapEnd = nil
            currentIndex = 0
            _availableCells = 0
//            bitmaps = []
            bitmaps = nil
            bitmapsSize = 0
            return
        }

        size = heapSize
        heap = Cell.Pointer.alloc(size)
        var ptr = heap
        for _ in 0..<size {
            ptr.initialize(Cell())
            ptr++
        }
        heapEnd = ptr.predecessor()

//        bitmaps = [UInt32](count: size / 32, repeatedValue: ~UInt32(0))
        bitmapsSize = size / 32
        bitmaps = UnsafeMutablePointer<UInt32>.alloc(bitmapsSize)
        var p = bitmaps
        for _ in 0..<(bitmapsSize) {
            p.initialize(~UInt32(0))
            p++
        }

        currentIndex = 0
        _availableCells = size
    }

    var availableCells: Int {
        return _availableCells
    }

    func dispose() {
        heap.dealloc(size)
        bitmaps.dealloc(bitmapsSize)
    }

    mutating
    func gc() {
        mark()
    }

    mutating
    func alloc() -> Cell.Pointer  {
        if let obj = lazySweep() {
            return obj
        }
        mark()
        if let obj = lazySweep() {
            return obj
        }
        fatalError("\(__FUNCTION__): Heap empty")
    }

    mutating
    func lazySweep() -> Cell.Pointer? {
        var i = currentIndex
        while bitmaps[i] == 0 {
            ++i
//            if i == bitmaps.count {
            if i == bitmapsSize {
                return nil
            }
        }
        if i != currentIndex {
            currentIndex = i
        }

        let bits = bitmaps[i]
        let n = bits.numberOfLeadingZeros
        setMark((i, n))
        return heap.advancedBy((i << 5) + n)
    }

    mutating
    func mark() {
        // reset bitmap table
        currentIndex = 0
        _availableCells = size
//        for i in 0..<bitmaps.count {
        for i in 0..<bitmapsSize {
            bitmaps[i] = ~UInt32(0)
        }

        markStack(0,0,0,0,0,0)
    }

    /// should check the disassembled code of the function prologue.
    /// The registers (rbx, r12~r15) have to be pushed to the stack.
    /// (X86_64 only).
    @inline(never) mutating
    func markStack(rdi: Int, _ rsi: Int, _ rdx: Int, _ rcx: Int, _ r8: Int, _ r9: Int) -> Int {
        // GC starts from the address of this local variable (on the stack).
        var v = UInt(0)
        var ptr = withUnsafePointer(&v) {$0}
        while ptr <= stackBase {
            let theValue = ptr.memory
            if theValue == 0 || !maybePointer(theValue) || !isHeapObject(theValue) {
                ptr++
                continue
            }

            // O.K. the object is alive, mark it.
            let obj = UnsafePointer<Cell.Pointer>(ptr).memory
            let loc = locateBitmap(obj)
            if !isMarked(loc) {
                setMark(loc)

                // usually in general...
                // for child in obj.children {
                //    mark(child)
                //}
                var cellPtr = obj.memory.nextPtr
                while cellPtr != nil {
                    let loc = locateBitmap(cellPtr)
                    if !isMarked(loc) {
                        setMark(loc)
                        cellPtr = cellPtr.memory.nextPtr
                    } else {
                        break
                    }
                }
            }
            ptr++
        }

        return rdi | rsi | rdx | rcx | r8 | r9 // no meaning
    }

    func maybePointer(val: UInt) -> Bool {
        // pointer is always a multiple of 8 (in case of 64bit arch.)
        return val & 0x7 == 0
    }

    func isHeapObject(val: UInt) -> Bool {
        let aPtr = UnsafePointer<Int8>(bitPattern: val)
        let heapPtr = UnsafePointer<Int8>(heap)
        let heapEndPtr = UnsafePointer<Int8>(heapEnd)

        // does the pointer point to something inside heap?
        if aPtr < heapPtr || aPtr > heapEndPtr {
            return false
        }
        // a multiple of the size of a cell?
        if sizeof(Cell.self) != 16 {
            fatalError("\(__FUNCTION__): sizeof(Cell) must be 16 (2^n). If you change the def of Cell and the size becomes not to fit 2^n, you can't use '&' operator. Re-implement to use '%' operator...")
        }
//        if (aPtr - heapPtr) % sizeof(Cell.self) != 0 {
        if (aPtr - heapPtr) & 0xf != 0 {
            return false
        }
        return true
    }

    typealias BitMapLocation = (index: Int, offset: Int)
    func locateBitmap(obj: Cell.Pointer) -> BitMapLocation {
        let index  = (obj - heap) >> 5
        let offset = (obj - heap) & 0x1f
        return (index, offset)
    }

    func isMarked(loc: BitMapLocation) -> Bool {
        // mark: 0, unmark:1
        return bitmaps[loc.index] & (UInt32(1) << UInt32(0x1f - loc.offset)) == 0

    }

    func isMarked(obj: Cell.Pointer) -> Bool {
        let index  = (obj - heap) >> 5
        let offset = (obj - heap) & 0x1f
        // mark: 0, unmark:1
        return bitmaps[index] & (UInt32(1) << UInt32(0x1f - offset)) == 0
    }

    mutating
    func setMark(loc: BitMapLocation) {
        // mark: 0, unmark:1
        bitmaps[loc.index] &= ~(UInt32(1) << UInt32(0x1f - loc.offset))
        _availableCells--
    }

    mutating
    func setMark(obj: Cell.Pointer) {
        let index  = (obj - heap) >> 5
        let offset = (obj - heap) & 0x1f
        // mark: 0, unmark:1
        bitmaps[index] &= ~(UInt32(1) << UInt32(0x1f - offset))
        _availableCells--
    }
}


extension UInt32 {
    /// see Hacker's Delight 2nd ed. 5-3
    var numberOfLeadingZeros: Int {
        if self == 0 { return 32 }
        var x = self
        var n:UInt32 = 1
        if x >> 16 == 0 {n += 16; x <<= 16}
        if x >> 24 == 0 {n +=  8; x <<=  8}
        if x >> 28 == 0 {n +=  4; x <<=  4}
        if x >> 30 == 0 {n +=  2; x <<=  2}
        return Int(n - (x >> UInt32(31)))
    }
}