(de-)marshal enums and bitfields

1 files changed, 79 insertions, 33 deletions

diff --git a/lib/Wayland/Display.lean b/lib/Wayland/Display.lean
index 70f5c67..63bd70d 100644
--- a/lib/Wayland/Display.lean
+++ b/lib/Wayland/Display.lean
@@ -43,7 +43,7 @@ structure AnyObject : Type where
   deriving Nonempty
 
 /-- Monad in which arguments can be marshalled one after another. -/
-class ArgumentMarshaller (m : Type → Type) extends Monad m where
+class ArgumentMarshaller (m : Type → Type) extends Monad m, MonadExcept String m where
   marshalInt                        : Int32             → m Unit
   marshalUInt                       : UInt32            → m Unit
   marshalFixed                      : Float32           → m Unit
@@ -64,10 +64,20 @@ class ArgumentDemarshaller (m : Type → Type) extends Monad m, MonadExcept Stri
   demarshalArray                      : m ByteArray
   demarshalFd                         : m Fd
 
+/-- Argument description in an interface's message. -/
+inductive MarshalType where
+  | int
+  | uint
+  | fixed
+  | string
+  | object
+  | array
+  | fd
+
 /-- A type is `Marshallable` if it can be marshalled and demarshalled using a corresponding monad. -/
 class Marshallable (α : Type) where
-  marshal {m : Type → Type} [ArgumentMarshaller m] : α → m Unit
-  demarshal {m : Type → Type} [ArgumentDemarshaller m] : m α
+  marshal {m : Type → Type} [ArgumentMarshaller m] : MarshalType → α → m Unit
+  demarshal {m : Type → Type} [ArgumentDemarshaller m] : MarshalType → m α
 
 /-- Nullability annotation in an interface's message. -/
 inductive Nullability where
@@ -250,7 +260,7 @@ deriving instance Monad for ArgumentBufferM
 private instance : MonadLift (ReaderT Native.ArgumentBuffer IO) ArgumentBufferM where
   monadLift := id
 
-private instance : ArgumentMarshaller ArgumentBufferM where
+private instance : ArgumentMarshaller (ExceptT String ArgumentBufferM) where
   marshalInt v := do (← read).writeInt v
   marshalUInt v := do (← read).writeUInt v
   marshalFixed v := do (← read).writeFixed v
@@ -279,55 +289,87 @@ private def ArgumentBufferM.run {α : Type} : ArgumentBufferM α → Native.Argu
   ReaderT.run
 
 instance : Marshallable Int32 where
-  marshal := ArgumentMarshaller.marshalInt
-  demarshal := ArgumentDemarshaller.demarshalInt
+  marshal
+  | .int => ArgumentMarshaller.marshalInt
+  | _    => fun _ ↦ throw "Marshalling unexpected type, expected int."
+  demarshal
+  | .int => ArgumentDemarshaller.demarshalInt
+  | _    => throw "Demarshalling unexpected type, expected int."
 
 instance : Marshallable UInt32 where
-  marshal := ArgumentMarshaller.marshalUInt
-  demarshal := ArgumentDemarshaller.demarshalUInt
+  marshal
+  | .uint => ArgumentMarshaller.marshalUInt
+  | _     => fun _ ↦ throw "Marshalling unexpected type, expected uint."
+  demarshal
+  | .uint => ArgumentDemarshaller.demarshalUInt
+  | _     => throw "Demarshalling unexpected type, expected uint."
 
 instance : Marshallable Float32 where
-  marshal := ArgumentMarshaller.marshalFixed
-  demarshal := ArgumentDemarshaller.demarshalFixed
+  marshal
+  | .fixed => ArgumentMarshaller.marshalFixed
+  | _      => fun _ ↦ throw "Marshalling unexpected type, expected fixed."
+  demarshal
+  | .fixed => ArgumentDemarshaller.demarshalFixed
+  | _      => throw "Demarshalling unexpected type, expected fixed."
 
 instance : Marshallable (Option String) where
-  marshal := ArgumentMarshaller.marshalString
-  demarshal := ArgumentDemarshaller.demarshalString
+  marshal
+  | .string => ArgumentMarshaller.marshalString
+  | _       => fun _ ↦ throw "Marshalling unexpected type, expected string."
+  demarshal
+  | .string => ArgumentDemarshaller.demarshalString
+  | _       => throw "Demarshalling unexpected type, expected string."
 
 instance (ι : InterfaceName) : Marshallable (Option (Object ι)) where
-  marshal := ArgumentMarshaller.marshalObject
-  demarshal := ArgumentDemarshaller.demarshalObject
+  marshal
+  | .object => ArgumentMarshaller.marshalObject
+  | _       => fun _ ↦ throw "Marshalling unexpected type, expected object."
+  demarshal
+  | .object => ArgumentDemarshaller.demarshalObject
+  | _       => throw "Demarshalling unexpected type, expected object."
 
 instance : Marshallable (Option AnyObject) where
