More advanced ThreadBox with affinity guards

threadbox.kt

// This is a class that attempts to stop you accessing variables outside a lock.
//
// It does not do a perfect job, but can catch some common kinds of mistake, in
// particular when you accidentally try to work with objects inside closures that
// end up running later, outside the locked region (or in a different thread).

// EXAMPLE

val bank = ThreadBox(object {
    val accounts by arrayListOf(10, 0, 0, 0).guard()
})

fun transfer(from: Int, to: Int, amount: Int) {
    bank.accounts[0] = 0   // Does not compile: accounts not in scope

    bank.locked {
        accounts[from] -= amount
        accounts[to] += amount

        val ref = accounts     // Bad idea
        thread {
            ref[0] = 0         // Doesn't throw: guard was stripped
            accounts[0] = 0    // Throws IllegalStateException
        }.join()
    }
}


// IMPLEMENTATION

class ThreadBox<T>(v: T) {
    companion object {
        val inLockedRegion: ThreadLocal<ThreadBox<*>?> = ThreadLocal()
        val owners = WeakHashMap<Any, ThreadBox<*>>()
    }
    private val value = v

    init {
        owners[value] = this
    }

    synchronized fun locked<R>(f: T.() -> R): R {
        inLockedRegion.set(this)
        return try {
            value.f()
        } finally {
            inLockedRegion.set(null)
        }
    }
}

class Guard<T>(initialValue: T) : ReadOnlyProperty<kotlin.Any, T> {
    val value = initialValue

    override fun get(thisRef: Any, desc: PropertyMetadata): T {
        check(ThreadBox.inLockedRegion.get() identityEquals ThreadBox.owners[thisRef]) { "Locking error: ${ThreadBox.inLockedRegion.get()} vs ${thisRef}" }
        return value
    }
}

fun <T> T.guard(): Guard<T> = Guard(this)