Add test for collection lambda snapshot
package io.improbable.keanu.network;
+import io.improbable.keanu.distributions.gradient.Gaussian;
import io.improbable.keanu.vertices.Vertex;
+import io.improbable.keanu.vertices.dbl.DoubleVertex;
import io.improbable.keanu.vertices.dbl.nonprobabilistic.operators.binary.AdditionVertex;
import io.improbable.keanu.vertices.dbl.nonprobabilistic.operators.unary.SinVertex;
import io.improbable.keanu.vertices.dbl.probabilistic.GaussianVertex;
import org.apache.commons.lang3.mutable.MutableInt;
import org.junit.Test;
+import java.util.Arrays;
import java.util.Set;
import static junit.framework.TestCase.assertEquals;
@@ -39,4 +42,31 @@ public class LambdaSectionEfficiencyTest {
assertEquals(3, callsToNext.intValue());
assertEquals(3, callsToPredicate.intValue());
}
+
+ @Test
+ public void doesNotVisitVerticesMoreThanOnceForCollection() {
+ GaussianVertex A = new GaussianVertex(0, 1);
+ GaussianVertex B = new GaussianVertex(A, 1);
+ DoubleVertex BSin = B.sin();
+ DoubleVertex BCos = B.cos();
+ GaussianVertex C = new GaussianVertex(BCos, 1);
+ GaussianVertex D = new GaussianVertex(C, BSin);
+
+ MutableInt callsToNext = new MutableInt(0);
+ MutableInt callsToPredicate = new MutableInt(0);
+
+ Set<Vertex> verticesDepthFirst = LambdaSection.getVertices(Arrays.asList(A, B, C), (v) -> {
+ callsToNext.increment();
+ return v.getChildren();
+ }, (v) -> {
+ callsToPredicate.increment();
+ return true;
+ });
+
+ assertEquals(6, verticesDepthFirst.size());
+ assertThat(verticesDepthFirst, containsInAnyOrder(A, B, BSin, BCos, C, D));
+
+ assertEquals(5, callsToNext.intValue());
+ assertEquals(4, callsToPredicate.intValue());
+ }
}
