From fc46219ca74e64868a5383e1bf54da33ca1be8a6 Mon Sep 17 00:00:00 2001
From: Ivan Smirnov <ifsmirnov@yandex.ru>
Date: Sun, 15 Oct 2017 11:32:38 +0300
Subject: [PATCH] Support printing tree in 'parents' format

---
 doc/printers.md          |   8 ++-
 doc/tree.md              |   2 +-
 generic_graph.h          |   6 +-
 impl/generic_graph_inl.h |  32 +++++-----
 impl/tree_inl.h          |  53 ++++++++++++++++-
 jngen.h                  | 124 +++++++++++++++++++++++++++++----------
 repr.h                   |  24 ++++++--
 tests/Makefile           |   2 +-
 tests/tree.cpp           |  25 ++++++++
 tree.h                   |   9 ++-
 10 files changed, 216 insertions(+), 69 deletions(-)

diff --git a/doc/printers.md b/doc/printers.md
index 3120d03..c1e237b 100644
--- a/doc/printers.md
+++ b/doc/printers.md
@@ -64,11 +64,13 @@ Note that Jngen does not interact with stl-defined operators. That mean that wri
 * Action: print number of edges of a graph.
 * Default: unset.
 #### printEdges(bool value = true)
-* Action: when printing a tree, print a list of edges. Opposite to *printParents*.
+* Action: when printing a tree, print a list of edges.
 * Default: set.
-#### printParents(bool value = true)
+#### printParents(int value = -1)
 * Action: when printing a tree, print a parent of each vertex. Opposite to *printEdges*.
-* Default: unset.
+* Arguments: *value* stands for the root of the tree. If *value* is *0* or greater, then the parent of each vertex is printed, having root's parent as
+    *-1* (*0* if *add1()* is present). *value = -1* is a special value: in this case tree is rooted at *0* and its parent is not printed (printing *n-1* values in total).
+* Note: this option and *printEdges* cancel each other.
 #### endl(bool value = true)
 * Action: separate elements of the array with line breaks instead of spaces.
 * Default: unset.
diff --git a/doc/tree.md b/doc/tree.md
index 2193e56..a81f39d 100644
--- a/doc/tree.md
+++ b/doc/tree.md
@@ -48,7 +48,7 @@ Note that all generators return trees with sorted edges to make tests more human
     * randomly swap egdes' endpoints.
 
 #### Array parents(int root) const
-* Returns: array of size *n*, where *i*-th element is a parent of vertex *i* if the tree is rooted at *root*. Parent of *root* is *root* itself.
+* Returns: array of size *n*, where *i*-th element is a parent of vertex *i* if the tree is rooted at *root*. Parent of *root* is *-1*.
 
 #### Tree link(int vInThis, const Tree& other, int vInOther)
 * Returns: a tree made of _*this_ and *other*, with an extra edge between two vertices with ids *vInThis* and *vInOther*, respectively.
