From 89d968cba8ec8232c29cbbc80c24e4496819e58e Mon Sep 17 00:00:00 2001
From: Joseph Tindall <joeytindall94@googlemail.com>
Date: Thu, 26 Oct 2023 16:34:14 -0400
Subject: [PATCH] Better algorithm specification

---
 src/traversals/trees_and_forests.jl | 42 ++++++++++++++++++-----------
 test/test_trees_and_forests.jl      | 16 +++++------
 2 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/src/traversals/trees_and_forests.jl b/src/traversals/trees_and_forests.jl
index 8486360..8362652 100644
--- a/src/traversals/trees_and_forests.jl
+++ b/src/traversals/trees_and_forests.jl
@@ -1,33 +1,43 @@
+abstract type SpanningTreeAlgorithm end
+
+struct BFS <: SpanningTreeAlgorithm end
+struct RandomBFS <: SpanningTreeAlgorithm end
+struct DFS <: SpanningTreeAlgorithm end
+
+default_spanning_tree_alg() = BFS()
+
 default_root_vertex(g) = last(findmax(eccentricities(g)))
 
+function spanning_tree(g::AbstractNamedGraph; root_vertex=default_root_vertex(g))
+  return spanning_tree(default_spanning_tree_alg(), g; root_vertex)
+end
+
+function spanning_tree(::BFS, g::AbstractNamedGraph; root_vertex=default_root_vertex(g))
+  @assert !NamedGraphs.is_directed(g)
+  return undirected_graph(bfs_tree(g, root_vertex))
+end
+
 function spanning_tree(
-  g::AbstractNamedGraph; root_vertex=default_root_vertex(g), alg::String="BFS"
+  ::RandomBFS, g::AbstractNamedGraph; root_vertex=default_root_vertex(g)
 )
   @assert !NamedGraphs.is_directed(g)
-  if alg == "BFS"
-    return undirected_graph(bfs_tree(g, root_vertex))
-  elseif alg == "RandomBFS"
-    return undirected_graph(random_bfs_tree(g, root_vertex))
-  elseif alg == "DFS"
-    return undirected_graph(dfs_tree(g, root_vertex))
-  else
-    error("Algorithm not current supported")
-  end
+  return undirected_graph(random_bfs_tree(g, root_vertex))
+end
+
+function spanning_tree(::DFS, g::AbstractNamedGraph; root_vertex=default_root_vertex(g))
+  @assert !NamedGraphs.is_directed(g)
+  return undirected_graph(dfs_tree(g, root_vertex))
 end
 
 #Given a graph, split it into its connected components, construct a spanning tree over each of them
 # and take the union.
-function spanning_forest(
-  g::AbstractNamedGraph; spanning_tree_function=g -> spanning_tree(g)
-)
+function spanning_forest(g::AbstractNamedGraph; spanning_tree_function=spanning_tree)
   return reduce(union, (spanning_tree_function(g[vs]) for vs in connected_components(g)))
 end
 
 #Given an undirected graph g with vertex set V, build a set of forests (each with vertex set V) which covers all edges in g
 # (see https://en.wikipedia.org/wiki/Arboricity) We do not find the minimum but our tests show this algorithm performs well
-function build_forest_cover(
-  g::AbstractNamedGraph; spanning_tree_function=g -> spanning_tree(g)
-)
+function forest_cover(g::AbstractNamedGraph; spanning_tree_function=spanning_tree)
   edges_collected = edgetype(g)[]
   remaining_edges = edges(g)
   forests = NamedGraph[]
diff --git a/test/test_trees_and_forests.jl b/test/test_trees_and_forests.jl
index 1a76e9d..222d794 100644
--- a/test/test_trees_and_forests.jl
+++ b/test/test_trees_and_forests.jl
@@ -2,7 +2,7 @@ using Test
 using Graphs
 using NamedGraphs
 using NamedGraphs:
-  hexagonal_lattice_graph, triangular_lattice_graph, build_forest_cover, spanning_tree
+  hexagonal_lattice_graph, triangular_lattice_graph, forest_cover, spanning_tree
 
 @testset "Test Spanning Trees" begin
   gs = [
@@ -13,10 +13,10 @@ using NamedGraphs:
     named_grid((10, 10)),
     triangular_lattice_graph(5, 5; periodic=true),
   ]
-  algs = ["BFS", "DFS", "RandomBFS"]
+  algs = [NamedGraphs.BFS(), NamedGraphs.DFS(), NamedGraphs.RandomBFS()]
   for g in gs
     for alg in algs
-      s_tree = spanning_tree(g; alg)
+      s_tree = spanning_tree(alg, g)
       @test is_tree(s_tree)
       @test Set(vertices(s_tree)) == Set(vertices(g))
       @test issubset(Set(edges(s_tree)), Set(edges(g)))
@@ -34,12 +34,12 @@ end
     triangular_lattice_graph(5, 5; periodic=true),
   ]
   for g in gs
-    forest_cover = build_forest_cover(g)
-    cover_edges = reduce(vcat, edges.(forest_cover))
+    cover = forest_cover(g)
+    cover_edges = reduce(vcat, edges.(cover))
     @test issetequal(cover_edges, edges(g))
-    @test all(issetequal(vertices(f), vertices(g)) for f in forest_cover)
-    for f in forest_cover
-      trees = NamedGraph[f[vs] for vs in connected_components(f)]
+    @test all(issetequal(vertices(forest), vertices(g)) for forest in cover)
+    for forest in cover
+      trees = NamedGraph[forest[vs] for vs in connected_components(forest)]
       @test all(is_tree.(trees))
     end
   end