1 files changed, 35 insertions, 9 deletions
diff --git a/tools/Scanner/Processor.lean b/tools/Scanner/Processor.lean
index 88853a6..d591ee7 100644
--- a/tools/Scanner/Processor.lean
+++ b/tools/Scanner/Processor.lean
@@ -1,5 +1,14 @@
 import Scanner.Protocol
 import Std.Data.TreeMap
+import Std.Data.TreeSet
+
+/-!
+# Processing of Protocol Descriptions
+
+This is mainly used to resolve interface/enum references within protocols, converting a `Raw.ProtocolDesc` into a
+`Client.ProtocolDesc`. Note that this requires a global view of all protocols since interfaces can mention interfaces
+from other protocols. Thus, this needs to be done after parsing all protocols.
+-/
 
 namespace WaylandScanner
 
@@ -12,7 +21,7 @@ private abbrev InterfaceMap := Std.TreeMap String InterfaceInfo
 private structure Context where
   interfaces : InterfaceMap
   currentInterface : Client.InterfaceRef
-  referencedInterfaces : Array Client.InterfaceRef
+  referencedInterfaces : Std.TreeSet Client.InterfaceRef
 
 private abbrev Processor : Type → Type := StateT Context (Except String)
 
@@ -37,7 +46,7 @@ private def resolveInterface (name : String) : Processor Client.InterfaceRef :=
     typeName := name
   }
   let old ← get
-  set { old with referencedInterfaces := old.referencedInterfaces ++ #[ref] }
+  set { old with referencedInterfaces := old.referencedInterfaces.insert ref }
   return ref
 
 private def resolveEnum (name : String) : Processor Client.EnumRef := do
@@ -114,7 +123,7 @@ private def processInterface (raw : Raw.InterfaceDesc) : Processor Client.Interf
   let requests ← raw.requests.mapM processRequest
   let events ← raw.events.mapM processEvent
   let state ← get
-  let referencedInterfaces := state.referencedInterfaces ++ #[state.currentInterface]
+  let referencedInterfaces := state.referencedInterfaces.insert state.currentInterface
   return {
     name := raw.name
     version := raw.version
@@ -123,9 +132,6 @@ private def processInterface (raw : Raw.InterfaceDesc) : Processor Client.Interf
     requests := requests
     events := events
     referencedInterfaces := referencedInterfaces
-      |>.toList
-      |>.eraseDups
-      |>.toArray
   }
 
 private def processProtocol (interfaces : InterfaceMap) (raw : Raw.ProtocolDesc) : Except String Client.ProtocolDesc := do
@@ -137,14 +143,34 @@ private def processProtocol (interfaces : InterfaceMap) (raw : Raw.ProtocolDesc)
       fun iface ↦ (processInterface iface).run' {
         interfaces := interfaces
         currentInterface := { moduleName := raw.name, typeName := iface.name }
-        referencedInterfaces := #[]
+        referencedInterfaces := .empty
       }
   }
 
-def processProtocols (raw : Array Raw.ProtocolDesc) : Except String (Array Client.ProtocolDesc) :=
+/-! Compute the transitive closure w.r.t. `referencedInterfaces`. -/
+private partial def transitiveClosure (protocols : Array Client.ProtocolDesc) : Array Client.ProtocolDesc :=
+  let successors : Std.TreeMap Client.InterfaceRef (List Client.InterfaceRef) := .ofArray <| protocols.flatMap
+    fun protocol ↦ protocol.interfaces.map
+      fun iface ↦ ((⟨protocol.name, iface.name⟩ : Client.InterfaceRef), iface.referencedInterfaces.toList)
+  let rec dfs : Std.TreeSet Client.InterfaceRef → List Client.InterfaceRef → Std.TreeSet Client.InterfaceRef
+  | visited, [] => visited
+  | visited, (x::xs) =>
+    if visited.contains x
+    then dfs visited xs
+    else dfs (visited.insert x) (successors.get! x ++ xs)
+  let interfaceClosure (iface : Client.InterfaceDesc) : Client.InterfaceDesc :=
+    { iface with referencedInterfaces := dfs .empty iface.referencedInterfaces.toList }
+  let protocolClosure (protocol : Client.ProtocolDesc) : Client.ProtocolDesc :=
+    { protocol with interfaces := protocol.interfaces.map interfaceClosure }
+  protocols.map protocolClosure
+
+-- TODO: Correctness proof for `transitiveClosure`
+
+def processProtocols (raw : Array Raw.ProtocolDesc) : Except String (Array Client.ProtocolDesc) := do
   let interfaces : InterfaceMap := .ofArray <| raw.flatMap
     fun protocol ↦ protocol.interfaces.map
       fun iface ↦ (iface.name, (⟨protocol.name, iface.enums.map Raw.EnumDesc.name⟩ : InterfaceInfo))
-  raw.mapM (processProtocol interfaces)
+  let res ← raw.mapM (processProtocol interfaces)
+  return transitiveClosure res
 
 end WaylandScanner