diff --git a/src/smt/theory_diff_logic.h b/src/smt/theory_diff_logic.h
index 1a7d2baac..f904a600a 100644
--- a/src/smt/theory_diff_logic.h
+++ b/src/smt/theory_diff_logic.h
@@ -397,33 +397,54 @@ namespace smt {
         // Denote costs c_ij on edge (i, j)
         vector<rational> m_costs;
         // Denote supply/demand b_i on node i
-        vector<rational> m_capacities;
+        vector<rational> m_potentials;
 
         // Keep optimal solution of the min cost flow problem
         inf_rational m_objective;
 
+        // Data structure of spanning trees
+        svector<dl_var> m_pred;
+        svector<int> m_depth;
+        svector<dl_var> m_thread;
+
     public:
-        // Create a spanning tree using Kruskal algorithm
-        virtual svector<edge_id> & create_spanning_tree();
+        // Initialize the network with a feasible spanning tree
+        virtual void initialize();
+
+        virtual void compute_potentials();
+
+        virtual void compute_flows();
             
-        // A spanning tree is a basis in network simplex.
-        // Check whether the edges' associated costs could be reduced
-        virtual rational calculate_reduced_costs(svector<edge_id> & spanning_tree);
-
         // If all reduced costs are non-negative, the current flow is optimal
-        virtual bool is_optimal(svector<edge_id> & spanning_tree);
-
-        // Choose an edge with negative reduced cost
-        virtual edge_id choose_entering_edge();
+        // If not optimal, return a violated edge in the corresponding variable
+        virtual bool is_optimal(edge_id & violated_edge);
 
         // Send as much flow as possible around the cycle, the first basic edge with flow 0 will leave
-        edge_id choose_leaving_edge();
+        virtual edge_id choose_leaving_edge();
+
+        virtual void update_tree(edge_id entering_edge, edge_id leaving_edge);
 
         virtual bool is_unbounded();
 
+        // Compute the optimal solution
+        virtual void compute_solution();
+
         // Minimize cost flows
         // Return true if found an optimal solution, and return false if unbounded
-        virtual bool minimize();
+        bool minimize() {
+            initialize();
+            compute_potentials();
+            compute_flows();
+            edge_id entering_edge;
+            while (!is_optimal(entering_edge)) {
+                edge_id leaving_edge = choose_leaving_edge();
+                update_tree(entering_edge, leaving_edge);
+
+                if (is_unbounded())
+                    return false;
+            }
+            return true;
+        }
     };
 
     /* Notes:
@@ -432,7 +453,7 @@ namespace smt {
         and another function to convert from min cost flow solution to DL solution.
 
         It remains unclear how to convert DL assignment to a basic feasible solution of Network Simplex.
-        A naive approach is to run Kruskal in order to get a spanning tree.
+        A naive approach is to run an algorithm on max flow in order to get a spanning tree.
 
         The network_simplex class hasn't had multiple pivoting strategies yet.
     */
diff --git a/src/smt/theory_diff_logic_def.h b/src/smt/theory_diff_logic_def.h
index 5f9e80d9b..684850d77 100644
--- a/src/smt/theory_diff_logic_def.h
+++ b/src/smt/theory_diff_logic_def.h
@@ -1038,8 +1038,8 @@ void theory_diff_logic<Ext>::internalize_objective(app * n, objective_term & obj
     // Compile term into objective_term format
     rational r;
     if (m_util.is_numeral(n, r)) {
-	    theory_var v = mk_num(n, r);
-	    objective.push_back(std::make_pair(v, r));
+        theory_var v = mk_num(n, r);
+        objective.push_back(std::make_pair(v, r));
     }
     else if (m_util.is_add(n)) {
         for (unsigned i = 0; i < n->get_num_args(); ++i) {
@@ -1047,22 +1047,23 @@ void theory_diff_logic<Ext>::internalize_objective(app * n, objective_term & obj
         };
     }
     else if (m_util.is_mul(n)) {
-	    SASSERT(n->get_num_args() == 2);
-	    rational r;
-	    app * lhs = to_app(n->get_arg(0));
-	    app * rhs = to_app(n->get_arg(1));
-	    SASSERT(m_util.is_numeral(lhs, r) || m_util.is_numeral(rhs, r));
+        SASSERT(n->get_num_args() == 2);
+        rational r;
+        app * lhs = to_app(n->get_arg(0));
+        app * rhs = to_app(n->get_arg(1));
+        SASSERT(m_util.is_numeral(lhs, r) || m_util.is_numeral(rhs, r));
 
         if (!m_util.is_numeral(lhs, r))
-		    std::swap(lhs, rhs);
-	    m_util.is_numeral(lhs, r);
-	    theory_var v = mk_var(rhs);
-	    objective.push_back(std::make_pair(v, r));
+	        std::swap(lhs, rhs);
+
+        m_util.is_numeral(lhs, r);
+        theory_var v = mk_var(rhs);
+        objective.push_back(std::make_pair(v, r));
     }
     else {
-	    theory_var v = mk_var(n);
-	    rational r(1);
-	    objective.push_back(std::make_pair(v, r));
+        theory_var v = mk_var(n);
+        rational r(1);
+        objective.push_back(std::make_pair(v, r));
     }
 }