diff --git a/main.cpp b/main.cpp
index 6e4a85f..3ace94d 100644
--- a/main.cpp
+++ b/main.cpp
@@ -4,22 +4,29 @@
 
 #include "areas.h"
 
+
+static std::vector<Rectangle> rects;
+
+void add_rect(Coordinate x1, Coordinate y1, Coordinate x2, Coordinate y2) {
+    rects.push_back({{x1,y1}, {x2,y2}, 0, 0});
+}
+
+void check_area(Unit area) {
+    Unit computed = get_area_union(rects);
+    std::cout << "Computed area is " << computed << std::endl;
+    assert(computed == area);
+}
+
+
 int main() {
     srand(987192345);
 
-    std::vector<Rectangle> rects;
-    rects.push_back({{0,0},{10,10},0,0});
-    Unit area = get_area_union(rects);
-    assert(area == 100);
+    add_rect(0,0,5,4);
+    check_area(20);
 
-    rects.push_back({{2,3},{4,12},0,0});
-    area = get_area_union(rects);
-    std::cout << "Area is " << area << std::endl;
-    assert(area == 104);
+    add_rect(3,2,6,6);
+    check_area(28);
 
-    std::vector<Rectangle> inst1 = get_random_instance(10,100);
-    area = get_area_union(inst1);
-    std::cout << "Area is " << area << "." << std::endl;
-    assert(area == 6559);
+    std::cout << "All checks passed" << std::endl;
     return 0;
 }
\ No newline at end of file
diff --git a/segment_tree.cpp b/segment_tree.cpp
index 9c9e20e..55a374d 100644
--- a/segment_tree.cpp
+++ b/segment_tree.cpp
@@ -59,6 +59,7 @@ void SegmentTree::remove_interval(Interval interval, Index node_idx) {
         if(right_child(node_idx).left_coord < interval.right) {
             remove_interval(interval, right_child_idx(node_idx));
         }
+        update_covered_length(node_idx);
     }
 }