ng_netflow: use small_color_map

diff --git a/src/nfagraph/ng_netflow.cpp b/src/nfagraph/ng_netflow.cpp
index cff263586df97ed2011597b62e5b3767ffa41591..780a319f5de5af73c7d20f80bb3428ee5948c5b0 100644 (file)
--- a/src/nfagraph/ng_netflow.cpp
+++ b/src/nfagraph/ng_netflow.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015-2016, Intel Corporation
+ * Copyright (c) 2015-2017, Intel Corporation
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
@@ -37,6 +37,7 @@
 #include "ue2common.h"
 #include "util/container.h"
 #include "util/graph_range.h"
+#include "util/graph_small_color_map.h"
 
 #include <algorithm>
 #include <boost/graph/boykov_kolmogorov_max_flow.hpp>
@@ -118,7 +119,7 @@ void removeEdgesFromIndex(NGHolder &g, vector<u64a> &capacityMap, u32 idx) {
  * colour map (from which we can find the min cut). */
 static
 u64a getMaxFlow(NGHolder &h, const vector<u64a> &capacityMap_in,
-                vector<default_color_type> &colorMap) {
+                decltype(make_small_color_map(NGHolder())) &colorMap) {
     vector<u64a> capacityMap = capacityMap_in;
     NFAVertex src = h.start;
     NFAVertex sink = h.acceptEod;
@@ -141,7 +142,6 @@ u64a getMaxFlow(NGHolder &h, const vector<u64a> &capacityMap_in,
     vector<u64a> edgeResiduals(numTotalEdges);
     vector<NFAEdge> predecessors(numVertices);
     vector<s32> distances(numVertices);
-    assert(colorMap.size() == numVertices);
 
     auto v_index_map = get(vertex_index, h);
     auto e_index_map = get(edge_index, h);
@@ -151,7 +151,7 @@ u64a getMaxFlow(NGHolder &h, const vector<u64a> &capacityMap_in,
          make_iterator_property_map(edgeResiduals.begin(), e_index_map),
          make_iterator_property_map(reverseEdges.begin(), e_index_map),
          make_iterator_property_map(predecessors.begin(), v_index_map),
-         make_iterator_property_map(colorMap.begin(), v_index_map),
+         colorMap,
          make_iterator_property_map(distances.begin(), v_index_map),
          v_index_map,
          src, sink);
@@ -169,8 +169,8 @@ vector<NFAEdge> findMinCut(NGHolder &h, const vector<u64a> &scores) {
     assert(hasCorrectlyNumberedEdges(h));
     assert(hasCorrectlyNumberedVertices(h));
 
-    vector<default_color_type> colorMap(num_vertices(h));
-    u64a flow = getMaxFlow(h, scores, colorMap);
+    auto colors = make_small_color_map(h);
+    u64a flow = getMaxFlow(h, scores, colors);
 
     vector<NFAEdge> picked_white;
     vector<NFAEdge> picked_black;
@@ -185,17 +185,17 @@ vector<NFAEdge> findMinCut(NGHolder &h, const vector<u64a> &scores) {
             continue; // skips, among other things, reverse edges
         }
 
-        default_color_type fromColor = colorMap[h[from].index];
-        default_color_type toColor = colorMap[h[to].index];
+        auto fromColor = get(colors, from);
+        auto toColor = get(colors, to);
 
-        if (fromColor != boost::white_color && toColor == boost::white_color) {
+        if (fromColor != small_color::white && toColor == small_color::white) {
             assert(ec <= INVALID_EDGE_CAP);
             DEBUG_PRINTF("found white cut edge %zu->%zu cap %llu\n",
                      h[from].index, h[to].index, ec);
             observed_white_flow += ec;
             picked_white.push_back(e);
         }
-        if (fromColor == boost::black_color && toColor != boost::black_color) {
+        if (fromColor == small_color::black && toColor != small_color::black) {
             assert(ec <= INVALID_EDGE_CAP);
             DEBUG_PRINTF("found black cut edge %zu->%zu cap %llu\n",
                      h[from].index, h[to].index, ec);
@@ -206,7 +206,7 @@ vector<NFAEdge> findMinCut(NGHolder &h, const vector<u64a> &scores) {
 
     DEBUG_PRINTF("min flow = %llu b flow = %llu w flow %llu\n", flow,
                  observed_black_flow, observed_white_flow);
-    if (MIN(observed_white_flow, observed_black_flow) != flow) {
+    if (min(observed_white_flow, observed_black_flow) != flow) {
         DEBUG_PRINTF("bad cut\n");
     }