Clean up macros

macros/aggr_co_class/Cargo.toml

@@ -9,3 +9,4 @@ syn = { version = "1.0.5", features = ["full"] }
 quote = "1.0"
 proc-macro2 = "1.0.1"
 macro_utils = { path = "../macro_utils" }
+co_class = { path = "../co_class" }

macros/aggr_co_class/src/class_factory.rs

@@ -8,14 +8,14 @@ pub fn generate(struct_item: &ItemStruct) -> HelperTokenStream {
     let struct_ident = &struct_item.ident;
     let class_factory_ident = macro_utils::class_factory_ident(&struct_ident);
 
-    quote!(
-        // We are not going to bother with using an Init_ struct here,
-        // as we are not relying on the production macros.
-        #[repr(C)]
-        pub struct #class_factory_ident {
-            inner: <dyn com::IClassFactory as com::ComInterface>::VPtr,
-            ref_count: u32,
-        }
+    let struct_definition =
+        co_class::class_factory::gen_class_factory_struct_definition(&class_factory_ident);
+    let lock_server = co_class::class_factory::gen_lock_server();
+    let iunknown_impl = co_class::class_factory::gen_iunknown_impl(&class_factory_ident);
+    let class_factory_impl = co_class::class_factory::gen_class_factory_impl(&class_factory_ident);
+
+    quote! {
+        #struct_definition
 
         impl com::IClassFactory for #class_factory_ident {
             unsafe fn create_instance(
@@ -54,65 +54,11 @@ pub fn generate(struct_item: &ItemStruct) -> HelperTokenStream {
                 hr
             }
 
-            // TODO: Implement correctly
-            fn lock_server(&mut self, _increment: winapi::shared::minwindef::BOOL) -> winapi::shared::winerror::HRESULT {
-                println!("LockServer called");
-                winapi::shared::winerror::S_OK
-            }
+            #lock_server
         }
 
-        impl com::IUnknown for #class_factory_ident {
-            fn query_interface(&mut self, riid: *const winapi::shared::guiddef::IID, ppv: *mut *mut winapi::ctypes::c_void) -> winapi::shared::winerror::HRESULT {
-                // Bringing trait into scope to access add_ref method.
-                use com::IUnknown;
+        #iunknown_impl
 
-                unsafe {
-                    println!("Querying interface on {}...", stringify!(#class_factory_ident));
-
-                    let riid = &*riid;
-                    if winapi::shared::guiddef::IsEqualGUID(riid, &<dyn com::IUnknown as com::ComInterface>::IID) | winapi::shared::guiddef::IsEqualGUID(riid, &<dyn com::IClassFactory as com::ComInterface>::IID) {
-                        *ppv = &self.inner as *const _ as *mut winapi::ctypes::c_void;
-                        self.add_ref();
-                        winapi::shared::winerror::NOERROR
-                    } else {
-                        *ppv = std::ptr::null_mut::<winapi::ctypes::c_void>();
-                        winapi::shared::winerror::E_NOINTERFACE
-                    }
-                }
-            }
-
-            fn add_ref(&mut self) -> u32 {
-                self.ref_count += 1;
-                println!("Count now {}", self.ref_count);
-                self.ref_count
-            }
-
-            fn release(&mut self) -> u32 {
-                self.ref_count -= 1;
-                println!("Count now {}", self.ref_count);
-                let count = self.ref_count;
-                if count == 0 {
-                    println!("Count is 0 for {}. Freeing memory...", stringify!(#class_factory_ident));
-                    unsafe { Box::from_raw(self as *const _ as *mut #class_factory_ident); }
-                }
-                count
-            }
-        }
-
-        // Not using allocate function here, since we are not using macros.
-        impl #class_factory_ident {
-            pub(crate) fn new() -> Box<#class_factory_ident> {
-                use com::IClassFactory;
-
-                println!("Allocating new Vtable for {}...", stringify!(#class_factory_ident));
-                let class_vtable = com::vtable!(#class_factory_ident: IClassFactory);
-                let vptr = Box::into_raw(Box::new(class_vtable));
-                let class_factory = #class_factory_ident {
-                    inner: vptr,
-                    ref_count: 0,
-                };
-                Box::new(class_factory)
-            }
-        }
-    )
+        #class_factory_impl
+    }
 }

macros/aggr_co_class/src/iunknown_impl.rs

@@ -15,13 +15,13 @@ pub fn generate(struct_item: &ItemStruct) -> HelperTokenStream {
 
     quote!(
         impl com::IUnknown for #struct_ident {
-            fn query_interface(
+            unsafe fn query_interface(
                 &mut self,
                 riid: *const winapi::shared::guiddef::IID,
                 ppv: *mut *mut winapi::ctypes::c_void
             ) -> winapi::shared::winerror::HRESULT {
                 println!("Delegating QI");
-                let mut iunknown_to_use: com::ComPtr<dyn com::IUnknown> = unsafe { com::ComPtr::new(self.#iunknown_to_use_field_ident as *mut winapi::ctypes::c_void) };
+                let mut iunknown_to_use: com::ComPtr<dyn com::IUnknown> = com::ComPtr::new(self.#iunknown_to_use_field_ident as *mut winapi::ctypes::c_void);
                 let hr = iunknown_to_use.query_interface(riid, ppv);
                 core::mem::forget(iunknown_to_use);
 
@@ -36,8 +36,8 @@ pub fn generate(struct_item: &ItemStruct) -> HelperTokenStream {
                 res
             }
 
-            fn release(&mut self) -> u32 {
-                let mut iunknown_to_use: com::ComPtr<dyn com::IUnknown> = unsafe { com::ComPtr::new(self.#iunknown_to_use_field_ident as *mut winapi::ctypes::c_void) };
+            unsafe fn release(&mut self) -> u32 {
+                let mut iunknown_to_use: com::ComPtr<dyn com::IUnknown> = com::ComPtr::new(self.#iunknown_to_use_field_ident as *mut winapi::ctypes::c_void);
                 let res = iunknown_to_use.release();
                 core::mem::forget(iunknown_to_use);
 

macros/co_class/src/class_factory.rs

@@ -1,4 +1,4 @@
-use proc_macro2::TokenStream as HelperTokenStream;
+use proc_macro2::{Ident, TokenStream as HelperTokenStream};
 use quote::quote;
 use syn::ItemStruct;
 
@@ -8,14 +8,13 @@ pub fn generate(struct_item: &ItemStruct) -> HelperTokenStream {
     let struct_ident = &struct_item.ident;
     let class_factory_ident = macro_utils::class_factory_ident(&struct_ident);
 
-    quote!(
-        // We are not going to bother with using an Init_ struct here,
-        // as we are not relying on the production macros.
-        #[repr(C)]
-        pub struct #class_factory_ident {
-            inner: <dyn com::IClassFactory as com::ComInterface>::VPtr,
-            ref_count: u32,
-        }
+    let struct_definition = gen_class_factory_struct_definition(&class_factory_ident);
+    let lock_server = gen_lock_server();
+    let iunknown_impl = gen_iunknown_impl(&class_factory_ident);
+    let class_factory_impl = gen_class_factory_impl(&class_factory_ident);
+
+    quote! {
+        #struct_definition
 
         impl com::IClassFactory for #class_factory_ident {
             unsafe fn create_instance(
@@ -41,65 +40,87 @@ pub fn generate(struct_item: &ItemStruct) -> HelperTokenStream {
                 hr
             }
 
-            // TODO: Implement correctly
-            fn lock_server(&mut self, _increment: winapi::shared::minwindef::BOOL) -> winapi::shared::winerror::HRESULT {
-                println!("LockServer called");
-                winapi::shared::winerror::S_OK
-            }
+            #lock_server
         }
 
+        #iunknown_impl
+
+        #class_factory_impl
+    }
+}
+
+pub fn gen_class_factory_struct_definition(class_factory_ident: &Ident) -> HelperTokenStream {
+    let ref_count_ident = macro_utils::ref_count_ident();
+    quote! {
+        #[repr(C)]
+        pub struct #class_factory_ident {
+            inner: <dyn com::IClassFactory as com::ComInterface>::VPtr,
+            #ref_count_ident: u32,
+        }
+    }
+}
+
+pub fn gen_lock_server() -> HelperTokenStream {
+    quote! {
+        // TODO: Implement correctly
+        fn lock_server(&mut self, _increment: winapi::shared::minwindef::BOOL) -> winapi::shared::winerror::HRESULT {
+            println!("LockServer called");
+            winapi::shared::winerror::S_OK
+        }
+    }
+}
+
+pub fn gen_iunknown_impl(class_factory_ident: &Ident) -> HelperTokenStream {
+    let query_interface = gen_query_interface(class_factory_ident);
+    let add_ref = crate::iunknown_impl::gen_add_ref();
+    let release = crate::iunknown_impl::gen_release(class_factory_ident);
+    quote! {
         impl com::IUnknown for #class_factory_ident {
-            fn query_interface(&mut self, riid: *const winapi::shared::guiddef::IID, ppv: *mut *mut winapi::ctypes::c_void) -> winapi::shared::winerror::HRESULT {
-                // Bringing trait into scope to access add_ref method.
-                use com::IUnknown;
+            #query_interface
+            #add_ref
+            #release
+        }
+    }
+}
 
-                unsafe {
-                    println!("Querying interface on {}...", stringify!(#class_factory_ident));
+fn gen_query_interface(class_factory_ident: &Ident) -> HelperTokenStream {
+    quote! {
+        unsafe fn query_interface(&mut self, riid: *const winapi::shared::guiddef::IID, ppv: *mut *mut winapi::ctypes::c_void) -> winapi::shared::winerror::HRESULT {
+            // Bringing trait into scope to access add_ref method.
+            use com::IUnknown;
 
-                    let riid = &*riid;
-                    if winapi::shared::guiddef::IsEqualGUID(riid, &<dyn com::IUnknown as com::ComInterface>::IID) | winapi::shared::guiddef::IsEqualGUID(riid, &<dyn com::IClassFactory as com::ComInterface>::IID) {
-                        *ppv = &self.inner as *const _ as *mut winapi::ctypes::c_void;
-                        self.add_ref();
-                        winapi::shared::winerror::NOERROR
-                    } else {
-                        *ppv = std::ptr::null_mut::<winapi::ctypes::c_void>();
-                        winapi::shared::winerror::E_NOINTERFACE
-                    }
-                }
-            }
+            println!("Querying interface on {}...", stringify!(#class_factory_ident));
 
-            fn add_ref(&mut self) -> u32 {
-                self.ref_count += 1;
-                println!("Count now {}", self.ref_count);
-                self.ref_count
-            }
-
-            fn release(&mut self) -> u32 {
-                self.ref_count -= 1;
-                println!("Count now {}", self.ref_count);
-                let count = self.ref_count;
-                if count == 0 {
-                    println!("Count is 0 for {}. Freeing memory...", stringify!(#class_factory_ident));
-                    unsafe { Box::from_raw(self as *const _ as *mut #class_factory_ident); }
-                }
-                count
+            let riid = &*riid;
+            if winapi::shared::guiddef::IsEqualGUID(riid, &<dyn com::IUnknown as com::ComInterface>::IID) | winapi::shared::guiddef::IsEqualGUID(riid, &<dyn com::IClassFactory as com::ComInterface>::IID) {
+                *ppv = &self.inner as *const _ as *mut winapi::ctypes::c_void;
+                self.add_ref();
+                winapi::shared::winerror::NOERROR
+            } else {
+                *ppv = std::ptr::null_mut::<winapi::ctypes::c_void>();
+                winapi::shared::winerror::E_NOINTERFACE
             }
         }
+    }
+}
 
-        // Not using allocate function here, since we are not using macros.
+pub fn gen_class_factory_impl(class_factory_ident: &Ident) -> HelperTokenStream {
+    let ref_count_ident = macro_utils::ref_count_ident();
+    quote! {
         impl #class_factory_ident {
             pub(crate) fn new() -> Box<#class_factory_ident> {
                 use com::IClassFactory;
 
                 println!("Allocating new Vtable for {}...", stringify!(#class_factory_ident));
                 let class_vtable = com::vtable!(#class_factory_ident: IClassFactory);
+                // allocate directly since no macros generated an `allocate` function
                 let vptr = Box::into_raw(Box::new(class_vtable));
                 let class_factory = #class_factory_ident {
                     inner: vptr,
-                    ref_count: 0,
+                    #ref_count_ident: 0,
                 };
                 Box::new(class_factory)
             }
         }
-    )
+    }
 }

macros/co_class/src/iunknown_impl.rs

@@ -12,34 +12,78 @@ pub fn generate(
     struct_item: &ItemStruct,
 ) -> HelperTokenStream {
     let struct_ident = &struct_item.ident;
-    let ref_count_ident = macro_utils::ref_count_ident();
 
     let query_interface = gen_query_interface(base_interface_idents, aggr_interface_idents);
+    let add_ref = gen_add_ref();
+    let release = gen_release(struct_ident);
 
     quote!(
         impl com::IUnknown for #struct_ident {
-
             #query_interface
+            #add_ref
+            #release
+        }
+    )
+}
 
-            fn add_ref(&mut self) -> u32 {
-                self.#ref_count_ident += 1;
-                println!("Count now {}", self.#ref_count_ident);
-                self.#ref_count_ident
+pub fn gen_add_ref() -> HelperTokenStream {
+    let ref_count_ident = macro_utils::ref_count_ident();
+    quote! {
+        fn add_ref(&mut self) -> u32 {
+            self.#ref_count_ident = self.#ref_count_ident.checked_add(1).expect("Overflow of reference count");
+            println!("Count now {}", self.#ref_count_ident);
+            self.#ref_count_ident
+        }
+    }
+}
+
+pub fn gen_release(struct_ident: &Ident) -> HelperTokenStream {
+    let ref_count_ident = macro_utils::ref_count_ident();
+    quote! {
+        unsafe fn release(&mut self) -> u32 {
+            self.#ref_count_ident = self.#ref_count_ident.checked_sub(1).expect("Underflow of reference count");
+            println!("Count now {}", self.#ref_count_ident);
+            let count = self.#ref_count_ident;
+            if count == 0 {
+                println!("Count is 0 for {}. Freeing memory...", stringify!(#struct_ident));
+                Box::from_raw(self as *const _ as *mut #struct_ident);
+            }
+            count
+        }
+    }
+}
+
+fn gen_query_interface(
+    base_interface_idents: &[Ident],
+    aggr_interface_idents: &HashMap<Ident, Vec<Ident>>,
+) -> HelperTokenStream {
+    let first_vptr_field = macro_utils::vptr_field_ident(&base_interface_idents[0]);
+
+    // Generate match arms for implemented interfaces
+    let base_match_arms = gen_base_match_arms(base_interface_idents);
+
+    // Generate match arms for aggregated interfaces
+    let aggr_match_arms = gen_aggregate_match_arms(aggr_interface_idents);
+
+    quote!(
+        unsafe fn query_interface(
+            &mut self,
+            riid: *const winapi::shared::guiddef::IID,
+            ppv: *mut *mut winapi::ctypes::c_void
+        ) -> winapi::shared::winerror::HRESULT {
+            let riid = &*riid;
+
+            if winapi::shared::guiddef::IsEqualGUID(riid, &com::IID_IUNKNOWN) {
+                *ppv = &self.#first_vptr_field as *const _ as *mut winapi::ctypes::c_void;
+            } #base_match_arms #aggr_match_arms else {
+                *ppv = std::ptr::null_mut::<winapi::ctypes::c_void>();
+                println!("Returning NO INTERFACE.");
+                return winapi::shared::winerror::E_NOINTERFACE;
             }
 
-            fn release(&mut self) -> u32 {
-                self.#ref_count_ident -= 1;
-                println!("Count now {}", self.#ref_count_ident);
-                let count = self.#ref_count_ident;
-                if count == 0 {
-                    println!("Count is 0 for {}. Freeing memory...", stringify!(#struct_ident));
-                    // drop(self)
-                    // TODO: This doesn't free the original heap allocation, as &mut self results
-                    // in a copy from the raw pointer.
-                    unsafe { Box::from_raw(self as *const _ as *mut #struct_ident); }
-                }
-                count
-            }
+            println!("Successful!.");
+            self.add_ref();
+            NOERROR
         }
     )
 }
@@ -101,40 +145,3 @@ fn gen_aggregate_match_arms(
 
     quote!(#(#aggr_match_arms)*)
 }
-
-fn gen_query_interface(
-    base_interface_idents: &[Ident],
-    aggr_interface_idents: &HashMap<Ident, Vec<Ident>>,
-) -> HelperTokenStream {
-    let first_vptr_field = macro_utils::vptr_field_ident(&base_interface_idents[0]);
-
-    // Generate match arms for implemented interfaces
-    let base_match_arms = gen_base_match_arms(base_interface_idents);
-
-    // Generate match arms for aggregated interfaces
-    let aggr_match_arms = gen_aggregate_match_arms(aggr_interface_idents);
-
-    quote!(
-        fn query_interface(
-            &mut self,
-            riid: *const winapi::shared::guiddef::IID,
-            ppv: *mut *mut winapi::ctypes::c_void
-        ) -> winapi::shared::winerror::HRESULT {
-            unsafe {
-                let riid = &*riid;
-
-                if winapi::shared::guiddef::IsEqualGUID(riid, &com::IID_IUNKNOWN) {
-                    *ppv = &self.#first_vptr_field as *const _ as *mut winapi::ctypes::c_void;
-                } #base_match_arms #aggr_match_arms else {
-                    *ppv = std::ptr::null_mut::<winapi::ctypes::c_void>();
-                    println!("Returning NO INTERFACE.");
-                    return winapi::shared::winerror::E_NOINTERFACE;
-                }
-
-                println!("Successful!.");
-                self.add_ref();
-                NOERROR
-            }
-        }
-    )
-}

macros/co_class/src/lib.rs

@@ -4,11 +4,11 @@ use syn::ItemStruct;
 
 use std::iter::FromIterator;
 
-mod class_factory;
+pub mod class_factory;
 mod com_struct;
 mod com_struct_impl;
 mod drop_impl;
-mod iunknown_impl;
+pub mod iunknown_impl;
 
 // Macro expansion entry point.
 pub fn expand_co_class(_attr: TokenStream, item: TokenStream) -> TokenStream {

src/inproc.rs

@@ -181,16 +181,16 @@ macro_rules! com_inproc_dll_module {
             if $crate::_winapi::shared::guiddef::IsEqualGUID(rclsid, &$clsid_one) {
                 let mut instance = <$classtype_one>::get_class_object();
                 instance.add_ref();
-                let hr = instance.query_interface(riid, ppv);
-                instance.release();
+                let hr = unsafe { instance.query_interface(riid, ppv) };
+                unsafe { instance.release() };
                 Box::into_raw(instance);
 
                 hr
             } $(else if $crate::_winapi::shared::guiddef::IsEqualGUID(rclsid, &$clsid) {
                 let mut instance = <$classtype>::get_class_object();
                 instance.add_ref();
-                let hr = instance.query_interface(riid, ppv);
-                instance.release();
+                let hr = unsafe { instance.query_interface(riid, ppv) };
+                unsafe { instance.release() };
                 Box::into_raw(instance);
 
                 hr