Improve filterAncestors to avoid walk to rootHEADmaster

1 files changed, 11 insertions, 6 deletions

diff --git a/src/Erebos/Storage/Graph.hs b/src/Erebos/Storage/Graph.hs
index 99ad3e6..853143f 100644
--- a/src/Erebos/Storage/Graph.hs
+++ b/src/Erebos/Storage/Graph.hs
@@ -127,12 +127,17 @@ precedesOrEquals x y = filterAncestors [ x, y ] == [ y ]
 filterAncestors :: Storable a => [ Stored a ] -> StoredTips a
 filterAncestors [ x ] = [ x ]
 filterAncestors xs = let xs' = uniq $ sort xs
-                      in helper xs' xs'
-    where helper remains walk = case generationMax walk of
-                                     Just x -> let px = previous x
-                                                   remains' = filter (\r -> all (/=r) px) remains
-                                                in helper remains' $ uniq $ sort (px ++ filter (/=x) walk)
-                                     Nothing -> remains
+                      in sort $ helper xs' xs'
+  where
+    helper [] _ = []
+    helper remains walk =
+        case generationMax walk of
+            Just x ->
+                let px = previous x
+                    youngerThenX r = Just GT == compareGeneration (storedGeneration r) (storedGeneration x)
+                    ( returned, remains' ) = partition youngerThenX $ filter (\r -> all (/= r) px) remains
+                 in returned ++ helper remains' (uniq $ sort (px ++ filter (/= x) walk))
+            Nothing -> remains
 
 commonAncestors :: Storable a => [ Stored a ] -> [ Stored a ] -> StoredTips a
 commonAncestors [] _ = []