5 files changed, 24 insertions, 10 deletions
diff --git a/lib/CWayland/CWayland.cc b/lib/CWayland/CWayland.cc
index 3e6bd0e..cc9a9de 100644
--- a/lib/CWayland/CWayland.cc
+++ b/lib/CWayland/CWayland.cc
@@ -174,8 +174,6 @@ lean_obj_res waylean_display_new(b_lean_obj_arg name)
     if (!disp) { return io_result_mk_from_errno(name_cstr ? name_cstr : "<default>"); }
     lean_object* waylean_disp = lean_alloc_external(waylean_display_class(), disp);
 
-    // wl_proxy_add_dispatcher((wl_proxy*)disp, dispatcher_func, waylean_disp, nullptr);
-
     return lean_io_result_mk_ok(waylean_disp);
 }
 
@@ -1271,4 +1269,11 @@ lean_obj_res waylean_proxy_listen(b_lean_obj_arg proxy_obj, lean_obj_arg handler
     if (old_data) { lean_dec(old_data); }
     wl_proxy_set_user_data(proxy, handler);
     return lean_io_result_mk_ok(lean_box(0));
+}
+
+// Proxy.keep : (proxy : @& Proxy) → IO Unit
+lean_obj_res waylean_proxy_keep(b_lean_obj_arg proxy)
+{
+    (void)proxy;
+    return lean_io_result_mk_ok(lean_box(0));
 }
 \ No newline at end of file
diff --git a/lib/CWayland/CWayland.h b/lib/CWayland/CWayland.h
index 44424a2..1b8e0c8 100644
--- a/lib/CWayland/CWayland.h
+++ b/lib/CWayland/CWayland.h
@@ -80,6 +80,7 @@ lean_obj_res waylean_proxy_send_constructor(b_lean_obj_arg proxy,
                                             b_lean_obj_arg iface_name);
 lean_obj_res waylean_proxy_listen          (b_lean_obj_arg proxy,
                                             lean_obj_arg   handler);
+lean_obj_res waylean_proxy_keep            (b_lean_obj_arg proxy);
 
 #ifdef __cplusplus
 }
diff --git a/lib/Wayland/Display.lean b/lib/Wayland/Display.lean
index f07c5ff..0a049b3 100644
--- a/lib/Wayland/Display.lean
+++ b/lib/Wayland/Display.lean
@@ -317,6 +317,10 @@ def Object.listen {ι : InterfaceName} [iface : Interface ι] (obj : Object ι)
     ReaderT.run (handler event) global
   global.listen obj.proxy actualHandler
 
+/-- Shorthand to ensure the reference count of an `Object` does not drop to `0`. -/
+def Object.keep {ι : InterfaceName} (obj : Object ι) : DisplayM σ Unit :=
+  obj.proxy.keep
+
 end Proxies
 
 end Wayland
diff --git a/lib/Wayland/Native.lean b/lib/Wayland/Native.lean
index dd138e4..f8d4c53 100644
--- a/lib/Wayland/Native.lean
+++ b/lib/Wayland/Native.lean
@@ -136,4 +136,7 @@ opaque Proxy.sendConstructor : (proxy : @& Proxy) → (opcode : UInt32) → (arg
 @[extern "waylean_proxy_listen"]
 opaque Proxy.listen : (proxy : @& Proxy) → (handler : UInt32 → ArgumentBuffer → EventLoop M × M) → IO Unit
 
+@[extern "waylean_proxy_keep"]
+opaque Proxy.keep : (proxy : @& Proxy) → IO Unit
+
 end Wayland.Native
diff --git a/src/Waylean.lean b/src/Waylean.lean
index 3f020a1..b595672 100644
--- a/src/Waylean.lean
+++ b/src/Waylean.lean
@@ -57,13 +57,16 @@ instance : ResolvedInterface WLDisplay.name where
 instance : ResolvedInterface WLCallback.name where
   referencedInterfaces := #[⟨WLCallback.name, inferInstance⟩]
 
-def WLDisplay.sync {σ : Type} (display : WLDisplay) : DisplayM σ WLCallback :=
-  display.sendConstructor do
+def WLDisplay.sync {σ : Type} (display : WLDisplay) (handler : DisplayM σ Unit) : DisplayM σ Unit := do
+  let callback : WLCallback ← display.sendConstructor do
     return 0
+  callback.listen fun
+  | .done _ => do
+      callback.keep
+      handler
 
 structure State where
   count : Int
-  callback : Option WLCallback
 
 partial def nextTimer : DisplayM Int Unit := do
   _ ← DisplayM.registerTimer 16ms do
@@ -79,10 +82,8 @@ def launch : DisplayM State Unit := do
   DisplayM.zoom (mk_lens count) nextTimer
 
   let disp : WLDisplay ← DisplayM.getDisplay
-  let callback : WLCallback ← disp.sync
-  callback.listen fun ev ↦ match ev with
-    | .done _ => DisplayM.liftIO (IO.println "Roundtrip complete")
-  set { (← get) with callback := callback }
+  disp.sync do
+    DisplayM.liftIO (IO.println "Roundtrip complete")
 
 def main : IO Unit :=
-  DisplayM.run launch { count := 0, callback := .none }
+  DisplayM.run launch { count := 0 }