-  marshal := ArgumentMarshaller.marshalAnyObject
-  demarshal := ArgumentDemarshaller.demarshalAnyObject
+  marshal
+  | .object => ArgumentMarshaller.marshalAnyObject
+  | _       => fun _ ↦ throw "Marshalling unexpected type, expected object."
+  demarshal
+  | .object => ArgumentDemarshaller.demarshalAnyObject
+  | _       => throw "Demarshalling unexpected type, expected object."
 
 instance : Marshallable ByteArray where
-  marshal := ArgumentMarshaller.marshalArray
-  demarshal := ArgumentDemarshaller.demarshalArray
+  marshal
+  | .array => ArgumentMarshaller.marshalArray
+  | _      => fun _ ↦ throw "Marshalling unexpected type, expected array."
+  demarshal
+  | .array => ArgumentDemarshaller.demarshalArray
+  | _      => throw "Demarshalling unexpected type, expected array."
 
 instance : Marshallable Fd where
-  marshal := ArgumentMarshaller.marshalFd
-  demarshal := ArgumentDemarshaller.demarshalFd
+  marshal
+  | .fd => ArgumentMarshaller.marshalFd
+  | _   => fun _ ↦ throw "Marshalling unexpected type, expected fd."
+  demarshal
+  | .fd => ArgumentDemarshaller.demarshalFd
+  | _   => throw "Demarshalling unexpected type, expected fd."
 
 instance : Marshallable String where
-  marshal s := Marshallable.marshal (Option.some s)
-  demarshal := do
-    match ((← Marshallable.demarshal) : Option String) with
+  marshal a s := Marshallable.marshal a (Option.some s)
+  demarshal a := do
+    match ((← Marshallable.demarshal a) : Option String) with
     | .some s => pure s
     | .none   => throw "Received null for a non-nullable string."
 
 instance (ι : InterfaceName) : Marshallable (Object ι) where
-  marshal o := Marshallable.marshal (Option.some o)
-  demarshal := do
-    match ((← Marshallable.demarshal) : Option (Object ι)) with
+  marshal a o := Marshallable.marshal a (Option.some o)
+  demarshal a := do
+    match ((← Marshallable.demarshal a) : Option (Object ι)) with
     | .some o => pure o
     | .none   => throw "Received null for a non-nullable object."
 
 instance : Marshallable AnyObject where
-  marshal o := Marshallable.marshal (Option.some o)
-  demarshal := do
-    match ((← Marshallable.demarshal) : Option AnyObject) with
+  marshal a o := Marshallable.marshal a (Option.some o)
+  demarshal a := do
+    match ((← Marshallable.demarshal a) : Option AnyObject) with
     | .some o => pure o
     | .none   => throw "Received null for a non-nullable object."
 
@@ -373,18 +415,22 @@ def DisplayM.getDisplay' : DisplayM σ (Object "wl_display") := do
 
 variable {ι ι' : InterfaceName}
 
+private def marshalArguments (marshaller : ∀ {α} [ArgumentMarshaller α], α OpCode) : IO (Native.ArgumentBuffer × OpCode) := do
+  let argBuffer ← Native.ArgumentBuffer.new
+  let opcode : Except String OpCode ← ArgumentBufferM.run (ExceptT.run marshaller) argBuffer
+  let opcode : OpCode ← IO.ofExcept opcode
+  return (argBuffer, opcode)
+
 /-- Send a request for an object to the server that doesn't construct any new object. -/
 def Object.send (obj : Object ι) (marshaller : ∀ {α} [ArgumentMarshaller α], α OpCode) : DisplayM σ Unit := do
-  let argBuffer ← Native.ArgumentBuffer.new
-  let opcode ← ArgumentBufferM.run marshaller argBuffer
+  let (argBuffer, opcode) ← marshalArguments marshaller
   obj.proxy.send opcode argBuffer
 
 /-- Send a request for an object to the server, allocating a new id for a constructed object in the process. -/
 def Object.sendConstructor [iface : ResolvedInterface ι'] (obj : Object ι) (marshaller : ∀ {α} [ArgumentMarshaller α], α OpCode) : DisplayM σ (Object ι') := do
   DisplayM.ensureInterface ι'
   let typeMap := (← read).typeMap
-  let argBuffer ← Native.ArgumentBuffer.new
-  let opcode ← ArgumentBufferM.run marshaller argBuffer
+  let (argBuffer, opcode) ← marshalArguments marshaller
   let proxy ← obj.proxy.sendConstructor opcode argBuffer typeMap ι'
   return ⟨proxy⟩