[ETCM-533] return root nodeas part of the proof, add tests

Author: bsuieric

src/main/scala/io/iohk/ethereum/mpt/MerklePatriciaTrie.scala

@@ -105,7 +105,7 @@ class MerklePatriciaTrie[K, V] private (private[mpt] val rootNode: Option[MptNod
   def getProof(key: K): Option[Vector[MptNode]] = {
     pathTraverse[Vector[MptNode]](Vector.empty, mkKeyNibbles(key)) { case (acc, node) =>
       node match {
-        case Some(nextNodeOnExt @ (_: BranchNode | _: ExtensionNode | _: LeafNode)) => acc :+ nextNodeOnExt
+        case Some(nextNodeOnExt @ (_: BranchNode | _: ExtensionNode | _: LeafNode | _: HashNode)) => acc :+ nextNodeOnExt
         case _ => acc
       }
     }
@@ -155,7 +155,7 @@ class MerklePatriciaTrie[K, V] private (private[mpt] val rootNode: Option[MptNod
 
     rootNode match {
       case Some(root) =>
-        pathTraverse(acc, root, searchKey, op)
+        pathTraverse(op(acc, Some(root)), root, searchKey, op)
       case None =>
         None
     }

src/test/scala/io/iohk/ethereum/ObjectGenerators.scala

@@ -50,9 +50,10 @@ trait ObjectGenerators {
     } yield (aByteList.toArray, t)
   }
 
-  def keyValueListGen(): Gen[List[(Int, Int)]] = {
+  def keyValueListGen(minValue: Int = Int.MinValue, maxValue: Int = Int.MaxValue): Gen[List[(Int, Int)]] = {
     for {
-      aKeyList <- Gen.nonEmptyListOf(Arbitrary.arbitrary[Int]).map(_.distinct)
+      values <- Gen.chooseNum(minValue, maxValue)
+      aKeyList <- Gen.nonEmptyListOf(values).map(_.distinct)
     } yield aKeyList.zip(aKeyList)
   }
 

src/test/scala/io/iohk/ethereum/mpt/MerklePatriciaTrieSuite.scala

@@ -554,7 +554,40 @@ class MerklePatriciaTrieSuite extends AnyFunSuite with ScalaCheckPropertyChecks
     assert(proof.isEmpty)
   }
 
-  test("getProof returns proof result for non-existing address") {
+  test("PatriciaTrie can get proof(at least the root node) for all inserted key-value pairs") {
+    forAll(keyValueListGen()) { keyValueList: Seq[(Int, Int)] =>
+      val trie = addEveryKeyValuePair(keyValueList)
+      assertCanGetProofForEveryKeyValue(trie, keyValueList)
+    }
+  }
+
+  test("PatriciaTrie return root as proof when no common nibbles are found between MPT root hash and search key") {
+    forAll(keyValueListGen(1, 10)) { keyValueList: Seq[(Int, Int)] =>
+      val trie = addEveryKeyValuePair(keyValueList)
+      val wrongKey = 22
+      val proof = trie.getProof(wrongKey)
+      assert(proof.getOrElse(Vector.empty).toList match {
+        case _ @ HashNode(_) :: Nil => true
+        case _ => false
+      })
+    }
+  }
+
+  test("PatriciaTrie return proof when having all nibbles in common except the last one between MPT root hash and search key") {
+
+    val key = 1111
+    val wrongKey = 1112
+    val emptyTrie = MerklePatriciaTrie[Int, Int](emptyEphemNodeStorage)
+      .put(key, 1)
+      .put(wrongKey, 2)
+    val proof = emptyTrie.getProof(key = wrongKey)
+    assert(proof.getOrElse(Vector.empty).toList match {
+      case _ @ HashNode(_) :: tail => tail.nonEmpty
+      case _ => false
+    })
+  }
+
+  test("getProof returns proof result for non-existing key") {
       // given
       val EmptyTrie = MerklePatriciaTrie[Array[Byte], Array[Byte]](emptyEphemNodeStorage)
       val key1: Array[Byte] = Hex.decode("10000001")
@@ -611,6 +644,12 @@ class MerklePatriciaTrieSuite extends AnyFunSuite with ScalaCheckPropertyChecks
       assert(obtained.get == value)
     }
 
+  private def assertCanGetProofForEveryKeyValue[K, V](trie: MerklePatriciaTrie[K, V], kvs: Seq[(K, V)]): Unit =
+    kvs.foreach { case (key, _) =>
+      val obtained = trie.getProof(key)
+      assert(obtained.getOrElse(Vector.empty).nonEmpty)
+    }
+
   private def assertCanGetEveryKeyValues[K, V](trie: MerklePatriciaTrie[K, Array[V]], kvs: List[(K, Array[V])]): Unit =
     kvs.foreach { case (key, value) =>
       val obtained = trie.get(key)