flatMap for LazyGraph.t! hey, why not?

lazyGraph.ml

@@ -356,7 +356,7 @@ let a_star graph
   ?(heuristic=(fun v -> 0.))
   ?(distance=(fun v1 e v2 -> 1.))
   ~goal
-  ~start =
+  start =
   fun () ->
     (* map node -> 'came_from' cell *)
     let nodes = mk_map ~eq:graph.eq ~hash:graph.hash in
@@ -428,7 +428,7 @@ let dijkstra graph ?on_explore ?(ignore=fun v -> false)
   ?(distance=fun v1 e v2 -> 1.) v1 v2 =
   let paths =
     a_star graph ?on_explore ~ignore ~distance ~heuristic:(fun _ -> 0.)
-       ~goal:(fun v -> graph.eq v v2) ~start:v1 in
+       ~goal:(fun v -> graph.eq v v2) v1 in
   let paths = Gen.start paths in
   try
     Gen.Gen.next paths
@@ -472,6 +472,20 @@ let map ~vertices ~edges g =
       Node (v, vertices l, edges_enum')
   in { eq=g.eq; hash=g.hash; force; }
 
+(** Replace each vertex by some vertices. By mapping [v'] to [f v'=v1,...,vn],
+    whenever [v] ---e---> [v'], then [v --e--> vi] for i=1,...,n. *)
+let flatMap f g =
+  let force v =
+    match g.force v with
+    | Empty -> Empty
+    | Node (_, l, edges_enum) ->
+      let edges_enum' = Gen.flatMap
+        (fun (e, v') ->
+          Gen.map (fun v'' -> e, v'') (f v'))
+        edges_enum in
+      Node (v, l, edges_enum')
+  in { eq=g.eq; hash=g.hash; force; }
+
 let filter ?(vertices=(fun v l -> true)) ?(edges=fun v1 e v2 -> true) g =
   let force v =
     match g.force v with

lazyGraph.mli

@@ -135,7 +135,7 @@ val a_star : ('id, 'v, 'e) t ->
              ?heuristic:('id -> float) ->
              ?distance:('id -> 'e -> 'id -> float) ->
              goal:('id -> bool) ->
-             start:'id ->
+             'id ->
              (float * ('id, 'e) path) Gen.t
   (** Shortest path from the first node to nodes that satisfy [goal], according
       to the given (positive!) distance function. The distance is also returned.
@@ -176,6 +176,13 @@ val map : vertices:('v -> 'v2) -> edges:('e -> 'e2) ->
           ('id, 'v, 'e) t -> ('id, 'v2, 'e2) t
   (** Map vertice and edge labels *)
 
+val flatMap : ('id -> 'id Gen.t) ->
+              ('id, 'v, 'e) t ->
+              ('id, 'v, 'e) t
+  (** Replace each vertex by some vertices. By mapping [v'] to [f v'=v1,...,vn],
+      whenever [v] ---e---> [v'], then [v --e--> vi] for i=1,...,n. Optional
+      functions can be used to transform labels for edges and vertices. *)
+
 val filter : ?vertices:('id -> 'v -> bool) ->
              ?edges:('id -> 'e -> 'id -> bool) ->
              ('id, 'v, 'e) t -> ('id, 'v, 'e) t