Automatically generate MainThreadMarker argument in methods

crates/header-translator/src/cache.rs

@@ -1,4 +1,4 @@
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 use std::mem;
 
 use crate::config::Config;
@@ -7,16 +7,38 @@ use crate::id::ItemIdentifier;
 use crate::method::Method;
 use crate::output::Output;
 use crate::stmt::Stmt;
+use crate::Mutability;
 
 /// A helper struct for doing global analysis on the output.
 #[derive(Debug, PartialEq, Clone)]
 pub struct Cache<'a> {
     config: &'a Config,
+    mainthreadonly_classes: BTreeSet<ItemIdentifier>,
 }
 
 impl<'a> Cache<'a> {
-    pub fn new(_output: &Output, config: &'a Config) -> Self {
-        Self { config }
+    pub fn new(output: &Output, config: &'a Config) -> Self {
+        let mut mainthreadonly_classes = BTreeSet::new();
+
+        for library in output.libraries.values() {
+            for file in library.files.values() {
+                for stmt in file.stmts.iter() {
+                    if let Stmt::ClassDecl {
+                        id,
+                        mutability: Mutability::MainThreadOnly,
+                        ..
+                    } = stmt
+                    {
+                        mainthreadonly_classes.insert(id.clone());
+                    }
+                }
+            }
+        }
+
+        Self {
+            config,
+            mainthreadonly_classes,
+        }
     }
 
     pub fn update(&self, output: &mut Output) {
@@ -68,6 +90,67 @@ impl<'a> Cache<'a> {
             }
         }
 
+        // Add `mainthreadonly` to relevant methods
+        for stmt in file.stmts.iter_mut() {
+            match stmt {
+                Stmt::Methods {
+                    cls: id, methods, ..
+                }
+                | Stmt::ProtocolDecl { id, methods, .. } => {
+                    for method in methods.iter_mut() {
+                        let mut result_type_contains_mainthreadonly: bool = false;
+                        method.result_type.visit_required_types(&mut |id| {
+                            if self.mainthreadonly_classes.contains(id) {
+                                result_type_contains_mainthreadonly = true;
+                            }
+                        });
+
+                        match (method.is_class, self.mainthreadonly_classes.contains(id)) {
+                            // MainThreadOnly class with static method
+                            (true, true) => {
+                                // Assume the method needs main thread
+                                result_type_contains_mainthreadonly = true;
+                            }
+                            // Class with static method
+                            (true, false) => {
+                                // Continue with the normal check
+                            }
+                            // MainThreadOnly class with non-static method
+                            (false, true) => {
+                                // Method is already required to run on main
+                                // thread, so no need to add MainThreadMarker
+                                continue;
+                            }
+                            // Class with non-static method
+                            (false, false) => {
+                                // Continue with the normal check
+                            }
+                        }
+
+                        if result_type_contains_mainthreadonly {
+                            let mut any_argument_contains_mainthreadonly: bool = false;
+                            for (_, argument) in method.arguments.iter() {
+                                // Important: We only visit the top-level types, to not
+                                // include e.g. `Option<&NSView>` or `&NSArray<NSView>`.
+                                argument.visit_toplevel_types(&mut |id| {
+                                    if self.mainthreadonly_classes.contains(id) {
+                                        any_argument_contains_mainthreadonly = true;
+                                    }
+                                });
+                            }
+
+                            // Apply main thread only, unless a (required)
+                            // argument was main thread only.
+                            if !any_argument_contains_mainthreadonly {
+                                method.mainthreadonly = true;
+                            }
+                        }
+                    }
+                }
+                _ => {}
+            }
+        }
+
         // Fix up a few typedef + enum declarations
         let mut iter = mem::take(&mut file.stmts).into_iter().peekable();
         while let Some(stmt) = iter.next() {

crates/header-translator/src/method.rs

@@ -249,14 +249,14 @@ pub struct Method {
     pub is_class: bool,
     is_optional_protocol: bool,
     memory_management: MemoryManagement,
-    arguments: Vec<(String, Ty)>,
+    pub(crate) arguments: Vec<(String, Ty)>,
     pub result_type: Ty,
     safe: bool,
     mutating: bool,
     is_protocol: bool,
     // Thread-safe, even on main-thread only (@MainActor/@UIActor) classes
     non_isolated: bool,
-    mainthreadonly: bool,
+    pub(crate) mainthreadonly: bool,
 }
 
 impl Method {
@@ -636,6 +636,11 @@ impl fmt::Display for Method {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let _span = debug_span!("method", self.fn_name).entered();
 
+        // TODO: Use this somehow?
+        // if self.non_isolated {
+        //     writeln!(f, "// non_isolated")?;
+        // }
+
         //
         // Attributes
         //
@@ -689,7 +694,11 @@ impl fmt::Display for Method {
         // Arguments
         for (param, arg_ty) in &self.arguments {
             let param = handle_reserved(&crate::to_snake_case(param));
-            write!(f, "{param}: {arg_ty},")?;
+            write!(f, "{param}: {arg_ty}, ")?;
+        }
+        // FIXME: Skipping main thread only on protocols for now
+        if self.mainthreadonly && !self.is_protocol {
+            write!(f, "mtm: MainThreadMarker")?;
         }
         write!(f, ")")?;
 

crates/header-translator/src/rust_type.rs

@@ -345,6 +345,12 @@ impl IdType {
             }
         }
     }
+
+    fn visit_toplevel_types(&self, f: &mut impl FnMut(&ItemIdentifier)) {
+        if let Some(id) = self._id() {
+            f(id);
+        }
+    }
 }
 
 impl fmt::Display for IdType {
@@ -1036,6 +1042,25 @@ impl Inner {
             _ => {}
         }
     }
+
+    pub fn visit_toplevel_types(&self, f: &mut impl FnMut(&ItemIdentifier)) {
+        match self {
+            Self::Id { ty, .. } => {
+                ty.visit_toplevel_types(f);
+            }
+            Self::Pointer {
+                // Only visit non-null types
+                nullability: Nullability::NonNull,
+                is_const: _,
+                pointee,
+            } => {
+                pointee.visit_toplevel_types(f);
+            }
+            // TODO
+            Self::TypeDef { id } => f(id),
+            _ => {}
+        }
+    }
 }
 
 /// This is sound to output in (almost, c_void is not a valid return type) any
@@ -1492,6 +1517,14 @@ impl Ty {
 
         self.ty.visit_required_types(f);
     }
+
+    pub fn visit_toplevel_types(&self, f: &mut impl FnMut(&ItemIdentifier)) {
+        if let TyKind::MethodReturn { with_error: true } = &self.kind {
+            f(&ItemIdentifier::nserror());
+        }
+
+        self.ty.visit_toplevel_types(f);
+    }
 }
 
 impl Ty {

crates/icrate/Cargo.toml

@@ -131,8 +131,9 @@ unstable-example-basic_usage = [
 unstable-example-delegate = [
     "apple",
     "Foundation",
-    "Foundation_NSString",
     "Foundation_NSNotification",
+    "Foundation_NSString",
+    "Foundation_NSThread",
     "AppKit",
     "AppKit_NSApplication",
 ]
@@ -165,6 +166,7 @@ unstable-example-browser = [
     "Foundation",
     "Foundation_NSNotification",
     "Foundation_NSString",
+    "Foundation_NSThread",
     "Foundation_NSURL",
     "Foundation_NSURLRequest",
     "WebKit",
@@ -177,10 +179,11 @@ unstable-example-metal = [
     "AppKit_NSWindow",
     "Foundation",
     "Foundation_NSCoder",
+    "Foundation_NSDate",
     "Foundation_NSError",
     "Foundation_NSNotification",
     "Foundation_NSString",
-    "Foundation_NSDate",
+    "Foundation_NSThread",
     "Metal",
     "Metal_MTLCompileOptions",
     "Metal_MTLRenderPassDescriptor",

crates/icrate/examples/browser.rs

@@ -12,15 +12,15 @@ use icrate::{
         NSWindowStyleMaskResizable, NSWindowStyleMaskTitled,
     },
     Foundation::{
-        ns_string, NSNotification, NSObject, NSObjectProtocol, NSPoint, NSRect, NSSize,
-        NSURLRequest, NSURL,
+        ns_string, MainThreadMarker, NSNotification, NSObject, NSObjectProtocol, NSPoint, NSRect,
+        NSSize, NSURLRequest, NSURL,
     },
     WebKit::{WKNavigation, WKNavigationDelegate, WKWebView},
 };
 use objc2::{
     declare::{Ivar, IvarDrop},
     declare_class, msg_send, msg_send_id,
-    mutability::InteriorMutable,
+    mutability::MainThreadOnly,
     rc::Id,
     runtime::{AnyObject, ProtocolObject, Sel},
     sel, ClassType,
@@ -47,7 +47,7 @@ declare_class!(
 
     unsafe impl ClassType for Delegate {
         type Super = NSObject;
-        type Mutability = InteriorMutable;
+        type Mutability = MainThreadOnly;
         const NAME: &'static str = "Delegate";
     }
 
@@ -68,9 +68,9 @@ declare_class!(
         #[method(applicationDidFinishLaunching:)]
         #[allow(non_snake_case)]
         unsafe fn applicationDidFinishLaunching(&self, _notification: &NSNotification) {
+            let mtm = MainThreadMarker::from(self);
             // create the app window
             let window = {
-                let this = NSWindow::alloc();
                 let content_rect = NSRect::new(NSPoint::new(0., 0.), NSSize::new(1024., 768.));
                 let style = NSWindowStyleMaskClosable
                     | NSWindowStyleMaskResizable
@@ -79,7 +79,7 @@ declare_class!(
                 let flag = false;
                 unsafe {
                     NSWindow::initWithContentRect_styleMask_backing_defer(
-                        this,
+                        mtm.alloc(),
                         content_rect,
                         style,
                         backing_store_type,
@@ -90,16 +90,14 @@ declare_class!(
 
             // create the web view
             let web_view = {
-                let this = WKWebView::alloc();
                 let frame_rect = NSRect::ZERO;
-                unsafe { WKWebView::initWithFrame(this, frame_rect) }
+                unsafe { WKWebView::initWithFrame(mtm.alloc(), frame_rect) }
             };
 
             // create the nav bar view
             let nav_bar = {
                 let frame_rect = NSRect::ZERO;
-                let this = NSStackView::alloc();
-                let this = unsafe { NSStackView::initWithFrame(this, frame_rect) };
+                let this = unsafe { NSStackView::initWithFrame(mtm.alloc(), frame_rect) };
                 unsafe {
                     this.setOrientation(NSUserInterfaceLayoutOrientationHorizontal);
                     this.setAlignment(NSLayoutAttributeHeight);
@@ -112,8 +110,7 @@ declare_class!(
             // create the nav buttons view
             let nav_buttons = {
                 let frame_rect = NSRect::ZERO;
-                let this = NSStackView::alloc();
-                let this = unsafe { NSStackView::initWithFrame(this, frame_rect) };
+                let this = unsafe { NSStackView::initWithFrame(mtm.alloc(), frame_rect) };
                 unsafe {
                     this.setOrientation(NSUserInterfaceLayoutOrientationHorizontal);
                     this.setAlignment(NSLayoutAttributeHeight);
@@ -130,7 +127,7 @@ declare_class!(
                 let target = Some::<&AnyObject>(&web_view);
                 let action = Some(sel!(goBack));
                 let this =
-                    unsafe { NSButton::buttonWithTitle_target_action(title, target, action) };
+                    unsafe { NSButton::buttonWithTitle_target_action(title, target, action, mtm) };
                 unsafe { this.setBezelStyle(NSBezelStyleShadowlessSquare) };
                 this
             };
@@ -142,7 +139,7 @@ declare_class!(
                 let target = Some::<&AnyObject>(&web_view);
                 let action = Some(sel!(goForward));
                 let this =
-                    unsafe { NSButton::buttonWithTitle_target_action(title, target, action) };
+                    unsafe { NSButton::buttonWithTitle_target_action(title, target, action, mtm) };
                 unsafe { this.setBezelStyle(NSBezelStyleShadowlessSquare) };
                 this
             };
@@ -155,8 +152,7 @@ declare_class!(
             // create the url text field
             let nav_url = {
                 let frame_rect = NSRect::ZERO;
-                let this = NSTextField::alloc();
-                let this = unsafe { NSTextField::initWithFrame(this, frame_rect) };
+                let this = unsafe { NSTextField::initWithFrame(mtm.alloc(), frame_rect) };
                 unsafe {
                     this.setDrawsBackground(true);
                     this.setBackgroundColor(Some(&NSColor::lightGrayColor()));
@@ -173,8 +169,7 @@ declare_class!(
             // create the window content view
             let content_view = {
                 let frame_rect = unsafe { window.frame() };
-                let this = NSStackView::alloc();
-                let this = unsafe { NSStackView::initWithFrame(this, frame_rect) };
+                let this = unsafe { NSStackView::initWithFrame(mtm.alloc(), frame_rect) };
                 unsafe {
                     this.setOrientation(NSUserInterfaceLayoutOrientationVertical);
                     this.setAlignment(NSLayoutAttributeWidth);
@@ -217,7 +212,7 @@ declare_class!(
                 menu_app_item.setSubmenu(Some(&menu_app_menu));
                 menu.addItem(&menu_app_item);
 
-                let app = NSApplication::sharedApplication();
+                let app = NSApplication::sharedApplication(mtm);
                 app.setMainMenu(Some(&menu));
             }
 
@@ -286,19 +281,20 @@ declare_class!(
 );
 
 impl Delegate {
-    pub fn new() -> Id<Self> {
-        unsafe { msg_send_id![Self::alloc(), init] }
+    pub fn new(mtm: MainThreadMarker) -> Id<Self> {
+        unsafe { msg_send_id![mtm.alloc(), init] }
     }
 }
 
 unsafe impl NSObjectProtocol for Delegate {}
 
 fn main() {
-    let app = unsafe { NSApplication::sharedApplication() };
+    let mtm = MainThreadMarker::new().unwrap();
+    let app = unsafe { NSApplication::sharedApplication(mtm) };
     unsafe { app.setActivationPolicy(NSApplicationActivationPolicyRegular) };
 
     // initialize the delegate
-    let delegate = Delegate::new();
+    let delegate = Delegate::new(mtm);
 
     // configure the application delegate
     unsafe {