some initial cleanup

1 files changed, 64 insertions, 83 deletions

diff --git a/src/Waylean.lean b/src/Waylean.lean
index dd43179..09950ab 100644
--- a/src/Waylean.lean
+++ b/src/Waylean.lean
@@ -4,25 +4,11 @@ import Std.Data.HashMap
 
 open Wayland Protocols Waylean
 
-def Array.findAndModifyM {α : Type} {m : Type → Type} [Monad m] (a : Array α) (pred : α → Bool) (update : α → m α) : m (Array α × Option α) := do
-  let idx := a.findIdx pred
-  if h : idx < a.size then
-    let new ← update a[idx]
-    return (a.setIfInBounds idx new, new)
-  else return (a, .none)
-
-def Array.findAndModify {α : Type} (a : Array α) (pred : α → Bool) (update : α → α) : Array α × Option α :=
-  (a.findAndModifyM pred (update : α → Id α) : Id (Array α × Option α))
-
 def Std.HashMap.mapM {α β γ : Type} [BEq α] [Hashable α] {m : Type → Type} [Monad m] (map : Std.HashMap α β) (update : α → β → m γ) : m (Std.HashMap α γ) := do
   let a := map.toArray
   let a' ← a.mapM fun (x, y) ↦ do return (x, ← update x y)
   return Std.HashMap.ofArray a'
 
-def Std.HashMap.lens {α β : Type} [BEq α] [Hashable α] (key : α) : Lens (Std.HashMap α β) (Option β) where
-  get x := x.get? key
-  set v x := x.alter key fun _ ↦ v
-
 def Std.HashMap.modifyM {α β : Type} [BEq α] [Hashable α] {m : Type → Type} [Monad m] (map : Std.HashMap α β) (key : α) (update : β → m β) : m (Std.HashMap α β) := do
   if h : map.contains key then
     let new ← update (map.get key h)
@@ -85,7 +71,8 @@ structure Seat where
 
 structure State where
   registry : Registry.State := .new
-  windows : Array Window := #[]
+  windows : Std.HashMap RiverWindowV1 Window := {}
+  windowOrder : Array RiverWindowV1 := #[]
   outputs : Std.HashMap RiverOutputV1 Output := {}
   seats : Std.HashMap RiverSeatV1 Seat := {}
 
@@ -122,23 +109,23 @@ def handleOutput (ev : RiverOutputV1.Event) : DisplayM (Option Output) Unit := d
       | .removed => set (Option.some { o with removed := true })
       | _        => pure ()
 
-def handleWindow (proxy : RiverWindowV1) (ev : RiverWindowV1.Event) : DisplayM State Unit := do
-  let windows := (← get).windows.findAndModify (fun win ↦ win.proxy == proxy) fun win ↦
+def handleWindow (ev : RiverWindowV1.Event) : DisplayM (Option Window) Unit := do
+  for w in ← get do
     match ev with
-    | .closed                     => { win with closed := true }
-    | .dimensions w h             => { win with width := w, height := h }
-    | .pointerMoveRequested s     => { win with pointerMoveRequested := s }
-    | .pointerResizeRequested s e => { win with pointerResizeRequested := s, pointerResizeRequestedEdges := e }
-    | _                           => win
-  set { ← get with windows := windows.fst }
+      | .closed                     => set (Option.some { w with closed := true })
+      | .dimensions width height    => set (Option.some { w with width, height })
+      | .pointerMoveRequested s     => set (Option.some { w with pointerMoveRequested := s })
+      | .pointerResizeRequested s e => set (Option.some { w with pointerResizeRequested := s, pointerResizeRequestedEdges := e })
+      | _                           => pure ()
 
 structure Context where
   wm : RiverWindowManagerV1
   xkbBindings : RiverXkbBindingsV1
 
-def Window.setPosition {σ : Type} (w : Window) (x y : Int32) : DisplayM σ Window := do
-  w.node.setPosition x y
-  return { w with x, y}
+def Window.setPosition (x y : Int32) : DisplayM (Option Window) Unit := do
+  for w in ← get do
+    w.node.setPosition x y
+    set (Option.some { w with x, y})
 
 def Seat.createXkbBinding (ctx : Context) (mods : RiverSeatV1.Modifiers) (keysym : UInt32) (action : Action) : DisplayM (Option Seat) Unit := do
   for s in ← get do
