From d987b39bbd876e2d1ade2bb91e1a9ce90273704c Mon Sep 17 00:00:00 2001
From: Pavel Patsey <patsey.pavel@gmail.com>
Date: Fri, 24 Jan 2025 10:05:11 +0300
Subject: [PATCH] add merge_nodes function

---
 tests/cidr4_merge/test_cidr4_merger.py  | 61 +++++++++++++++++++------
 vpn_manager/cidr4_merge/cidr4_merger.py | 48 ++++++++++---------
 2 files changed, 75 insertions(+), 34 deletions(-)

diff --git a/tests/cidr4_merge/test_cidr4_merger.py b/tests/cidr4_merge/test_cidr4_merger.py
index 32add5f..2c87518 100644
--- a/tests/cidr4_merge/test_cidr4_merger.py
+++ b/tests/cidr4_merge/test_cidr4_merger.py
@@ -6,6 +6,7 @@ from vpn_manager.cidr4_merge.cidr4_merger import (
     cidr4_to_node,
     find_parent,
     make_cidr4,
+    merge_nodes,
     merge_two_nodes,
     sort_nodes,
 )
@@ -89,20 +90,54 @@ def test_find_parent__with_exception():
 
 
 def test_calc_dip():
-    assert calc_dip(32, 32) == 0
-    assert calc_dip(32, 31) == 2
-    assert calc_dip(32, 30) == 4
-    assert calc_dip(31, 30) == 2
-    assert calc_dip(30, 29) == 4
-    assert calc_dip(29, 31) == 6
-    assert calc_dip(26, 26) == 0
-    assert calc_dip(3, 2) == 2**30 - 2**29
-    assert calc_dip(3, 1) == 2**31 - 2**29
+    assert calc_dip(32, 32, 31) == 0
+    assert calc_dip(32, 32, 30) == 2
+    assert calc_dip(32, 32, 29) == 6
+    assert calc_dip(32, 30, 29) == 3
+    assert calc_dip(32, 31, 29) == 5
+    assert calc_dip(32, 31, 29) == 5
+
+    assert calc_dip(2, 2, 1) == 0
+    assert calc_dip(2, 2, 0) == 2**31
+    assert calc_dip(3, 3, 1) == 2**30
 
 
 def test_merge_two_nodes():
     assert merge_two_nodes((0, 32), (1, 32)) == ((0, 31), 0)
-    assert merge_two_nodes((0, 32), (2, 32)) == ((0, 30), 4)
-    assert merge_two_nodes((0, 32), (5, 32)) == ((0, 29), 8)
-    assert merge_two_nodes((3, 32), (4, 32)) == ((0, 29), 8)
-    assert merge_two_nodes((0, 32), (4, 30)) == ((0, 29), 4)
+    assert merge_two_nodes((0, 32), (2, 32)) == ((0, 30), 2)
+    assert merge_two_nodes((0, 32), (5, 32)) == ((0, 29), 6)
+    assert merge_two_nodes((3, 32), (4, 32)) == ((0, 29), 6)
+    assert merge_two_nodes((0, 32), (4, 30)) == ((0, 29), 3)
+    assert merge_two_nodes((0, 32), (6, 31)) == ((0, 29), 5)
+
+
+def test_merge_nodes():
+    assert merge_nodes(
+        [
+            (0, 32),
+            (3, 32),
+            (4, 32),
+        ],
+        2,
+    ) == ([(0, 30), (4, 32)], 2)
+
+    # т.е. сейчас алгоритм работает так, что он может оставить соседей
+    assert merge_nodes(
+        [
+            (0, 32),
+            (3, 32),
+            (4, 32),
+            (7, 32),
+        ],
+        2,
+    ) == ([(0, 30), (4, 30)], 4)
+
+    assert merge_nodes(
+        [
+            (0, 32),
+            (3, 32),
+            (4, 32),
+            (7, 32),
+        ],
+        1,
+    ) == ([(0, 29)], 4)
diff --git a/vpn_manager/cidr4_merge/cidr4_merger.py b/vpn_manager/cidr4_merge/cidr4_merger.py
index c3bfeab..61d5893 100644
--- a/vpn_manager/cidr4_merge/cidr4_merger.py
+++ b/vpn_manager/cidr4_merge/cidr4_merger.py
@@ -42,13 +42,17 @@ def find_parent(a: Node, b: Node) -> Node:
     return parent_node
 
 
-def calc_dip(mask_len_a: int, mask_len_b: int) -> int:
-    # a = 2 ** (32 - mask_len_a) if mask_len_a != 32 else 0
-    # b = 2 ** (32 - mask_len_b) if mask_len_b != 32 else 0
-    a = 1 << (32 - mask_len_a) if mask_len_a != 32 else 0
-    b = 1 << (32 - mask_len_b) if mask_len_b != 32 else 0
-    dip = abs(a - b)
-    return dip
+def calc_dip(mask_len_a: int, mask_len_b: int, mask_len_p: int) -> int:
+    mask_len = mask_len_p + 1
+    dip_a = 0
+    while mask_len_a > mask_len:
+        dip_a += 2 ** (32 - mask_len_a)
+        mask_len_a -= 1
+    dip_b = 0
+    while mask_len_b > mask_len:
+        dip_b += 2 ** (32 - mask_len_b)
+        mask_len_b -= 1
+    return dip_a + dip_b
 
 
 def merge_two_nodes(node_a: Node, node_b: Node) -> tuple[Node, int]:
@@ -56,9 +60,7 @@ def merge_two_nodes(node_a: Node, node_b: Node) -> tuple[Node, int]:
     ip_b, mask_len_b = node_b
     parent_node = find_parent(node_a, node_b)
     _, mask_len_p = parent_node
-    dip_a = calc_dip(mask_len_a, mask_len_p + 1)
-    dip_b = calc_dip(mask_len_b, mask_len_p + 1)
-    dip = dip_a + dip_b
+    dip = calc_dip(mask_len_a, mask_len_b, mask_len_p)
     return parent_node, dip
 
 
@@ -68,18 +70,22 @@ def make_cidr4(ip, mask_len) -> str:
     return f"{ip_address}/{mask_len}"
 
 
-def merge_nodes(
-    nodes_to_merge: list[Node], required_len: int
-) -> tuple[list[Node], int]:
-    nodes = [x for x in nodes_to_merge]
+def merge_nodes(nodes: list[Node], required_len: int) -> tuple[list[Node], int]:
     sum_dip = 0
-    # Преобразовать список нод в список туплов: родитель (ip, mask len), кол-во добавляемых адресов d_ip
-    # найти подходящего родителя (с минимальным значением d_ip) и индекс
-    # затем мержить два узла:
-    # - т.е. удалить элемент с индексом i
-    # - еще раз удалить элемент с индексом i
-    # - добавить родителя перед элементом с индексом i
-    # повторить пока не достигнем нужного кол-ва узлов
+    while len(nodes) > required_len:
+        min_tuple = None, (None, float("inf"))
+        for i, (a, b) in enumerate(zip(nodes, nodes[1:])):
+            parent_node, dip = merge_two_nodes(a, b)
+            if dip < min_tuple[1][1]:
+                min_tuple = i, (parent_node, dip)
+        idx, (parent_node, dip) = min_tuple
+        nodes = nodes[:idx] + [parent_node] + nodes[idx + 2 :]
+        sum_dip += dip
+
+        # for i, (a, b) in enumerate(zip(nodes, nodes[1:])):
+        #     parent_node, dip = merge_two_nodes(a, b)
+        #     assert parent_node != a
+        #     assert parent_node != b
 
     return nodes, sum_dip