From 4afa8a8a4e6437c2059b2a9c0fefc5973e89f017 Mon Sep 17 00:00:00 2001
From: Jake McGinty <me@jake.su>
Date: Wed, 23 Jun 2021 20:47:53 +0900
Subject: [PATCH] shared(types): CidrTree: re-add with_root, rewrite leaves()

---
 shared/src/types.rs | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/shared/src/types.rs b/shared/src/types.rs
index f37e478..789675a 100644
--- a/shared/src/types.rs
+++ b/shared/src/types.rs
@@ -237,6 +237,10 @@ impl<'a> CidrTree<'a> {
             .iter()
             .min_by_key(|c| c.cidr.prefix())
             .expect("failed to find root CIDR");
+        Self::with_root(cidrs, root)
+    }
+
+    pub fn with_root(cidrs: &'a [Cidr], root: &'a Cidr) -> Self {
         Self {
             cidrs,
             contents: root,
@@ -254,13 +258,11 @@ impl<'a> CidrTree<'a> {
     }
 
     pub fn leaves(&self) -> Vec<Cidr> {
-        let mut leaves = vec![];
-        for cidr in self.cidrs {
-            if !self.cidrs.iter().any(|c| c.parent == Some(cidr.id)) {
-                leaves.push(cidr.clone());
-            }
+        if !self.cidrs.iter().any(|cidr| cidr.parent == Some(self.id)) {
+            vec![self.contents.clone()]
+        } else {
+            self.children().flat_map(|child| child.leaves()).collect()
         }
-        leaves
     }
 }