Add support for weak references (#745)

crates/gen/src/types/delegate.rs

@@ -89,7 +89,7 @@ impl Delegate {
                 pub fn new<#fn_constraint>(invoke: F) -> Self {
                     let com = #box_name {
                         vtable: &#box_name::VTABLE,
-                        count: ::windows::RefCount::new(),
+                        count: ::windows::RefCount::new(1),
                         invoke,
                     };
                     unsafe {
@@ -140,6 +140,8 @@ impl Delegate {
                             ::std::ptr::null_mut()
                         };
 
+                    // TODO: implement IMarshal
+
                     if (*interface).is_null() {
                         ::windows::HRESULT(0x8000_4002) // E_NOINTERFACE
                     } else {

crates/macros/src/implement.rs

@@ -195,7 +195,7 @@ pub fn gen(
         struct #box_ident {
             vtable: (#(*const #vtable_idents,)*),
             inner: #inner_ident,
-            count: ::windows::RefCount,
+            count: ::windows::WeakRefCount,
         }
         impl #box_ident {
             const VTABLE: (#(#vtable_idents,)*) = (
@@ -205,7 +205,7 @@ pub fn gen(
                 Self {
                     vtable: (#(&Self::VTABLE.#vtable_ordinals,)*),
                     inner,
-                    count: ::windows::RefCount::new()
+                    count: ::windows::WeakRefCount::new()
                 }
             }
             fn QueryInterface(&mut self, iid: &::windows::Guid, interface: *mut ::windows::RawPtr) -> ::windows::HRESULT {
@@ -220,10 +220,16 @@ pub fn gen(
                         _ => ::std::ptr::null_mut(),
                     };
 
+                    if !(*interface).is_null() {
+                        self.count.add_ref();
+                        return ::windows::HRESULT(0);
+                    }
+
+                    *interface = self.count.query(iid, &mut self.vtable.0 as *mut _ as _);
+
                     if (*interface).is_null() {
                         ::windows::HRESULT(0x8000_4002) // E_NOINTERFACE
                     } else {
-                        self.count.add_ref();
                         ::windows::HRESULT(0)
                     }
                 }

src/interfaces/unknown.rs

@@ -8,14 +8,14 @@ use crate::*;
 pub struct IUnknown(std::ptr::NonNull<std::ffi::c_void>);
 
 #[repr(C)]
-pub struct IUnknown_vtable(
-    pub unsafe extern "system" fn(this: RawPtr, iid: &Guid, interface: *mut RawPtr) -> HRESULT,
+pub struct IUnknown_abi(
+    pub unsafe extern "system" fn(this: RawPtr, iid: *const Guid, interface: *mut RawPtr) -> HRESULT,
     pub unsafe extern "system" fn(this: RawPtr) -> u32,
     pub unsafe extern "system" fn(this: RawPtr) -> u32,
 );
 
 unsafe impl Interface for IUnknown {
-    type Vtable = IUnknown_vtable;
+    type Vtable = IUnknown_abi;
 
     const IID: Guid = Guid::from_values(
         0x0000_0000,

src/lib.rs

@@ -17,14 +17,20 @@ mod result;
 mod runtime;
 mod traits;
 
+use interfaces::*;
 use runtime::*;
 
+#[doc(hidden)]
 pub use bindings::Windows::Win32::Com::IAgileObject;
-pub use interfaces::{IActivationFactory, IUnknown, Object};
+
+#[doc(hidden)]
+pub use interfaces::IActivationFactory;
+
+pub use interfaces::{IUnknown, Object};
 pub use result::{Error, Result, HRESULT};
 pub use runtime::{
     create_instance, factory, initialize_mta, initialize_sta, Array, FactoryCache, Guid, HString,
-    Param, RefCount, Waiter, Weak,
+    Param, RefCount, Waiter, Weak, WeakRefCount,
 };
 pub use traits::{Abi, Interface, IntoParam, RuntimeName, RuntimeType};
 

src/runtime/hstring.rs

@@ -247,7 +247,7 @@ impl Header {
             (*header).flags = 0;
             (*header).len = len;
             (*header).data = &mut (*(*header).shared.as_mut_ptr()).buffer_start;
-            (*(*header).shared.as_mut_ptr()).count = RefCount::new();
+            (*(*header).shared.as_mut_ptr()).count = RefCount::new(1);
         }
         header
     }

src/runtime/mod.rs

@@ -9,6 +9,7 @@ mod param;
 mod ref_count;
 mod waiter;
 mod weak;
+mod weak_ref_count;
 
 pub use array::*;
 pub use com::*;
@@ -21,3 +22,4 @@ pub use param::*;
 pub use ref_count::*;
 pub use waiter::*;
 pub use weak::*;
+pub use weak_ref_count::*;

src/runtime/ref_count.rs

@@ -1,34 +1,30 @@
-use std::sync::atomic::{self, AtomicI32, Ordering};
+use std::sync::atomic::{fence, AtomicI32, Ordering};
 
-/// A thread-safe reference count for use with COM implementations.
+/// A thread-safe reference count for use with COM/HSTRING implementations.
 #[repr(transparent)]
 #[derive(Default)]
-pub struct RefCount {
-    value: AtomicI32,
-}
+pub struct RefCount(pub(crate) AtomicI32);
 
 impl RefCount {
     /// Creates a new `RefCount` with an initial value of `1`.
-    pub fn new() -> RefCount {
-        RefCount {
-            value: AtomicI32::new(1),
-        }
+    pub fn new(count: u32) -> Self {
+        Self(AtomicI32::new(count as _))
     }
 
     /// Increments the reference count, returning the new value.
     pub fn add_ref(&self) -> u32 {
-        (self.value.fetch_add(1, Ordering::Relaxed) + 1) as u32
+        (self.0.fetch_add(1, Ordering::Relaxed) + 1) as u32
     }
 
     /// Decrements the reference count, returning the new value.
     ///
     /// This operation inserts an `Acquire` fence when the reference count reaches zero.
     /// This prevents reordering before the object is destroyed.
     pub fn release(&self) -> u32 {
-        let remaining = self.value.fetch_sub(1, Ordering::Release) - 1;
+        let remaining = self.0.fetch_sub(1, Ordering::Release) - 1;
 
         if remaining == 0 {
-            atomic::fence(Ordering::Acquire);
+            fence(Ordering::Acquire);
         } else if remaining < 0 {
             panic!("Object has been over-released.");
         }

src/runtime/weak.rs

@@ -15,6 +15,7 @@ impl<I: Interface> Weak<I> {
     /// Attempts to upgrade the weak reference to a strong reference.
     pub fn upgrade(&self) -> Option<I> {
         self.0.as_ref().and_then(|inner| unsafe {
+            // Calls IWeakReference's Resolve ABI directly as the metadata incorrectly types it as returning IInspectable.
             let mut result = None;
             let _ = (inner.vtable().3)(inner.abi(), &I::IID, &mut result as *mut _ as _);
             result