From eb3977aab73bc626a99abc3870c5a8ae553d14bc Mon Sep 17 00:00:00 2001
From: numirias <numirias@users.noreply.github.com>
Date: Mon, 18 Dec 2017 16:40:30 +0100
Subject: [PATCH] Fix restore resizing

---
 plasma/node.py     |  9 ++++++---
 tests/test_node.py | 15 +++++++++++++++
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/plasma/node.py b/plasma/node.py
index 944130e..7f1c60c 100644
--- a/plasma/node.py
+++ b/plasma/node.py
@@ -479,7 +479,6 @@ def close_right(self):
         return self.close_neighbor(RIGHT)
 
     def add_child(self, node, idx=None):
-        # TODO Currently we assume that the node has no size.
         if idx is None:
             idx = len(self)
         self.children.insert(idx, node)
@@ -497,12 +496,14 @@ def remove_child(self, node):
         node.force_size(0)
         self.children.remove(node)
         if len(self) == 1:
+            child = self[0]
             if self.is_root:
                 # A single child doesn't need a fixed size
-                self[0].reset_size()
+                child.reset_size()
             else:
                 # Collapse tree with a single child
-                self.parent.replace_child(self, self[0])
+                self.parent.replace_child(self, child)
+                Node.fit_into(child, self.capacity)
 
     def remove(self):
         self.parent.remove_child(self)
@@ -555,8 +556,10 @@ def restore(self, node):
             raise NotRestorableError()
         node.reset_size()
         if flip:
+            old_parent_size = parent.size
             parent.flip_with(node, reverse=(idx == 0))
             node.size, parent.size = sizes
+            Node.fit_into(parent, old_parent_size)
         else:
             parent.add_child(node, idx=idx)
             node.size = sizes[0]
diff --git a/tests/test_node.py b/tests/test_node.py
index db3946f..9a4efee 100644
--- a/tests/test_node.py
+++ b/tests/test_node.py
@@ -893,6 +893,21 @@ def test_resize_max(self, root, tiny_grid):
         assert a.width == 110
         assert b.width == c.width == 10
 
+    def test_resize_with_collapse_and_restore(self, root, small_grid):
+        a, b, c, d = small_grid
+        root.height = 30
+        c.size = 30
+        d.size += 10
+        b.remove()
+        assert c.size == c.height == 10
+        assert d.size == d.height == 20
+        root.restore(b)
+        assert b.height == 15
+        assert b.width == 60
+        assert c.height == d.height == 15
+        assert c.width == 20
+        assert d.width == 40
+
 class TestRestore:
 
     def test_restore(self, root, grid):