@@ -176,18 +163,15 @@ def Seat.pointerResize {σ : Type} (seat : Seat) (win : Window) (edges : RiverWi
   }
 
 def Seat.focusTop (s : Seat) : DisplayM State Seat := do
-  let windows := (← get).windows
-  if h : windows.size == 0 then
-    s.proxy.clearFocus
-    return { s with focused := .none }
-  else
-    let win := windows[windows.size - 1]'(by
-      apply Nat.sub_one_lt
-      rw [beq_iff_eq] at h
-      assumption)
-    s.proxy.focusWindow win.proxy
-    win.node.placeTop
-    return { s with focused := win.proxy }
+  let windows := (← get).windowOrder
+  if h : windows.size != 0 then
+    let proxy := windows[windows.size - 1]'(Nat.sub_one_lt (bne_iff_ne.mp h))
+    for win in (← get).windows.get? proxy do
+      s.proxy.focusWindow win.proxy
+      win.node.placeTop
+      return { s with focused := win.proxy }
+  s.proxy.clearFocus
+  return { s with focused := .none }
 
 def Seat.doAction (s : Seat) (ctx : Context) : DisplayM State Seat := do
   let action := s.pendingAction
@@ -204,21 +188,19 @@ def Seat.doAction (s : Seat) (ctx : Context) : DisplayM State Seat := do
         for window in s.focused do
           window.close
     | .focusNext =>
-        let mut windows := (← get).windows
-        if h : 0 < windows.size then
-          let win := windows[0]
-          windows := windows.eraseIdx 0
-          windows := windows.push win
-        set { ← get with windows }
+        let mut windowOrder := (← get).windowOrder
+        if h : 0 < windowOrder.size then
+          let win := windowOrder[0]
+          windowOrder := windowOrder.eraseIdx 0
+          windowOrder := windowOrder.push win
+        set { ← get with windowOrder }
         s ← focusTop s
     | .move =>
-        for proxy in s.hovered do
-          for win in (← get).windows.find? (fun win ↦ win.proxy == proxy) do
-            s ← s.pointerMove win
+        for win in s.hovered >>= (← get).windows.get? do
+          s ← s.pointerMove win
     | .resize =>
-        for proxy in s.hovered do
-          for win in (← get).windows.find? (fun win ↦ win.proxy == proxy) do
-            s ← s.pointerResize win { bottom := true, right := true }
+        for win in s.hovered >>= (← get).windows.get? do
+          s ← s.pointerResize win { bottom := true, right := true }
     | .exit =>
         ctx.wm.exitSession
   return s
@@ -251,9 +233,11 @@ def removeOutputs : DisplayM State Unit := do
   set { ← get with outputs := aliveOutputs }
 
 def removeWindows : DisplayM State Unit := do
-  let (closedWindows, aliveWindows) := (← get).windows.partition Window.closed
+  let (closedWindows, aliveWindows) := (← get).windows.partition fun _ win ↦ win.closed
   let mut seats := (← get).seats
-  for win in closedWindows do
+  let mut windowOrder := (← get).windowOrder
+  for win in closedWindows.values do
+    windowOrder := windowOrder.erase win.proxy
     seats ← seats.mapM fun _ seat ↦ do
       let op  ← match seat.op with
         | .none => pure SeatOp.none
@@ -274,7 +258,7 @@ def removeWindows : DisplayM State Unit := do
         interacted := if seat.interacted == win.proxy then .none else seat.interacted
       }
     win.proxy.destroy
-  set { ← get with windows := aliveWindows, seats }
+  set { ← get with windows := aliveWindows, windowOrder, seats }
 
 def removeSeats : DisplayM State Unit := do
   let (removedSeats, aliveSeats) := (← get).seats.partition fun _ seat ↦ seat.removed
@@ -287,14 +271,11 @@ def removeSeats : DisplayM State Unit := do
   set { ← get with seats := aliveSeats }
 
 def initNewWindows : DisplayM State Unit := do
