Consolidate/Simplify the foreign bindings handlemap code

CHANGELOG.md

@@ -16,6 +16,7 @@
 
 - The callback interface code was reworked to use vtables rather than a single callback method.
   See https://github.com/mozilla/uniffi-rs/pull/1818 for details and how the other bindings were updated.
+- Removed `FfiType::RustFutureContinuationData`.  `RustFutureContinuation` callbacks now take a `u64` handle.
 
 ## v0.26.0 (backend crates: v0.26.0) - (_2024-01-23_)
 

uniffi_bindgen/src/bindings/kotlin/gen_kotlin/mod.rs

@@ -379,7 +379,6 @@ impl KotlinCodeOracle {
             FfiType::Struct(name) => self.ffi_struct_name(name),
             FfiType::Reference(inner) => self.ffi_type_label_by_reference(inner),
             FfiType::VoidPointer | FfiType::RustFutureHandle => "Pointer".to_string(),
-            FfiType::RustFutureContinuationData => "USize".to_string(),
         }
     }
 

uniffi_bindgen/src/bindings/kotlin/templates/Async.kt

@@ -3,18 +3,18 @@
 internal const val UNIFFI_RUST_FUTURE_POLL_READY = 0.toByte()
 internal const val UNIFFI_RUST_FUTURE_POLL_MAYBE_READY = 1.toByte()
 
