multisig: collecting PositionedLeafs

interpreter/shared/src/main/scala/sigmastate/NodePosition.scala

@@ -0,0 +1,41 @@
+package sigmastate
+
+/**
+  * Data type which encodes position of a node in a tree.
+  *
+  * Position is encoded like following (the example provided is for CTHRESHOLD(2, Seq(pk1, pk2, pk3 && pk4)) :
+  *
+  *      0
+  *    / | \
+  *   /  |  \
+  * 0-0 0-1 0-2
+  *         /|
+  *        / |
+  *       /  |
+  *      /   |
+  *  0-2-0  0-2-1
+  *
+  * So a hint associated with pk1 has a position "0-0", pk4 - "0-2-1" .
+  *
+  * Please note that "0" prefix is for a crypto tree. There are several kinds of trees during evaluation.
+  * Initial mixed tree (ergoTree) would have another prefix.
+  *
+  * @param positions - positions from root (inclusive) in top-down order
+  */
+case class NodePosition(positions: Seq[Int]) {
+  def child(childIdx: Int): NodePosition = NodePosition(positions :+ childIdx)
+  def ++(path: Seq[Int]): NodePosition = NodePosition(positions ++ path)
+  override def toString: String = positions.mkString("-")
+}
+
+object NodePosition {
+  /**
+    * Prefix to encode node positions in a crypto tree.
+    */
+  val CryptoTreePrefix = NodePosition(Seq(0))
+
+  /**
+    * Prefix to encode node positions in an ErgoTree instance.
+    */
+  val ErgoTreePrefix = NodePosition(Seq(1))
+}

interpreter/shared/src/main/scala/sigmastate/UnprovenTree.scala

@@ -34,47 +34,6 @@ trait ProofTreeConjecture extends ProofTree {
   val children: Seq[ProofTree]
 }
 
-/**
-  * Data type which encodes position of a node in a tree.
-  *
-  * Position is encoded like following (the example provided is for CTHRESHOLD(2, Seq(pk1, pk2, pk3 && pk4)) :
-  *
-  *            0
-  *          / | \
-  *         /  |  \
-  *       0-0 0-1 0-2
-  *               /|
-  *              / |
-  *             /  |
-  *            /   |
-  *          0-2-0 0-2-1
-  *
-  * So a hint associated with pk1 has a position "0-0", pk4 - "0-2-1" .
-  *
-  * Please note that "0" prefix is for a crypto tree. There are several kinds of trees during evaluation.
-  * Initial mixed tree (ergoTree) would have another prefix.
-  *
-  * @param positions - positions from root (inclusive) in top-down order
-  */
-case class NodePosition(positions: Seq[Int]) {
-
-  def child(childIdx: Int): NodePosition = NodePosition(positions :+ childIdx)
-
-  override def toString: String = positions.mkString("-")
-}
-
-object NodePosition {
-  /**
-    * Prefix to encode node positions in a crypto tree.
-    */
-  val CryptoTreePrefix = NodePosition(Seq(0))
-
-  /**
-    * Prefix to encode node positions in an ErgoTree instance.
-    */
-  val ErgoTreePrefix = NodePosition(Seq(1))
-}
-
 /**
   * A node of a sigma-tree used by the prover. See ProverInterpreter comments and the
   * ErgoScript white-paper https://ergoplatform.org/docs/ErgoScript.pdf , Appendix A, for details

interpreter/shared/src/main/scala/sigmastate/Values.scala

@@ -739,8 +739,11 @@ object Values {
     /** Size of the proposition tree (number of nodes). */
     def size: Int
 
