Merge pull request #41 from SciML/hr/performance

Author: ranocha

src/RootedTrees.jl

@@ -682,7 +682,8 @@ function Base.iterate(forest::PartitionForestIterator, state)
   # and `edge_set`.
   deleteat!(level_sequence, subtree_root_index:subtree_last_index)
   deleteat!(edge_set, subtree_root_index-1:subtree_last_index-1)
-  edge_to_remove = findlast(==(false), edge_set)
+
+  edge_to_remove = findprev(==(false), edge_set, edge_to_remove - 1)
   if edge_to_remove === nothing
     edge_to_remove = typemin(Int)
   end
@@ -755,7 +756,7 @@ function partition_skeleton!(level_sequence, edge_set)
     deleteat!(level_sequence, subtree_root_index)
     deleteat!(edge_set, edge_to_contract)
 
-    edge_to_contract = findlast(edge_set)
+    edge_to_contract = findprev(edge_set, edge_to_contract - 1)
   end
 
   # The level sequence `level_sequence` will not automatically be a canonical
@@ -788,14 +789,27 @@ function all_partitions(t::RootedTree)
   skeletons = [partition_skeleton(t, edge_set)]
 
   for edge_set_value in 1:(2^length(edge_set) - 1)
-    digits!(edge_set, edge_set_value, base=2)
+    binary_digits!(edge_set, edge_set_value)
     push!(forests,   partition_forest(t, edge_set))
     push!(skeletons, partition_skeleton(t, edge_set))
   end
 
   return (; forests, skeletons)
 end
 
+# A helper function to comute the binary representation of an integer `n` as
+# a vector of `Bool`s. This is a more efficient version of
+#   binary_digits!(digits, n) = digits!(digits, n, base=2)
+function binary_digits!(digits::Vector{Bool}, n::Int)
+  bit = 1
+  for i in eachindex(digits)
+    digits[i] = n & bit > 0
+    bit = bit << 1
+  end
+  digits
+end
+
+
 
 """
     PartitionIterator(t::RootedTree)
@@ -855,7 +869,7 @@ function Base.iterate(partitions::PartitionIterator, edge_set_value)
   edge_set     = partitions.edge_set
   edge_set_tmp = partitions.edge_set_tmp
 
-  digits!(edge_set, edge_set_value, base=2)
+  binary_digits!(edge_set, edge_set_value)
 
   # Compute the partition skeleton.
   # The following is a more efficient version of
@@ -918,7 +932,7 @@ function all_splittings(t::RootedTree)
   subtrees = Vector{RootedTree{T, Vector{T}}}() # ordered subtrees
 
   for node_set_value in 0:(2^order(t) - 1)
-    digits!(node_set, node_set_value, base=2)
+    binary_digits!(node_set, node_set_value)
 
     # Check that if a node is removed then all of its descendants are removed
     subtree_root_index = 1
@@ -1000,7 +1014,7 @@ function Base.iterate(splittings::SplittingIterator, node_set_value)
   forest = Vector{RootedTree{T, Vector{T}}}()
 
   while node_set_value <= splittings.max_node_set_value
-    digits!(node_set, node_set_value, base=2)
+    binary_digits!(node_set, node_set_value)
 
     # Check that if a node is removed then all of its descendants are removed
     subtree_root_index = 1