-  let windows ← (← get).windows.mapM fun win ↦ do
-    let mut win := win
-    if win.new then
-      win ← win.setPosition win.x win.y
-      win.proxy.proposeDimensions win.width win.height
-      win := { win with new := false }
-    pure win
-  set { ← get with windows }
+  for w in (← get).windows.values.filter Window.new do
+    DisplayM.zoom (mk_lens windows ≫ Std.HashMap.lens w.proxy) do
+      set (Option.some { w with new := false })
+      Window.setPosition w.x w.y
+      w.proxy.proposeDimensions w.width w.height
 
 def initNewSeats (ctx : Context) : DisplayM State Unit := do
   -- See xkbcommon/xkbcommon-keysyms.h
@@ -318,7 +299,7 @@ def initNewSeats (ctx : Context) : DisplayM State Unit := do
       Seat.createPointerBinding mods btnRight .resize
 
 def manageWindows : DisplayM State Unit := do
-  let windows ← (← get).windows.mapM fun win ↦ do
+  let windows ← (← get).windows.mapM fun _ win ↦ do
     let mut seats := (← get).seats
     for proxy in win.pointerMoveRequested do
       seats ← seats.modifyM proxy (fun s ↦ s.pointerMove win)
@@ -334,13 +315,13 @@ def manageWindows : DisplayM State Unit := do
 def manageSeats (ctx : Context) : DisplayM State Unit := do
   let seats ← (← get).seats.mapM fun _ seat ↦ do
     for proxy in seat.interacted do
-      let mut windows := (← get).windows
-      let winIdx := windows.findIdx (fun win ↦ win.proxy == proxy)
-      if h : winIdx < windows.size then
-        let win := windows[winIdx]
-        windows := windows.eraseIdx winIdx
-        windows := windows.push win
-      set { ← get with windows }
+      let mut windowOrder := (← get).windowOrder
+      let winIdx := windowOrder.findIdx (· == proxy)
+      if h : winIdx < windowOrder.size then
+        let win := windowOrder[winIdx]
+        windowOrder := windowOrder.eraseIdx winIdx
+        windowOrder := windowOrder.push win
+      set { ← get with windowOrder }
     let mut seat := seat
     seat ← seat.focusTop
     seat ← seat.doAction ctx
@@ -374,28 +355,28 @@ def handle (ctx : Context) : RiverWindowManagerV1.Event → DisplayM State Unit
         match seat.op with
           | .none => pure ()
           | .move proxy startX startY =>
-              let windows ← (← get).windows.findAndModifyM
-                (fun win ↦ win.proxy == proxy)
-                (fun win ↦ win.setPosition (startX + seat.opDeltaX) (startY + seat.opDeltaY))
-              set { ← get with windows := windows.fst }
+              DisplayM.zoom (mk_lens windows ≫ Std.HashMap.lens proxy) do
+                Window.setPosition (startX + seat.opDeltaX) (startY + seat.opDeltaY)
           | .resize proxy startX startY startWidth startHeight edges =>
-              let windows ← (← get).windows.findAndModifyM
-                (fun win ↦ win.proxy == proxy)
-                (fun win ↦
-                  let x := startX + if edges.left then (startWidth - win.width) else 0
-                  let y := startY + if edges.top then (startHeight - win.height) else 0
-                  win.setPosition x y)
-              set { ← get with windows := windows.fst }
+              DisplayM.zoom (mk_lens windows ≫ Std.HashMap.lens proxy) do
+                let mut x := startX
+                let mut y := startY
+                for w in ← get do
+                  if edges.left then x := x + startWidth  - w.width
+                  if edges.top  then y := y + startHeight - w.height
+                Window.setPosition x y
       ctx.wm.renderFinish
   | .sessionLocked => pure ()
   | .sessionUnlocked => pure ()
   | .window proxy => do
-      let windows := (← get).windows.push {
+      let windows := (← get).windows.insert proxy {
         proxy
         node := ← proxy.getNode
       }
-      set { ← get with windows }
-      proxy.listen (handleWindow proxy)
+      let windowOrder := (← get).windowOrder.push proxy
+      set { ← get with windows, windowOrder }
+      DisplayM.zoom (mk_lens windows ≫ Std.HashMap.lens proxy) do
+        proxy.listen handleWindow
   | .output proxy => do
       let outputs := (← get).outputs.insert proxy { proxy }
       set { ← get with outputs }