prog: include number of calls into exec encoding

Prepend total number of calls to the exec encoding.
This will allow pkg/ipc to better parse executor response
without full parsing of the encoded program.

3 files changed, 23 insertions, 2 deletions

diff --git a/prog/decodeexec.go b/prog/decodeexec.go
index 2356e051b..81c70cb2b 100644
--- a/prog/decodeexec.go
+++ b/prog/decodeexec.go
@@ -72,6 +72,17 @@ type ExecCsumChunk struct {
 	Size  uint64
 }
 
+func ExecCallCount(exec []byte) (int, error) {
+	v, n := binary.Varint(exec)
+	if n <= 0 {
+		return 0, fmt.Errorf("not enough data in the buffer")
+	}
+	if v > MaxCalls {
+		return 0, fmt.Errorf("too many calls (%v)", v)
+	}
+	return int(v), nil
+}
+
 func (target *Target) DeserializeExec(exec []byte, stats map[string]int) (ExecProg, error) {
 	dec := &execDecoder{target: target, data: exec, stats: stats}
 	dec.parse()
@@ -101,6 +112,7 @@ type execDecoder struct {
 }
 
 func (dec *execDecoder) parse() {
+	ncalls := dec.read("header")
 	for dec.err == nil {
 		switch instr := dec.read("instr/opcode"); instr {
 		case execInstrCopyin:
@@ -117,6 +129,9 @@ func (dec *execDecoder) parse() {
 			})
 		case execInstrEOF:
 			dec.commitCall()
+			if ncalls != uint64(len(dec.calls)) {
+				dec.err = fmt.Errorf("bad number of calls: %v/%v", ncalls, len(dec.calls))
+			}
 			return
 		case execInstrSetProps:
 			dec.commitCall()
diff --git a/prog/encodingexec.go b/prog/encodingexec.go
index a51aa54d7..40cfc9592 100644
--- a/prog/encodingexec.go
+++ b/prog/encodingexec.go
@@ -75,6 +75,7 @@ func (p *Prog) SerializeForExec() ([]byte, error) {
 		buf:    make([]byte, 0, 4<<10),
 		args:   make(map[Arg]argInfo),
 	}
+	w.write(uint64(len(p.Calls)))
 	for _, c := range p.Calls {
 		w.csumMap, w.csumUses = calcChecksumsCall(c)
 		w.serializeCall(c)
diff --git a/prog/encodingexec_test.go b/prog/encodingexec_test.go
index 726cbe49b..357cca18b 100644
--- a/prog/encodingexec_test.go
+++ b/prog/encodingexec_test.go
@@ -26,10 +26,15 @@ func TestSerializeForExecRandom(t *testing.T) {
 		if err != nil {
 			t.Fatalf("failed to serialize: %v", err)
 		}
-		_, err = target.DeserializeExec(buf, sizes)
+		got, err := target.DeserializeExec(buf, sizes)
 		if err != nil {
 			t.Fatal(err)
 		}
+		if n, err := ExecCallCount(buf); err != nil {
+			t.Fatal(err)
+		} else if n != len(got.Calls) {
+			t.Fatalf("mismatching number of calls: %v/%v", n, len(got.Calls))
+		}
 		totalSize += len(buf)
 		execSizes.Add(float64(len(buf)))
 		textSizes.Add(float64(len(p.Serialize())))
@@ -660,7 +665,7 @@ test$res1(r0)
 			if err != nil {
 				t.Fatalf("failed to serialize: %v", err)
 			}
-			var want []byte
+			want := binary.AppendVarint(nil, int64(len(p.Calls)))
 			for _, e := range test.serialized {
 				switch elem := e.(type) {
 				case uint64: