From c320411a4092508f8113698177c19833e70fee93 Mon Sep 17 00:00:00 2001
From: Adrian Kuta <adrian.kuta93@gmail.com>
Date: Tue, 21 Jan 2020 15:38:44 +0100
Subject: [PATCH] `child` function is now returning TreeNode

---
 .../tree/ChildDeclarationInterface.kt            |  3 ++-
 .../adriankuta/datastructure/tree/TreeNode.kt    |  3 ++-
 .../datastructure/tree/TreeNodeTest.kt           | 16 ----------------
 3 files changed, 4 insertions(+), 18 deletions(-)

diff --git a/treedatastructure/src/main/java/com/github/adriankuta/datastructure/tree/ChildDeclarationInterface.kt b/treedatastructure/src/main/java/com/github/adriankuta/datastructure/tree/ChildDeclarationInterface.kt
index 8d31cd9..fc0888a 100644
--- a/treedatastructure/src/main/java/com/github/adriankuta/datastructure/tree/ChildDeclarationInterface.kt
+++ b/treedatastructure/src/main/java/com/github/adriankuta/datastructure/tree/ChildDeclarationInterface.kt
@@ -15,7 +15,8 @@ interface ChildDeclarationInterface<T> {
      *     }
      * }
      * ```
+     * @return New created TreeNode.
      */
     @JvmSynthetic
-    fun child(value: T, childDeclaration: ChildDeclaration<T>? = null)
+    fun child(value: T, childDeclaration: ChildDeclaration<T>? = null): TreeNode<T>
 }
\ No newline at end of file
diff --git a/treedatastructure/src/main/java/com/github/adriankuta/datastructure/tree/TreeNode.kt b/treedatastructure/src/main/java/com/github/adriankuta/datastructure/tree/TreeNode.kt
index 4f8461f..1d19079 100644
--- a/treedatastructure/src/main/java/com/github/adriankuta/datastructure/tree/TreeNode.kt
+++ b/treedatastructure/src/main/java/com/github/adriankuta/datastructure/tree/TreeNode.kt
@@ -27,11 +27,12 @@ open class TreeNode<T>(val value: T) : Iterable<TreeNode<T>>, ChildDeclarationIn
     }
 
     @JvmSynthetic
-    override fun child(value: T, childDeclaration: ChildDeclaration<T>?) {
+    override fun child(value: T, childDeclaration: ChildDeclaration<T>?): TreeNode<T> {
         val newChild = TreeNode(value)
         if(childDeclaration != null)
             newChild.childDeclaration()
         _children.add(newChild)
+        return newChild
     }
 
     /**
diff --git a/treedatastructure/src/test/java/com/github/adriankuta/datastructure/tree/TreeNodeTest.kt b/treedatastructure/src/test/java/com/github/adriankuta/datastructure/tree/TreeNodeTest.kt
index c1f17eb..1d455e1 100644
--- a/treedatastructure/src/test/java/com/github/adriankuta/datastructure/tree/TreeNodeTest.kt
+++ b/treedatastructure/src/test/java/com/github/adriankuta/datastructure/tree/TreeNodeTest.kt
@@ -139,20 +139,4 @@ class TreeNodeTest {
         }
         assertEquals(root.prettyString(), rootExt.prettyString())
     }
-
-//    @Test
-//    fun getVisibleNodes() {
-//        val root = tree("World") {
-//            child("Level 1") {
-//                child("Level 2") {
-//                    child("Level 3") {
-//                        child("Level 4")
-//                    }
-//                }
-//            }
-//        }
-//
-//        print(root.prettyString())
-//    }
-
 }
\ No newline at end of file