1 files changed, 68 insertions, 15 deletions

diff --git a/src/main.c b/src/main.c
index a710ecb..f3e455b 100644
--- a/src/main.c
+++ b/src/main.c
@@ -109,20 +109,31 @@ struct AVLNode* left_rotate(struct AVLNode* parent)
   return child1;
 }
 
-struct AVLNode* insert_mask(struct AVLNode* node, struct MaskData* data)
+struct Result insert_mask(struct AVLNode* node, struct MaskData* data)
 {
+  struct Result result;
   // 1. Standard BST insertion
   if (node == NULL) {
-    return create_avl_mask_node(data);
+    return (struct Result) {create_avl_mask_node(data), TRUE};
   }
 
   struct MaskData *node_data = (struct MaskData*)node->data;
   if (node->compare(data, node_data)) {
-    node->left = insert_mask(node->left, data);
+    result = insert_mask(node->left, data);
+    if (!result.success) {
+      fprintf(stderr, "Failed to insert!");
+      return result;
+    }
+    node->left = (struct AVLNode*)result.data;
   } else if (node->compare(node->data, data)) {
-    node->right = insert_mask(node->right, data);
+    result = insert_mask(node->right, data);
+    if (!result.success) {
+      fprintf(stderr, "Failed to insert!");
+      return result;
+    }
+    node->right = (struct AVLNode*)result.data;
   } else {
-    return node;
+    return (struct Result) {node, FALSE};
   }
 
   // 2. Update height of the ancestor node
@@ -134,21 +145,31 @@ struct AVLNode* insert_mask(struct AVLNode* node, struct MaskData* data)
 
   // LeftLeft
   if ((balance > 1) && node->compare(data, node->left->data)) {
-    return right_rotate(node);
+    return (struct Result) {right_rotate(node), TRUE};
   }
   // RightRight
   if ((balance < -1) && node->compare(node->right->data, data)) {
-    return left_rotate(node);
+    return (struct Result) {left_rotate(node), TRUE};
   }
   // LeftRight
   if ((balance > 1) && node->compare(node->left->data, data)) {
-    return right_rotate(node);
+    return (struct Result) {right_rotate(node), TRUE};
   }
   // RightLeft
   if ((balance < -1) && node->compare(data,node->right->data)) {
-    return left_rotate(node);
+    return (struct Result) {left_rotate(node), TRUE};
   }
-  return node;
+  return (struct Result) {node, TRUE};
+}
+
+struct AVLNode* insert_mask_alloc(struct AVLNode* node, uint16_t label)
+{
+  struct MaskData* data = create_mask_data(label);
+  struct Result result = insert_mask(node, data);
+  if (!result.success) {
+    free(data);
+  }
+  return (struct AVLNode*)result.data;
 }
 
 void print_label(struct AVLNode* root) {
@@ -179,14 +200,46 @@ void free_avl_tree_nodes(struct AVLNode* root) {
   }
 }
 
+void increase_label_area(struct AVLNode* root, uint16_t label) {
+  if (root == NULL) {
+    return;
+  }
+  struct MaskData* data = (struct MaskData*)root->data;
+  if (data->label == label) {
+    data->area++;
+  }
+  else if (data->label > label) {
+    increase_label_area(root->left, label);
+  }
+  else if (data->label < label) {
+    increase_label_area(root->right, label);
+  }
+}
+
+void increase_label_perimeter(struct AVLNode* root, uint16_t label) {
+  if (root == NULL) {
+    return;
+  }
+  struct MaskData* data = (struct MaskData*)root->data;
+  if (data->label == label) {
+    data->perimeter++;
+  }
+  else if (data->label > label) {
+    increase_label_perimeter(root->left, label);
+  }
+  else if (data->label < label) {
+    increase_label_perimeter(root->right, label);
+  }
+}
+
 int main(int argc, char** argv)
 {
   struct AVLNode* root = NULL;
-  root = insert_mask(root, create_mask_data(2));
-  root = insert_mask(root, create_mask_data(5));
-  root = insert_mask(root, create_mask_data(1));
-  root = insert_mask(root, create_mask_data(3));
-  root = insert_mask(root, create_mask_data(4));
+  root = insert_mask_alloc(root, 2);
+  root = insert_mask_alloc(root, 5);
+  root = insert_mask_alloc(root, 1);
+  root = insert_mask_alloc(root, 3);
+  root = insert_mask_alloc(root, 4);
   printf("Inorder traversal of AVL tree: ");
   print_label(root);
   printf("\n");