diff --git a/src/main.c b/src/main.c
index 9a6df2f8be037308fbbe3e0779cdccbc8400d3e3..339ac4b3ebf32beff32016775358989e0f5b137c 100644
--- a/src/main.c
+++ b/src/main.c
@@ -10,7 +10,11 @@ typedef struct Node {
     struct Node* next;
 } Node;
 
-/*@ requires \valid(head_ref);
+/*@ predicate is_valid_linked_list(Node* head) =
+  @   head == \null || (is_valid_linked_list(head->next) && \valid(head));
+  @*/
+
+/*@ requires \valid(head_ref) && is_valid_linked_list(*head_ref);
   @ ensures *head_ref != \null && (*head_ref)->data == new_data && (*head_ref)->next == \old(*head_ref) && \fresh(*head_ref, sizeof(Node));
   @ allocates *head_ref;
   @ assigns *head_ref;
@@ -35,15 +39,16 @@ void popHead(Node** head_ref) {
     (*head_ref) = nextHead;
 }
 
-/*@ requires start >= 0 && end >= 0;
+/*@ requires start >= 0 && end >= 0 && end < 300000;
   @ requires graph != \null && \valid(levels + (0..end));
   @ requires \forall int i; 0 <= i < end + 1  ==> \valid(graph[i] + (0..end));
   @ assigns levels[0..end];
   @ ensures \result == 0 || \result == 1;
+  @ ensures \forall int i; 0 <= i <= end && levels[i] == -1 ==> graph[start][i] == 0.0 && graph[i][start] == 0.0;
   @*/
 int BFS(int start, int end, double **graph, int *levels){
     /*@ loop invariant 0 <= i <= end + 1;
-      @ loop invariant \forall int j; 0 <= j < i ==> levels[j] != -1;
+      @ loop invariant \forall int j; 0 <= j < i ==> levels[j] == -1;
       @ loop assigns i, levels[0..end];
       @ loop variant end - i;
       @*/
@@ -81,7 +86,7 @@ int BFS(int start, int end, double **graph, int *levels){
     }
 }
 
-/*@ requires 0 <= cur <= end && end >= 0;
+/*@ requires 0 <= cur <= end && end >= 0 && end <= 300000;
   @ requires graph != \null && \valid(levels + (0..end)) && \valid(visited + (0..end));
   @ requires \forall int i; 0 <= i <= end ==> \valid(graph[i] + (0..end));
   @ assigns visited[0..end], graph[0..end][0..end];
@@ -110,9 +115,20 @@ double DFS (int cur, int end, double **graph, int *levels, int *visited, double
     return 0;
 }
 
+/*@ axiomatic MemoryManagement {
+  @   logic boolean is_valid_graph(double **graph, integer nodes) reads graph[0][0..nodes-1];
+  @   axiom valid_graph:
+  @     \forall double **graph, integer nodes;
+  @       0 <= nodes && nodes <= 300000 ==>
+  @         is_valid_graph(graph, nodes) <==> (
+  @           \valid(graph) &&
+  @           \forall int i; 0 <= i < nodes ==> \valid(graph[i] + (0..nodes-1))
+  @         );
+  @ }
+  @*/
+
 /*@ requires \valid(graph);
-  @ ensures \result == -1 || (2 <= \result && \result <= 300000);
-  @ ensures *graph != \null && \forall int i; 0 <= i < \result ==> \valid((*graph)[i] + (0..\result-1));
+  @ ensures (\result == -1 && *graph == \null) || (\result >= 2 && \result <= 300000 && is_valid_graph(*graph, \result));
   @ assigns *graph;
   @ allocates graph;
   @ allocates graph[0..\result-1];
@@ -196,7 +212,7 @@ int readInput( double ***graph)
   @ ensures valid_flow(\result, graph, nodes);
   @ assigns graph[0..nodes-1][0..nodes-1];
   @*/
-double dinicMaxFlow(int nodes, double **graph){
+double dinitzMaxFlow(int nodes, double **graph){
     double total = 0;
     int *vertexLevel = malloc(nodes * sizeof(int));
     int *visited;
@@ -225,7 +241,12 @@ double dinicMaxFlow(int nodes, double **graph){
             }
             else
             {
-                total += flow;
+                if (DBL_MAX - total < flow) {
+                    total = DBL_MAX;
+                    break;
+                } else {
+                    total += flow;
+                }
             }
             // Reset visited array
             memset(visited, 0, nodes * sizeof(int));
@@ -244,14 +265,22 @@ double dinicMaxFlow(int nodes, double **graph){
   @ requires nodes > 0;
   @ requires \valid(graph + (0 .. nodes-1));
   @ requires \forall int i; 0 <= i < nodes ==> \valid(graph[i] + (0..nodes-1));
+  @ ensures \forall integer i; 0 <= i < nodes ==> graph[i] == \null;
   @ ensures !\valid(graph) && (\forall int i; 0 <= i < nodes ==> !\valid(graph[i]));
   @ frees graph[0..nodes-1];
   @ frees graph;
-  @ assigns \nothing;
+  @ assigns graph[0..nodes-1];
   @*/
 void cleanUp(int nodes, double** graph){
-    for (int i = 0; i < nodes; i++)
+    /*@ loop invariant 0 <= i <= nodes;
+      @ loop invariant \forall integer k; 0 <= k < i ==> graph[k] == \null;
+      @ loop assigns i, graph[0..nodes-1];
+      @ loop variant nodes - i;
+      @*/
+    for (int i = 0; i < nodes; i++){
         free(graph[i]);
+        graph[i] = NULL;
+    }
     free(graph);
 }
 
@@ -267,7 +296,7 @@ int main()
     if(nodes == -1){
         return 1;
     }
-    total = dinicMaxFlow(nodes, graph);
+    total = dinitzMaxFlow(nodes, graph);
     if (total == -1.0)
     {
         cleanUp(nodes, graph);
@@ -279,4 +308,4 @@ int main()
         cleanUp(nodes, graph);
     }
     return 0;
-}
\ No newline at end of file
+}