diff --git a/generic_graph.h b/generic_graph.h
index 66c437c..40b85d0 100644
--- a/generic_graph.h
+++ b/generic_graph.h
@@ -25,8 +25,8 @@ class GenericGraph {
     virtual void addEdge(int u, int v, const Weight& w = Weight{});
     virtual bool isConnected() const { return dsu_.isConnected(); }
 
-    virtual int vertexLabel(int v) const { return vertexLabel_[v]; }
-    virtual int vertexByLabel(int v) const { return vertexByLabel_[v]; }
+    virtual int vertexLabel(int v) const { return vertexLabel_.at(v); }
+    virtual int vertexByLabel(int v) const { return vertexByLabel_.at(v); }
 
     // v: label
     // return: array<label>
@@ -60,6 +60,8 @@ class GenericGraph {
     virtual bool operator>=(const GenericGraph& other) const;
 
 protected:
+    static WeightArray prepareWeightArray(WeightArray a, int requiredSize);
+
     void doShuffle();
 
     void extend(size_t size);
diff --git a/impl/generic_graph_inl.h b/impl/generic_graph_inl.h
index b1fc862..16bd3ec 100644
--- a/impl/generic_graph_inl.h
+++ b/impl/generic_graph_inl.h
@@ -75,6 +75,20 @@ Arrayp GenericGraph::edges() const {
     return edges;
 }
 
+WeightArray GenericGraph::prepareWeightArray(WeightArray a, int requiredSize) {
+    ENSURE(a.hasNonEmpty(), "Attempt to print empty weight array");
+
+    a.extend(requiredSize);
+    int type = a.anyType();
+    for (auto& x: a) {
+        if (x.empty()) {
+            x.setType(type);
+        }
+    }
+
+    return a;
+}
+
 void GenericGraph::doShuffle() {
     // this if is to be removed after all checks pass
     if (vertexLabel_.size() < static_cast<size_t>(n())) {
@@ -173,24 +187,6 @@ void GenericGraph::addEdge(int u, int v, const Weight& w) {
     }
 }
 
-namespace {
-
-WeightArray prepareWeightArray(WeightArray a, int requiredSize) {
-    ENSURE(a.hasNonEmpty(), "Attempt to print empty weight array");
-
-    a.extend(requiredSize);
-    int type = a.anyType();
-    for (auto& x: a) {
-        if (x.empty()) {
-            x.setType(type);
-        }
-    }
-
-    return a;
-}
-
-} // namespace
-
 void GenericGraph::doPrintEdges(
     std::ostream& out, const OutputModifier& mod) const
 {
diff --git a/impl/tree_inl.h b/impl/tree_inl.h
index 6c106f6..0beee21 100644
--- a/impl/tree_inl.h
+++ b/impl/tree_inl.h
@@ -30,7 +30,7 @@ Array Tree::parents(int root) const {
     root = vertexByLabel(root);
 
     Array parents(n());
-    parents[root] = root;
+    parents[root] = -1;
     std::vector<int> used(n());
     std::vector<int> queue{root};
     for (size_t i = 0; i < queue.size(); ++i) {
@@ -45,7 +45,9 @@ Array Tree::parents(int root) const {
     }
 
     for (auto& x: parents) {
-        x = vertexLabel(x);
+        if (x != -1) {
+            x = vertexLabel(x);
+        }
     }
 
     return parents;
@@ -218,3 +220,50 @@ Tree Tree::kary(int size, int k) {
     return t;
 }
 
+void Tree::doPrintParents(std::ostream& out, const OutputModifier& mod) const {
+    int root = mod.printParents;
+    if (root == -1) {
+        root = 0;
+    }
+
+    auto parents = this->parents(root);
+    if (mod.printParents == -1) {
+        parents.erase(parents.begin());
+    }
+
+    if (mod.printN) {
+        out << n() << "\n";
+    }
+
+    // TODO: avoid copy-paste from doPrintEdges
+    if (mod.printWeights && vertexWeights_.hasNonEmpty()) {
+        auto vertexWeights = prepareWeightArray(vertexWeights_, n());
+        for (int i = 0; i < n(); ++i) {
+            if (i > 0) {
+                out << " ";
+            }
+            JNGEN_PRINT_NO_MOD(vertexWeights[vertexByLabel(i)]);
+        }
+        out << "\n";
+    }
+
+    auto t(mod);
+    {
+        auto mod(t);
+        mod.printN = false;
+
+        if (mod.printWeights && edgeWeights_.hasNonEmpty()) {
+            ensure(false, "Printing parents and edge weights is not supported");
+            ensure(
+                mod.printParents == -1,
+                "Root must not be set to any exact value when printing a tree "
+                "with edge weights. To fix it, either set printParents() "
+                "or printWeights(false)");
+            ENSURE(root == 0);
+            // TODO: some code to be here
+        } else {
+            JNGEN_PRINT(parents);
+        }
+    }
+}
+
diff --git a/jngen.h b/jngen.h
index ac14e75..398473d 100644
--- a/jngen.h
+++ b/jngen.h
@@ -2388,8 +2388,9 @@ struct OutputModifier {
     bool printN = false;
     bool printM = false;
 
-    bool printParents = false;
+    int printParents;
     bool printEdges = true;
+    bool printWeights = true;
 
     char sep = ' ';
 };
@@ -2437,15 +2438,22 @@ class Repr {
         return *this;
     }
 
-    Repr<T>& printParents(bool value = true) {
+    Repr<T>& printParents(int value = -1) {
         mod_.printParents = value;
-        mod_.printEdges = !value;
+        mod_.printEdges = false;
         return *this;
     }
 
     Repr<T>& printEdges(bool value = true) {
         mod_.printEdges = value;
-        mod_.printParents = !value;
+        if (!value) {
+            mod_.printParents = -1;
+        }
+        return *this;
+    }
+
+    Repr<T>& printWeights(bool value = true) {
+        mod_.printWeights = value;
         return *this;
     }
 
@@ -2493,7 +2501,7 @@ class ReprProxy : public BaseReprProxy {
         return repr;
     }
 
-    Repr<T> printParents(bool value = true) {
+    Repr<T> printParents(int value = -1) {
         Repr<T> repr(static_cast<const T&>(*this));
         repr.printParents(value);
         return repr;
@@ -2505,6 +2513,12 @@ class ReprProxy : public BaseReprProxy {
         return repr;
     }
 
+    Repr<T> printWeights(bool value = true) {
+        Repr<T> repr(static_cast<const T&>(*this));
+        repr.printWeights(value);
+        return repr;
+    }
+
     Repr<T> endl(bool value = true) {
         Repr<T> repr(static_cast<const T&>(*this));
         repr.endl(value);
@@ -5219,8 +5233,8 @@ class GenericGraph {
     virtual void addEdge(int u, int v, const Weight& w = Weight{});
     virtual bool isConnected() const { return dsu_.isConnected(); }
 
-    virtual int vertexLabel(int v) const { return vertexLabel_[v]; }
-    virtual int vertexByLabel(int v) const { return vertexByLabel_[v]; }
+    virtual int vertexLabel(int v) const { return vertexLabel_.at(v); }
+    virtual int vertexByLabel(int v) const { return vertexByLabel_.at(v); }
 
     // v: label
     // return: array<label>
@@ -5254,6 +5268,8 @@ class GenericGraph {
     virtual bool operator>=(const GenericGraph& other) const;
 
 protected:
+    static WeightArray prepareWeightArray(WeightArray a, int requiredSize);
+
     void doShuffle();
 
     void extend(size_t size);
@@ -5388,6 +5404,20 @@ Arrayp GenericGraph::edges() const {
     return edges;
 }
 
+WeightArray GenericGraph::prepareWeightArray(WeightArray a, int requiredSize) {
+    ENSURE(a.hasNonEmpty(), "Attempt to print empty weight array");
+
+    a.extend(requiredSize);
+    int type = a.anyType();
+    for (auto& x: a) {
+        if (x.empty()) {
+            x.setType(type);
+        }
+    }
+
+    return a;
+}
+
 void GenericGraph::doShuffle() {
     // this if is to be removed after all checks pass
     if (vertexLabel_.size() < static_cast<size_t>(n())) {
@@ -5486,24 +5516,6 @@ void GenericGraph::addEdge(int u, int v, const Weight& w) {
     }
 }
 
-namespace {
-
-WeightArray prepareWeightArray(WeightArray a, int requiredSize) {
-    ENSURE(a.hasNonEmpty(), "Attempt to print empty weight array");
-
-    a.extend(requiredSize);
-    int type = a.anyType();
-    for (auto& x: a) {
-        if (x.empty()) {
-            x.setType(type);
-        }
-    }
-
-    return a;
-}
-
-} // namespace
-
 void GenericGraph::doPrintEdges(
     std::ostream& out, const OutputModifier& mod) const
 {
@@ -5649,18 +5661,17 @@ class Tree : public ReprProxy<Tree>, public GenericGraph {
     static Tree caterpillar(int size, int length);
     static Tree binary(int size);
     static Tree kary(int size, int k);
+
+    void doPrintParents(std::ostream& out, const OutputModifier& mod) const;
 };
 
 JNGEN_DECLARE_SIMPLE_PRINTER(Tree, 2) {
     ensure(t.isConnected(), "Cannot print a tree: it is not connected");
 
-    if (mod.printParents) {
-        ensure(false, "Printing parents is not implemented");
-    } else if (mod.printEdges) {
+    if (mod.printEdges) {
         t.doPrintEdges(out, mod);
     } else {
-        ensure(false, "Print mode is not set, select one of 'printParents'"
-            " or 'printEdges'");
+        t.doPrintParents(out, mod);
     }
 }
 
@@ -5710,7 +5721,7 @@ Array Tree::parents(int root) const {
     root = vertexByLabel(root);
 
     Array parents(n());
-    parents[root] = root;
+    parents[root] = -1;
     std::vector<int> used(n());
     std::vector<int> queue{root};
     for (size_t i = 0; i < queue.size(); ++i) {
@@ -5725,7 +5736,9 @@ Array Tree::parents(int root) const {
     }
 
     for (auto& x: parents) {
-        x = vertexLabel(x);
+        if (x != -1) {
+            x = vertexLabel(x);
+        }
     }
 
     return parents;
@@ -5898,6 +5911,53 @@ Tree Tree::kary(int size, int k) {
     return t;
 }
 
+void Tree::doPrintParents(std::ostream& out, const OutputModifier& mod) const {
+    int root = mod.printParents;
+    if (root == -1) {
+        root = 0;
+    }
+
+    auto parents = this->parents(root);
+    if (mod.printParents == -1) {
+        parents.erase(parents.begin());
+    }
+
+    if (mod.printN) {
+        out << n() << "\n";
+    }
+
+    // TODO: avoid copy-paste from doPrintEdges
+    if (mod.printWeights && vertexWeights_.hasNonEmpty()) {
+        auto vertexWeights = prepareWeightArray(vertexWeights_, n());
+        for (int i = 0; i < n(); ++i) {
+            if (i > 0) {
+                out << " ";
+            }
+            JNGEN_PRINT_NO_MOD(vertexWeights[vertexByLabel(i)]);
+        }
+        out << "\n";
+    }
+
+    auto t(mod);
+    {
+        auto mod(t);
+        mod.printN = false;
+
+        if (mod.printWeights && edgeWeights_.hasNonEmpty()) {
+            ensure(false, "Printing parents and edge weights is not supported");
+            ensure(
+                mod.printParents == -1,
+                "Root must not be set to any exact value when printing a tree "
+                "with edge weights. To fix it, either set printParents() "
+                "or printWeights(false)");
+            ENSURE(root == 0);
+            // TODO: some code to be here
+        } else {
+            JNGEN_PRINT(parents);
+        }
+    }
+}
+
 #undef JNGEN_INCLUDE_TREE_INL_H
 #endif // JNGEN_DECLARE_ONLY
 
diff --git a/repr.h b/repr.h
index 4c63c27..471bc8a 100644
--- a/repr.h
+++ b/repr.h
@@ -16,8 +16,9 @@ struct OutputModifier {
     bool printN = false;
     bool printM = false;
 
-    bool printParents = false;
+    int printParents;
     bool printEdges = true;
+    bool printWeights = true;
 
     char sep = ' ';
 };
@@ -65,15 +66,22 @@ class Repr {
         return *this;
     }
 
-    Repr<T>& printParents(bool value = true) {
+    Repr<T>& printParents(int value = -1) {
         mod_.printParents = value;
-        mod_.printEdges = !value;
+        mod_.printEdges = false;
         return *this;
     }
 
     Repr<T>& printEdges(bool value = true) {
         mod_.printEdges = value;
-        mod_.printParents = !value;
+        if (!value) {
+            mod_.printParents = -1;
+        }
+        return *this;
+    }
+
+    Repr<T>& printWeights(bool value = true) {
+        mod_.printWeights = value;
         return *this;
     }
 
@@ -121,7 +129,7 @@ class ReprProxy : public BaseReprProxy {
         return repr;
     }
 
-    Repr<T> printParents(bool value = true) {
+    Repr<T> printParents(int value = -1) {
         Repr<T> repr(static_cast<const T&>(*this));
         repr.printParents(value);
         return repr;
@@ -133,6 +141,12 @@ class ReprProxy : public BaseReprProxy {
         return repr;
     }
 
+    Repr<T> printWeights(bool value = true) {
+        Repr<T> repr(static_cast<const T&>(*this));
+        repr.printWeights(value);
+        return repr;
+    }
+
     Repr<T> endl(bool value = true) {
         Repr<T> repr(static_cast<const T&>(*this));
         repr.endl(value);
diff --git a/tests/Makefile b/tests/Makefile
index 998adb3..eb411ad 100644
--- a/tests/Makefile
+++ b/tests/Makefile
@@ -1,4 +1,4 @@
-CXXFLAGS = -O2 -Wall -Wextra -Werror -std=c++14
+CXXFLAGS = -O2 -Wall -Wextra -Werror -std=c++11
 CXXFLAGS += -DJNGEN_DECLARE_ONLY
 LDFLAGS += -lboost_system -lboost_unit_test_framework
 
diff --git a/tests/tree.cpp b/tests/tree.cpp
index 341f724..bd49a4b 100644
--- a/tests/tree.cpp
+++ b/tests/tree.cpp
@@ -221,6 +221,31 @@ BOOST_AUTO_TEST_CASE(check_glue) {
     checkEquals(glued, t);
 }
 
+BOOST_AUTO_TEST_CASE(print_parents) {
+    rnd.seed(123);
+
+    std::ostringstream ss;
+
+    const std::string res1 = "5\ngrt kzw bar pja oap\n5 1 3 3\n";
+
+    Tree t = Tree::random(5);
+    t.setVertexWeights(TArray<std::string>::random(t.n(), "[a-z]{%d}", 3));
+    ss << t.add1().printN().printParents() << std::endl;
+
+    BOOST_TEST(ss.str() == res1);
+
+    t.setVertexWeights(WeightArray(t.n()));
+
+    ss.str("");
+    ss << t.printParents(0);
+    BOOST_TEST(ss.str() == "-1 4 0 2 2");
+
+    ss.str("");
+    ss << t.printParents(1);
+    BOOST_TEST(ss.str() == "2 -1 4 2 1");
+
+}
+
 // TODO: add tests to check random generators exactly
 
 BOOST_AUTO_TEST_SUITE_END()
diff --git a/tree.h b/tree.h
index f2fe234..e388c0f 100644
--- a/tree.h
+++ b/tree.h
@@ -38,18 +38,17 @@ class Tree : public ReprProxy<Tree>, public GenericGraph {
     static Tree caterpillar(int size, int length);
     static Tree binary(int size);
     static Tree kary(int size, int k);
+
+    void doPrintParents(std::ostream& out, const OutputModifier& mod) const;
 };
 
 JNGEN_DECLARE_SIMPLE_PRINTER(Tree, 2) {
     ensure(t.isConnected(), "Cannot print a tree: it is not connected");
 
-    if (mod.printParents) {
-        ensure(false, "Printing parents is not implemented");
-    } else if (mod.printEdges) {
+    if (mod.printEdges) {
         t.doPrintEdges(out, mod);
     } else {
-        ensure(false, "Print mode is not set, select one of 'printParents'"
-            " or 'printEdges'");
+        t.doPrintParents(out, mod);
     }
 }