Fixes and tests for missing data in messages

convex-core/src/main/java/convex/core/data/Format.java

@@ -1,11 +1,12 @@
 package convex.core.data;
-
+ 
 import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.function.Consumer;
+import java.util.function.Predicate;
 
 import convex.core.Belief;
 import convex.core.Block;
@@ -725,7 +726,7 @@ public static ACell[] decodeCells(Blob data) throws BadFormatException {
 
 		return cells.toArray(ACell[]::new);
 	}
-
+	
 	/**
 	 * Reads a cell from a Blob of data, allowing for non-embedded children following the first cell
 	 * @param data Data to decode
@@ -745,28 +746,11 @@ public static <T extends ACell> T decodeMultiCell(Blob data) throws BadFormatExc
 		if (rl==ml) return result; // Already complete
 
 		// read remaining cells
-		HashMap<Hash,Ref<?>> hm=new HashMap<>();
-		for (int ix=rl; ix<ml;) {
-			ACell c=Format.read(data,ix);
-			if (c==null) {
-				c=Format.read(data,ix);
-				throw new BadFormatException("Null child encoding in Message");
-			}
-			if (c.isEmbedded()) throw new BadFormatException("Embedded Cell provided in Message");
-			Hash h=c.getHash();
-
-			// Check store for Ref - avoids duplicate objects in many cases
-			//Ref<?> storeRef=store.checkCache(h);
-			//Ref<?> cr=(storeRef!=null)?storeRef:Ref.get(c);
-
-			Ref<?> cr=Ref.get(c);
-			// System.out.println("Decoding novelty: "+cr.getHash()+ " "+c.getClass().getSimpleName());
-			hm.put(h, cr);
-			ix+=c.getEncodingLength();
-		}
+		HashMap<Hash,ACell> hm=new HashMap<>();
+		decodeCells(hm,data.slice(rl,ml));
 
 		HashMap<Hash,ACell> done=new HashMap<Hash,ACell>();		
-		ArrayList<ACell> al=new ArrayList<>();
+		ArrayList<ACell> stack=new ArrayList<>();
 
 		IRefFunction func=new IRefFunction() {
 			@Override
@@ -779,15 +763,16 @@ public Ref<?> apply(Ref<?> r) {
 					return nc.getRef();
 				} else {
 					Hash h=r.getHash();
+
 					// if done, just replace with done version
 					ACell doneVal=done.get(h);
 					if (doneVal!=null) return doneVal.getRef();
 
 					// if in map, push cell to stack
-					Ref<?> partRef=hm.get(h);
-					if (partRef!=null) {
-						al.add(partRef.getValue());
-						return partRef;
+					ACell part=hm.get(h);
+					if (part!=null) {
+						stack.add(part);
+						return part.getRef();
 					}
 
 					// not in message, must be partial
@@ -796,27 +781,62 @@ public Ref<?> apply(Ref<?> r) {
 			} 
 		};
 
-		al.add(result);
-		Trees.visitStack(al, c->{
-			Hash h=c.getHash();
-			if (done.containsKey(h)) return;
-
-			al.add(c);
-			int pos=al.size();
-			ACell nc=c.updateRefs(func);
-			if (pos==al.size()) {
-				// we must be done
-				done.put(h,nc);
-			} else {
-				// something extra on the stack to handle first
+		stack.add(result); 
+		Trees.visitStackMaybePopping(stack, new Predicate<ACell>() {
+			@Override
+			public boolean test(ACell c) {
+				Hash h=c.getHash();
+				if (done.containsKey(h)) return true;
+
+				int pos=stack.size();
+				// Update Refs, adding new non-embedded cells to stack
+				ACell nc=c.updateRefs(func);
+
+				if (stack.size()==pos) {
+					// we must be done since nothing new added to stack
+					done.put(h,nc);
+					return true;
+				} else {
+					// something extra on the stack to handle first
+					stack.set(pos-1,nc);
+					return false;
+				}
 			}
 		});
-
+
+		// ACell cc=done.get(check);
 
 		result=(T) done.get(result.getHash());
 
 		return result;
 	}
+
+	/**
+	 * Decode encoded non-embedded Cells into an accumulator HashMap
+	 * @param acc Accumulator for Cells, keyed by Hash
+	 * @param data Encoding to read
+	 * @throws BadFormatException In case of bad format, including any embedded values
+	 */
+	public static void decodeCells(HashMap<Hash,ACell> acc, Blob data) throws BadFormatException {
+		long ml=data.count();
+		int ix=0;
+		while( ix<ml) {
+			ACell c=Format.read(data,ix);
+			if (c==null) {
+				throw new BadFormatException("Null child encoding in Message");
+			}
+			if (c.isEmbedded()) throw new BadFormatException("Embedded Cell provided in Message");
+			Hash h=c.getHash();
+
+			// Check store for Ref - avoids duplicate objects in many cases
+			//Ref<?> storeRef=store.checkCache(h);
+			//Ref<?> cr=(storeRef!=null)?storeRef:Ref.get(c);
+
+			acc.put(h, c);
+			ix+=c.getEncodingLength();
+		}
+		if (ix!=ml) throw new BadFormatException("Bad message length when decoding");
+	}
 
 	/**
 	 * Encode a Cell completely in multi-cell message format. Format places top level
@@ -854,7 +874,7 @@ public static Blob encodeMultiCell(ACell a) {
 		int messageLength=ml[0];
 		byte[] msg=new byte[messageLength];
 
-		// Ensure we add each unique child
+		// Write top encoding, ensure we add each unique child
 		topCellEncoding.getBytes(msg, 0);
 		int ix=Utils.checkedInt(topCellEncoding.count());
 		for (Ref<?> r: refs) {

convex-core/src/main/java/convex/core/lang/impl/Fn.java

@@ -210,7 +210,7 @@ public Fn<T> updateRefs(IRefFunction func) {
 		AOp<T> newBody = body.updateRefs(func);
 		AVector<ACell> newLexicalEnv = lexicalEnv.updateRefs(func);
 		if ((params == newParams) && (body == newBody) && (lexicalEnv == newLexicalEnv)) return this;
-		return new Fn<>(newParams, newBody, lexicalEnv);
+		return new Fn<>(newParams, newBody, newLexicalEnv);
 	}
 
 	@Override

convex-core/src/main/java/convex/core/util/Trees.java

@@ -2,6 +2,7 @@
 
 import java.util.List;
 import java.util.function.Consumer;
+import java.util.function.Predicate;
 
 /**
  * Utility class for tree handling functions
@@ -26,4 +27,24 @@ public static <T> void visitStack(List<T> stack, Consumer<T> visitor) {
 			visitor.accept(r);
 		}
 	}
+
+	/**
+	 * Visits elements on a stack, the element if predicate returns true.
+	 * Predicate function MAY add to the stack. Will terminate when stack is empty.
+	 * 
+	 * IMPORTANT: O(1) usage of JVM stack, may be necessary to use a function like this when 
+	 * visiting deeply nested trees in CVM code.
+	 * 
+	 * @param <T> Type of element to visit
+	 * @param stack Stack of values to visit, must be a mutable List
+	 * @param visitor Visitor function to call for each stack element.
+	 */
+	public static <T> void visitStackMaybePopping(List<T> stack, Predicate<T> visitor) {
+		while(!stack.isEmpty()) {
+			int pos=stack.size()-1;
+			T r=stack.get(pos);
+			boolean pop= visitor.test(r);
+			if (pop) stack.remove(pos);
+		}
+	}
 }

