1 files changed, 78 insertions, 49 deletions

diff --git a/src/codegen/mod.rs b/src/codegen/mod.rs
index 9a3b10d7..2be9794b 100644
--- a/src/codegen/mod.rs
+++ b/src/codegen/mod.rs
@@ -3394,9 +3394,12 @@ impl TryToRustTy for Type {
             TypeKind::ObjCSel => Ok(quote! {
                 objc::runtime::Sel
             }),
-            TypeKind::ObjCId | TypeKind::ObjCInterface(..) => Ok(quote! {
+            TypeKind::ObjCId => Ok(quote! {
                 id
             }),
+            TypeKind::ObjCInterface(..) => Ok(quote! {
+                objc::runtime::Object
+            }),
             ref u @ TypeKind::UnresolvedTypeRef(..) => {
                 unreachable!("Should have been resolved after parsing {:?}!", u)
             }
@@ -3644,7 +3647,7 @@ fn objc_method_codegen(
     method: &ObjCMethod,
     class_name: Option<&str>,
     prefix: &str,
-) -> (proc_macro2::TokenStream, proc_macro2::TokenStream) {
+) -> proc_macro2::TokenStream {
     let signature = method.signature();
     let fn_args = utils::fnsig_arguments(ctx, signature);
     let fn_ret = utils::fnsig_return_ty(ctx, signature);
@@ -3665,15 +3668,13 @@ fn objc_method_codegen(
     let methods_and_args = method.format_method_call(&fn_args);
 
     let body = if method.is_class_method() {
-        let class_name = class_name
-            .expect("Generating a class method without class name?")
-            .to_owned();
-        let expect_msg = proc_macro2::Literal::string(&format!(
-            "Couldn't find {}",
+        let class_name = ctx.rust_ident(
             class_name
-        ));
+                .expect("Generating a class method without class name?")
+                .to_owned(),
+        );
         quote! {
-            msg_send!(objc::runtime::Class::get(#class_name).expect(#expect_msg), #methods_and_args)
+            msg_send!(class!(#class_name), #methods_and_args)
         }
     } else {
         quote! {
@@ -3684,16 +3685,11 @@ fn objc_method_codegen(
     let method_name =
         ctx.rust_ident(format!("{}{}", prefix, method.rust_name()));
 
-    (
-        quote! {
-            unsafe fn #method_name #sig {
-                #body
-            }
-        },
-        quote! {
-            unsafe fn #method_name #sig ;
-        },
-    )
+    quote! {
+        unsafe fn #method_name #sig where <Self as std::ops::Deref>::Target: objc::Message + Sized {
+            #body
+        }
+    }
 }
 
 impl CodeGenerator for ObjCInterface {
@@ -3708,13 +3704,10 @@ impl CodeGenerator for ObjCInterface {
         debug_assert!(item.is_enabled_for_codegen(ctx));
 
         let mut impl_items = vec![];
-        let mut trait_items = vec![];
 
         for method in self.methods() {
-            let (impl_item, trait_item) =
-                objc_method_codegen(ctx, method, None, "");
+            let impl_item = objc_method_codegen(ctx, method, None, "");
             impl_items.push(impl_item);
-            trait_items.push(trait_item)
         }
 
         let instance_method_names: Vec<_> =
@@ -3724,61 +3717,97 @@ impl CodeGenerator for ObjCInterface {
             let ambiquity =
                 instance_method_names.contains(&class_method.rust_name());
             let prefix = if ambiquity { "class_" } else { "" };
-            let (impl_item, trait_item) = objc_method_codegen(
+            let impl_item = objc_method_codegen(
                 ctx,
                 class_method,
                 Some(self.name()),
                 prefix,
             );
             impl_items.push(impl_item);
-            trait_items.push(trait_item)
         }
 
         let trait_name = ctx.rust_ident(self.rust_name());
-
+        let trait_constraints = quote! {
+            Sized + std::ops::Deref
+        };
         let trait_block = if self.is_template() {
             let template_names: Vec<Ident> = self
                 .template_names
                 .iter()
                 .map(|g| ctx.rust_ident(g))
                 .collect();
+
             quote! {
-                pub trait #trait_name <#(#template_names),*>{
-                    #( #trait_items )*
+                pub trait #trait_name <#(#template_names),*> : #trait_constraints {
+                    #( #impl_items )*
                 }
             }
         } else {
             quote! {
-                pub trait #trait_name {
-                    #( #trait_items )*
+                pub trait #trait_name : #trait_constraints {
+                    #( #impl_items )*
                 }
             }
         };
 
-        let ty_for_impl = quote! {
-            id
-        };
-        let impl_block = if self.is_template() {
-            let template_names: Vec<Ident> = self
-                .template_names
-                .iter()
-                .map(|g| ctx.rust_ident(g))
-                .collect();
-            quote! {
-                impl <#(#template_names :'static),*> #trait_name <#(#template_names),*> for #ty_for_impl {
-                    #( #impl_items )*
+        let class_name = ctx.rust_ident(self.name());
+        if !self.is_category() && !self.is_protocol() {
+            let struct_block = quote! {
+                #[repr(transparent)]
+                #[derive(Clone, Copy)]
+                pub struct #class_name(pub id);
+                impl std::ops::Deref for #class_name {
+                    type Target = objc::runtime::Object;
+                    fn deref(&self) -> &Self::Target {
+                        unsafe {
+                            &*self.0
+                        }
+                    }
                 }
-            }
-        } else {
-            quote! {
-                impl #trait_name for #ty_for_impl {
-                    #( #impl_items )*
+                unsafe impl objc::Message for #class_name { }
+                impl #class_name {
+                    pub fn alloc() -> Self {
+                        Self(unsafe {
+                            msg_send!(objc::class!(#class_name), alloc)
+                        })
+                    }
                 }
+            };
+            result.push(struct_block);
+            for protocol_id in self.conforms_to.iter() {
+                let protocol_name = ctx.rust_ident(
+                    ctx.resolve_type(protocol_id.expect_type_id(ctx))
+                        .name()
+                        .unwrap(),
+                );
+                let impl_trait = quote! {
+                    impl #protocol_name for #class_name { }
+                };
+                result.push(impl_trait);
             }
-        };
+        }
+
+        if !self.is_protocol() {
+            let impl_block = if self.is_template() {
+                let template_names: Vec<Ident> = self
+                    .template_names
+                    .iter()
+                    .map(|g| ctx.rust_ident(g))
+                    .collect();
+                quote! {
+                    impl <#(#template_names :'static),*> #trait_name <#(#template_names),*> for #class_name {
+                    }
+                }
+            } else {
+                quote! {
+                    impl #trait_name for #class_name {
+                    }
+                }
+            };
+            result.push(impl_block);
+        }
 
         result.push(trait_block);
-        result.push(impl_block);
         result.saw_objc();
     }
 }