-    /** Recursively collect all the leaves of this sigma expression into `buf`. */
-    def collectLeaves(buf: mutable.ArrayBuffer[SigmaLeaf]): Unit
+    /** Recursively collect all the leaves of this sigma expression into `buf`.
+      * @param position - position of this node in the tree
+      * @param buf - buffer to collect leaves into
+      */
+    def collectLeaves(position: NodePosition, buf: mutable.ArrayBuffer[PositionedLeaf]): Unit
   }
 
   object SigmaBoolean {
@@ -1012,17 +1015,17 @@ object Values {
     }
 
     /** Traverses the tree and returns all leaves nodes of sigma proposition tree. */
-    def leaves(): Seq[SigmaLeaf] = {
-      val buf = mutable.ArrayBuffer.empty[SigmaLeaf]
-      sb.collectLeaves(buf)
+    def leaves(): Seq[PositionedLeaf] = {
+      val buf = mutable.ArrayBuffer.empty[PositionedLeaf]
+      sb.collectLeaves(NodePosition.CryptoTreePrefix, buf)
       buf.toSeq
     }
 
     /** Traverses the tree and returns all DISTINCT leaves of sigma proposition tree. */
     def distinctLeaves: Set[SigmaLeaf] = {
-      val buf = mutable.ArrayBuffer.empty[SigmaLeaf]
-      sb.collectLeaves(buf)
-      buf.iterator.toSet
+      val buf = mutable.ArrayBuffer.empty[PositionedLeaf]
+      sb.collectLeaves(NodePosition.CryptoTreePrefix, buf)
+      buf.iterator.map(_.leaf).toSet
     }
   }
 

interpreter/shared/src/main/scala/sigmastate/trees.scala

@@ -33,9 +33,9 @@ import scala.collection.mutable.ArrayBuffer
 trait SigmaConjecture extends SigmaBoolean {
   def children: Seq[SigmaBoolean]
 
-  override def collectLeaves(buf: mutable.ArrayBuffer[SigmaLeaf]): Unit = {
+  override def collectLeaves(position: NodePosition, buf: mutable.ArrayBuffer[PositionedLeaf]): Unit = {
     cfor(0)(_ < children.length, _ + 1) { i =>
-      children(i).collectLeaves(buf)
+      children(i).collectLeaves(position.child(i),  buf)
     }
   }
 }
@@ -44,10 +44,12 @@ trait SigmaConjecture extends SigmaBoolean {
   * Basic trait for leafs of crypto-trees, such as ProveDlog and ProveDiffieHellman instances
   */
 trait SigmaLeaf extends SigmaBoolean {
-  override def collectLeaves(buf: mutable.ArrayBuffer[SigmaLeaf]): Unit =
-    buf += this
+  override def collectLeaves(position: NodePosition, buf: mutable.ArrayBuffer[PositionedLeaf]): Unit =
+    buf += PositionedLeaf(position, this)
 }
 
+/** Represents leaf and its position in a SigmaBoolean tree. */
+case class PositionedLeaf(position: NodePosition, leaf: SigmaLeaf)
 
 /**
   * AND conjunction for sigma propositions
@@ -139,7 +141,7 @@ case class CTHRESHOLD(k: Int, children: Seq[SigmaBoolean]) extends SigmaConjectu
 abstract class TrivialProp(val condition: Boolean) extends SigmaBoolean with Product1[Boolean] {
   override def _1: Boolean = condition
   override def canEqual(that: Any): Boolean = that != null && that.isInstanceOf[TrivialProp]
-  override def collectLeaves(buf: mutable.ArrayBuffer[SigmaLeaf]): Unit = () // not a leaf
+  override def collectLeaves(position: NodePosition, buf: mutable.ArrayBuffer[PositionedLeaf]): Unit = () // not a leaf
 }
 object TrivialProp {
   // NOTE: the corresponding unapply is missing because any implementation (even using Nullable)

interpreter/shared/src/test/scala/sigmastate/SigmaProtocolSpecification.scala

@@ -45,19 +45,29 @@ class SigmaProtocolSpecification extends SigmaTestingData with ScalaCheckPropert
     val th = CTHRESHOLD(1, Seq(dlog1, dht1))
     val th2 = CTHRESHOLD(2, Seq(TrueProp, and, or, th, dlog1, dht1))
 
+    def position(path: Int*): NodePosition = NodePosition.CryptoTreePrefix ++ path
+
     val table = Table(("proposition", "leafs"),
       (TrueProp, Seq()),
       (FalseProp, Seq()),
-      (dlog1, Seq(dlog1)),
-      (dht1, Seq(dht1)),
-      (and, Seq(dlog1, dlog2)),
-      (or, Seq(dlog1, dlog2)),
-      (th, Seq(dlog1, dht1)),
-      (th2, Seq(dlog1, dlog2, dlog1, dlog2, dlog1, dht1, dlog1, dht1))
+      (dlog1, Seq(PositionedLeaf(NodePosition.CryptoTreePrefix, dlog1))),
+      (dht1, Seq(PositionedLeaf(NodePosition.CryptoTreePrefix, dht1))),
+      (and, Seq(PositionedLeaf(position(0), dlog1), PositionedLeaf(position(1), dlog2))),
+      (or, Seq(PositionedLeaf(position(0), dlog1), PositionedLeaf(position(1), dlog2))),
+      (th, Seq(PositionedLeaf(position(0), dlog1), PositionedLeaf(position(1), dht1))),
+      (th2, Seq(
+        PositionedLeaf(position(1, 0), dlog1),
+        PositionedLeaf(position(1, 1), dlog2),
+        PositionedLeaf(position(2, 0), dlog1),
+        PositionedLeaf(position(2, 1), dlog2),
+        PositionedLeaf(position(3, 0), dlog1),
+        PositionedLeaf(position(3, 1), dht1),
+        PositionedLeaf(position(4), dlog1),
+        PositionedLeaf(position(5), dht1)))
     )
-    forAll(table) { (prop: SigmaBoolean, leafs: Seq[SigmaLeaf]) =>
+    forAll(table) { (prop: SigmaBoolean, leafs: Seq[PositionedLeaf]) =>
       prop.leaves shouldBe leafs
     }
-    th2.leaves.iterator.distinct.toSet shouldBe Set(dlog1, dlog2, dht1)
+    th2.distinctLeaves shouldBe Set(dlog1, dlog2, dht1)
   }
 }