convex-core/src/test/java/convex/core/StateTest.java

@@ -5,6 +5,8 @@
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import java.util.HashMap;
+
 import org.junit.jupiter.api.Test;
 
 import convex.core.crypto.AKeyPair;
@@ -13,6 +15,7 @@
 import convex.core.data.AccountStatus;
 import convex.core.data.Blob;
 import convex.core.data.Format;
+import convex.core.data.Hash;
 import convex.core.data.Lists;
 import convex.core.data.RecordTest;
 import convex.core.data.Ref;
@@ -75,10 +78,30 @@ public void testRoundTrip() throws BadFormatException {
 	@Test
 	public void testMultiCellTrip() throws BadFormatException {
 		State s = INIT_STATE;
+		RefTreeStats rstats  = Refs.getRefTreeStats(s.getRef());
+
+		// Hash of a known value in the tree that should be encoded
+		Hash check=Hash.fromHex("1fe0a93790d5e2a6d6d31db57edc611b128afe97941af611f65b703006ba5387");
 
+		Refs.visitAllRefs(s.getRef(), r->{
+			if (r.getHash().equals(check)) {
+				System.out.println(r.getValue());
+			}
+		});
+
 		Blob b=Format.encodeMultiCell(s);
+
+		HashMap<Hash,ACell> acc=new HashMap<>();
+		Format.decodeCells(acc, b);
+		assertTrue(acc.containsKey(check));
+
 		State s2=Format.decodeMultiCell(b);
+		// System.err.println(Refs.printMissingTree(s2));
 		assertEquals(s,s2);
+
+		RefTreeStats rstats2  = Refs.getRefTreeStats(s2.getRef());
+		assertEquals(rstats2.total,rstats2.direct);
+		assertEquals(rstats.total,rstats2.total);
 	}
 
 	@SuppressWarnings("unused")

convex-core/src/test/java/convex/core/data/EncodingTest.java

@@ -21,6 +21,7 @@
 import convex.core.Block;
 import convex.core.Order;
 import convex.core.crypto.AKeyPair;
+import convex.core.data.Refs.RefTreeStats;
 import convex.core.data.prim.CVMBigInteger;
 import convex.core.data.prim.CVMLong;
 import convex.core.exceptions.BadFormatException;
@@ -430,4 +431,23 @@ private <T extends ACell> T doMultiEncodingTest(ACell a) throws BadFormatExcepti
 		// illegal child tag
 		assertThrows(BadFormatException.class,()->Format.decodeMultiCell(first.append(Blob.fromHex("00FF")).toFlatBlob()));
 	}
+
+	@Test public void testFullMessageEncoding() throws BadFormatException {
+		testFullencoding(Samples.BIG_BLOB_TREE);
+		testFullencoding(Samples.INT_SET_300);
+		testFullencoding(Samples.INT_LIST_300);
+	}
+
+	private void testFullencoding(ACell s) throws BadFormatException {
+		RefTreeStats rstats  = Refs.getRefTreeStats(s.getRef());
+		Blob b=Format.encodeMultiCell(s);
+
+		ACell s2=Format.decodeMultiCell(b);
+		// System.err.println(Refs.printMissingTree(s2));
+		assertEquals(s,s2);
+
+		RefTreeStats rstats2  = Refs.getRefTreeStats(s2.getRef());
+		assertEquals(rstats2.total,rstats2.direct);
+		assertEquals(rstats.total,rstats2.total);
+	}
 }