|
|
import Lean
|
|
|
|
|
|
open Lean Meta
|
|
|
|
|
|
/-! ## Graph -/
|
|
|
|
|
|
variable [inst : BEq α] [inst : Hashable α]
|
|
|
|
|
|
structure Graph (α β : Type) [inst : BEq α] [inst : Hashable α] where
|
|
|
nodes: HashMap α β := {}
|
|
|
edges: Array (α × α) := {}
|
|
|
deriving Inhabited
|
|
|
|
|
|
instance [ToJson β] : ToJson (Graph Name β) := {
|
|
|
toJson := fun graph => Json.mkObj [
|
|
|
("nodes", Json.mkObj (graph.nodes.toList.map fun (a,b) => (a.toString, toJson b))),
|
|
|
("edges", toJson graph.edges)
|
|
|
]
|
|
|
}
|
|
|
|
|
|
instance : EmptyCollection (Graph α β) := ⟨default⟩
|
|
|
|
|
|
def Graph.insertNode (g : Graph α β) (a : α) (b : β) :=
|
|
|
{g with nodes := g.nodes.insert a b}
|
|
|
|
|
|
/-- Check if graph contains loops -/
|
|
|
partial def Graph.hasLoops (g : Graph α β) (visited0 : HashSet α := {}): Bool := Id.run do
|
|
|
let mut visited : HashSet α := visited0
|
|
|
let all : Array α := g.nodes.toArray.map (·.1)
|
|
|
|
|
|
-- find some node that we haven't visited
|
|
|
let some x := all.find? fun x => ¬ visited.contains x
|
|
|
| return false -- We have visted all nodes and found no loops
|
|
|
visited := visited.insert x
|
|
|
|
|
|
match visitSuccessors x x visited with -- visit all recursive successors of x
|
|
|
| some visited' => visited := visited'
|
|
|
| none => return true -- none means a loop has been found
|
|
|
|
|
|
g.hasLoops visited -- continue looking for unvisited nodes
|
|
|
|
|
|
where
|
|
|
visitSuccessors (x : α) (x0 : α) (visited0 : HashSet α) : Option (HashSet α) := Id.run do
|
|
|
let mut visited : HashSet α := visited0
|
|
|
|
|
|
let directSuccessors := (g.edges.filter (·.1 == x)).map (·.2)
|
|
|
for y in directSuccessors do
|
|
|
if y == x0 then
|
|
|
return none -- loop found
|
|
|
if visited.contains y then
|
|
|
continue -- no loop possible here because the visited nodes do not lead to x0
|
|
|
visited := visited.insert y
|
|
|
match visitSuccessors y x0 visited with
|
|
|
| some visited' => visited := visited'
|
|
|
| none => return none
|
|
|
|
|
|
return some visited
|
|
|
|
|
|
-- #eval Graph.hasLoops ⟨HashMap.ofList [(2,2), (1,1)], #[(2,1)]⟩
|
|
|
|
|
|
partial def Graph.predecessors (g : Graph α β) (x : α) (acc : HashSet α := {}) : HashSet α := Id.run do
|
|
|
let mut res := acc
|
|
|
let directPredecessors := (g.edges.filter (·.2 == x)).map (·.1)
|
|
|
for y in directPredecessors do
|
|
|
if ¬ res.contains y then
|
|
|
res := res.insert y
|
|
|
res := g.predecessors y res
|
|
|
|
|
|
return res
|