-internal val uniffiContinuationHandleMap = UniFfiHandleMap<CancellableContinuation<Byte>>()
+internal val uniffiContinuationHandleMap = UniffiHandleMap<CancellableContinuation<Byte>>()
 
 // FFI type for Rust future continuations
 internal object uniffiRustFutureContinuationCallback: UniffiRustFutureContinuationCallback {
-    override fun callback(data: USize, pollResult: Byte) {
-        uniffiContinuationHandleMap.remove(data)?.resume(pollResult)
+    override fun callback(data: Long, pollResult: Byte) {
+        uniffiContinuationHandleMap.remove(data).resume(pollResult)
     }
 }
 
 internal suspend fun<T, F, E: Exception> uniffiRustCallAsync(
     rustFuture: Pointer,
-    pollFunc: (Pointer, UniffiRustFutureContinuationCallback, USize) -> Unit,
+    pollFunc: (Pointer, UniffiRustFutureContinuationCallback, Long) -> Unit,
     completeFunc: (Pointer, UniffiRustCallStatus) -> F,
     freeFunc: (Pointer) -> Unit,
     liftFunc: (F) -> T,

uniffi_bindgen/src/bindings/kotlin/templates/CallbackInterfaceRuntime.kt

@@ -1,37 +1,3 @@
-{{- self.add_import("java.util.concurrent.atomic.AtomicLong") }}
-{{- self.add_import("java.util.concurrent.locks.ReentrantLock") }}
-{{- self.add_import("kotlin.concurrent.withLock") }}
-
-internal typealias UniffiHandle = Long
-internal class ConcurrentHandleMap<T>(
-    private val leftMap: MutableMap<UniffiHandle, T> = mutableMapOf(),
-) {
-    private val lock = java.util.concurrent.locks.ReentrantLock()
-    private val currentHandle = AtomicLong(0L)
-    private val stride = 1L
-
-    fun insert(obj: T): UniffiHandle =
-        lock.withLock {
-            currentHandle.getAndAdd(stride)
-                .also { handle ->
-                    leftMap[handle] = obj
-                }
-            }
-
-    fun get(handle: UniffiHandle) = lock.withLock {
-        leftMap[handle] ?: throw InternalException("No callback in handlemap; this is a Uniffi bug")
-    }
-
-    fun delete(handle: UniffiHandle) {
-        this.remove(handle)
-    }
-
-    fun remove(handle: UniffiHandle): T? =
-        lock.withLock {
-            leftMap.remove(handle)
-        }
-}
-
 // Magic number for the Rust proxy to call using the same mechanism as every other method,
 // to free the callback once it's dropped by Rust.
 internal const val IDX_CALLBACK_FREE = 0
@@ -40,8 +6,8 @@ internal const val UNIFFI_CALLBACK_SUCCESS = 0
 internal const val UNIFFI_CALLBACK_ERROR = 1
 internal const val UNIFFI_CALLBACK_UNEXPECTED_ERROR = 2
 
-public abstract class FfiConverterCallbackInterface<CallbackInterface>: FfiConverter<CallbackInterface, UniffiHandle> {
-    internal val handleMap = ConcurrentHandleMap<CallbackInterface>()
+public abstract class FfiConverterCallbackInterface<CallbackInterface: Any>: FfiConverter<CallbackInterface, UniffiHandle> {
+    internal val handleMap = UniffiHandleMap<CallbackInterface>()
 
     internal fun drop(handle: UniffiHandle) {
         handleMap.remove(handle)

uniffi_bindgen/src/bindings/kotlin/templates/HandleMap.kt

@@ -0,0 +1,30 @@
+// Handle from a UniffiHandleMap
+internal typealias UniffiHandle = Long
+
+// Map handles to objects
+//
+// This is used pass an opaque 64-bit handle representing a foreign object to the Rust code.
+internal class UniffiHandleMap<T: Any> {
+    private val map = ConcurrentHashMap<UniffiHandle, T>()
+    private val counter = java.util.concurrent.atomic.AtomicLong(0)
+
+    val size: Int
+        get() = map.size
+
+    // Insert a new object into the handle map and get a handle for it
+    fun insert(obj: T): UniffiHandle {
+        val handle = counter.getAndAdd(1)
+        map.put(handle, obj)
+        return handle
+    }
+
+    // Get an object from the handle map
+    fun get(handle: UniffiHandle): T {
+        return map.get(handle) ?: throw InternalException("UniffiHandleMap.get: Invalid handle")
+    }
+
+    // Remove an entry from the handlemap and get the Kotlin object back
+    fun remove(handle: UniffiHandle): T {
+        return map.remove(handle) ?: throw InternalException("UniffiHandleMap: Invalid handle")
+    }
+}

uniffi_bindgen/src/bindings/kotlin/templates/Helpers.kt

@@ -77,51 +77,6 @@ private inline fun <U> uniffiRustCall(callback: (UniffiRustCallStatus) -> U): U
     return uniffiRustCallWithError(UniffiNullRustCallStatusErrorHandler, callback);
 }
 
-// IntegerType that matches Rust's `usize` / C's `size_t`
-public class USize(value: Long = 0) : IntegerType(Native.SIZE_T_SIZE, value, true) {
-    // This is needed to fill in the gaps of IntegerType's implementation of Number for Kotlin.
-    override fun toByte() = toInt().toByte()
-    // Needed until https://youtrack.jetbrains.com/issue/KT-47902 is fixed.
-    @Deprecated("`toInt().toChar()` is deprecated")
-    override fun toChar() = toInt().toChar()
-    override fun toShort() = toInt().toShort()
-
-    fun writeToBuffer(buf: ByteBuffer) {
-        // Make sure we always write usize integers using native byte-order, since they may be
-        // casted to pointer values
-        buf.order(ByteOrder.nativeOrder())
-        try {
-            when (Native.SIZE_T_SIZE) {
-                4 -> buf.putInt(toInt())
-                8 -> buf.putLong(toLong())
-                else -> throw RuntimeException("Invalid SIZE_T_SIZE: ${Native.SIZE_T_SIZE}")
-            }
-        } finally {
-            buf.order(ByteOrder.BIG_ENDIAN)
-        }
-    }
-
-    companion object {
-        val size: Int
-            get() = Native.SIZE_T_SIZE
-
-        fun readFromBuffer(buf: ByteBuffer) : USize {
-            // Make sure we always read usize integers using native byte-order, since they may be
-            // casted from pointer values
-            buf.order(ByteOrder.nativeOrder())
-            try {
-                return when (Native.SIZE_T_SIZE) {
-                    4 -> USize(buf.getInt().toLong())
-                    8 -> USize(buf.getLong())
-                    else -> throw RuntimeException("Invalid SIZE_T_SIZE: ${Native.SIZE_T_SIZE}")
-                }
-            } finally {
-                buf.order(ByteOrder.BIG_ENDIAN)
-            }
-        }
-    }
-}
-
 internal inline fun<T> uniffiTraitInterfaceCall(
     callStatus: UniffiRustCallStatus,
     makeCall: () -> T,
@@ -153,39 +108,3 @@ internal inline fun<T, reified E: Throwable> uniffiTraitInterfaceCallWithError(
         }
     }
 }
-
-// Map handles to objects
-//
-// This is used when the Rust code expects an opaque pointer to represent some foreign object.
-// Normally we would pass a pointer to the object, but JNA doesn't support getting a pointer from an
-// object reference , nor does it support leaking a reference to Rust.
-//
-// Instead, this class maps USize values to objects so that we can pass a pointer-sized type to
-// Rust when it needs an opaque pointer.
-//
-// TODO: refactor callbacks to use this class
-internal class UniFfiHandleMap<T: Any> {
-    private val map = ConcurrentHashMap<USize, T>()
-    // Use AtomicInteger for our counter, since we may be on a 32-bit system.  4 billion possible
-    // values seems like enough. If somehow we generate 4 billion handles, then this will wrap
-    // around back to zero and we can assume the first handle generated will have been dropped by
-    // then.
-    private val counter = java.util.concurrent.atomic.AtomicInteger(0)
-
-    val size: Int
-        get() = map.size
-
-    fun insert(obj: T): USize {
-        val handle = USize(counter.getAndAdd(1).toLong())
-        map.put(handle, obj)
-        return handle
-    }
-
-    fun get(handle: USize): T? {
-        return map.get(handle)
-    }
-
-    fun remove(handle: USize): T? {
-        return map.remove(handle)
-    }
-}

uniffi_bindgen/src/bindings/kotlin/templates/ObjectTemplate.kt

@@ -95,7 +95,6 @@
 // [1] https://stackoverflow.com/questions/24376768/can-java-finalize-an-object-when-it-is-still-in-scope/24380219
 //
 
-{{ self.add_import("java.util.concurrent.atomic.AtomicLong") }}
 {{ self.add_import("java.util.concurrent.atomic.AtomicBoolean") }}
 {%- if self.include_once_check("interface-support") %}
     {%- include "ObjectCleanerHelper.kt" %}
@@ -312,7 +311,7 @@ open class {{ impl_class_name }}: Disposable, AutoCloseable, {{ interface_name }
 
 public object {{ obj|ffi_converter_name }}: FfiConverter<{{ type_name }}, Pointer> {
     {%- if obj.has_callback_interface() %}
-    internal val handleMap = ConcurrentHandleMap<{{ interface_name }}>()
+    internal val handleMap = UniffiHandleMap<{{ type_name }}>()
     {%- endif %}
 
     override fun lower(value: {{ type_name }}): Pointer {

uniffi_bindgen/src/bindings/kotlin/templates/wrapper.kt

@@ -30,6 +30,7 @@ import java.nio.ByteBuffer
 import java.nio.ByteOrder
 import java.nio.CharBuffer
 import java.nio.charset.CodingErrorAction
+import java.util.concurrent.atomic.AtomicLong
 import java.util.concurrent.ConcurrentHashMap
 
 {%- for req in self.imports() %}
@@ -39,6 +40,7 @@ import java.util.concurrent.ConcurrentHashMap
 {% include "RustBufferTemplate.kt" %}
 {% include "FfiConverterTemplate.kt" %}
 {% include "Helpers.kt" %}
+{% include "HandleMap.kt" %}
 
 // Contains loading, initialization code,
 // and the FFI Function declarations in a com.sun.jna.Library.

uniffi_bindgen/src/bindings/python/gen_python/mod.rs

@@ -372,7 +372,6 @@ impl PythonCodeOracle {
             // Pointer to an `asyncio.EventLoop` instance
             FfiType::Reference(inner) => format!("ctypes.POINTER({})", self.ffi_type_label(inner)),
             FfiType::VoidPointer | FfiType::RustFutureHandle => "ctypes.c_void_p".to_string(),
-            FfiType::RustFutureContinuationData => "ctypes.c_size_t".to_string(),
         }
     }
 

uniffi_bindgen/src/bindings/python/templates/Async.py

@@ -3,13 +3,13 @@
 _UNIFFI_RUST_FUTURE_POLL_MAYBE_READY = 1
 
 # Stores futures for _uniffi_continuation_callback
-_UniffiContinuationPointerManager = _UniffiPointerManager()
+_UniffiContinuationHandleMap = _UniffiHandleMap()
 
 # Continuation callback for async functions
 # lift the return value or error and resolve the future, causing the async function to resume.
 @UNIFFI_RUST_FUTURE_CONTINUATION_CALLBACK
 def _uniffi_continuation_callback(future_ptr, poll_code):
-    (eventloop, future) = _UniffiContinuationPointerManager.release_pointer(future_ptr)
+    (eventloop, future) = _UniffiContinuationHandleMap.remove(future_ptr)
     eventloop.call_soon_threadsafe(_uniffi_set_future_result, future, poll_code)
 
 def _uniffi_set_future_result(future, poll_code):
@@ -26,7 +26,7 @@ async def _uniffi_rust_call_async(rust_future, ffi_poll, ffi_complete, ffi_free,
             ffi_poll(
                 rust_future,
                 _uniffi_continuation_callback,
-                _UniffiContinuationPointerManager.new_pointer((eventloop, future)),
+                _UniffiContinuationHandleMap.insert((eventloop, future)),
             )
             poll_code = await future
             if poll_code == _UNIFFI_RUST_FUTURE_POLL_READY:

uniffi_bindgen/src/bindings/python/templates/CallbackInterfaceRuntime.py

@@ -1,38 +1,3 @@
-import threading
-
-class ConcurrentHandleMap:
-    """
-    A map where inserting, getting and removing data is synchronized with a lock.
-    """
-
-    def __init__(self):
-        # type Handle = int
-        self._left_map = {}  # type: Dict[Handle, Any]
-
-        self._lock = threading.Lock()
-        self._current_handle = 0
-        self._stride = 1
-
-    def insert(self, obj):
-        with self._lock:
-            handle = self._current_handle
-            self._current_handle += self._stride
-            self._left_map[handle] = obj
-            return handle
-
-    def get(self, handle):
-        with self._lock:
-            obj = self._left_map.get(handle)
-        if not obj:
-            raise InternalError("No callback in handlemap; this is a uniffi bug")
-        return obj
-
-    def remove(self, handle):
-        with self._lock:
-            if handle in self._left_map:
-                obj = self._left_map.pop(handle)
-                return obj
-
 # Magic number for the Rust proxy to call using the same mechanism as every other method,
 # to free the callback once it's dropped by Rust.
 IDX_CALLBACK_FREE = 0
@@ -42,7 +7,7 @@ def remove(self, handle):
 _UNIFFI_CALLBACK_UNEXPECTED_ERROR = 2
 
 class UniffiCallbackInterfaceFfiConverter:
-    _handle_map = ConcurrentHandleMap()
+    _handle_map = _UniffiHandleMap()
 
     @classmethod
     def lift(cls, handle):

uniffi_bindgen/src/bindings/python/templates/HandleMap.py

@@ -0,0 +1,31 @@
+class _UniffiHandleMap:
+    """
+    A map where inserting, getting and removing data is synchronized with a lock.
+    """
+
+    def __init__(self):
+        # type Handle = int
+        self._map = {}  # type: Dict[Handle, Any]
+        self._lock = threading.Lock()
+        self._counter = itertools.count()
+
+    def insert(self, obj):
+        with self._lock:
+            handle = next(self._counter)
+            self._map[handle] = obj
+            return handle
+
+    def get(self, handle):
+        try:
+            with self._lock:
+                return self._map[handle]
+        except KeyError:
+            raise InternalError("UniffiHandleMap.get: Invalid handle")
+
+    def remove(self, handle):
+        try:
+            with self._lock:
+                return self._map.pop(handle)
+        except KeyError:
+            raise InternalError("UniffiHandleMap.remove: Invalid handle")
+

uniffi_bindgen/src/bindings/python/templates/NamespaceLibraryTemplate.py

@@ -9,7 +9,7 @@ def _uniffi_future_callback_t(return_type):
     """
     Factory function to create callback function types for async functions
     """
-    return ctypes.CFUNCTYPE(None, ctypes.c_size_t, return_type, _UniffiRustCallStatus)
+    return ctypes.CFUNCTYPE(None, ctypes.c_uint64, return_type, _UniffiRustCallStatus)
 
 def _uniffi_load_indirect():
     """

uniffi_bindgen/src/bindings/python/templates/ObjectTemplate.py

@@ -84,7 +84,7 @@ def __ne__(self, other: object) -> {{ ne.return_type().unwrap()|type_name }}:
 
 class {{ ffi_converter_name }}:
     {%- if obj.has_callback_interface() %}
-    _handle_map = ConcurrentHandleMap()
+    _handle_map = _UniffiHandleMap()
     {%- endif %}
 
     @staticmethod