1 files changed, 9 insertions, 8 deletions
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 }