AVM: Track Scratch Slot Types

data/transactions/logic/assembler.go

@@ -252,12 +252,18 @@ type OpStream struct {
 // newOpStream constructs OpStream instances ready to invoke assemble. A new
 // OpStream must be used for each call to assemble().
 func newOpStream(version uint64) OpStream {
-	return OpStream{
+	o := OpStream{
 		labels:       make(map[string]int),
 		OffsetToLine: make(map[int]int),
 		typeTracking: true,
 		Version:      version,
 	}
+
+	for i := range o.known.scratchSpace {
+		o.known.scratchSpace[i] = StackUint64
+	}
+
+	return o
 }
 
 // ProgramKnowledge tracks statically known information as we assemble
@@ -279,6 +285,8 @@ type ProgramKnowledge struct {
 	// deadcode indicates that the program is in deadcode, so no type checking
 	// errors should be reported.
 	deadcode bool
+
+	scratchSpace [256]StackType
 }
 
 func (pgm *ProgramKnowledge) pop() StackType {
@@ -312,6 +320,9 @@ func (pgm *ProgramKnowledge) reset() {
 	pgm.stack = nil
 	pgm.bottom = StackAny
 	pgm.deadcode = false
+	for i := range pgm.scratchSpace {
+		pgm.scratchSpace[i] = StackAny
+	}
 }
 
 // createLabel inserts a label to point to the next instruction, reporting an
@@ -334,7 +345,7 @@ func (ops *OpStream) referToLabel(pc int, label string) {
 	ops.labelReferences = append(ops.labelReferences, labelReference{ops.sourceLine, pc, label})
 }
 
-type refineFunc func(pgm ProgramKnowledge, immediates []string) (StackTypes, StackTypes)
+type refineFunc func(pgm *ProgramKnowledge, immediates []string) (StackTypes, StackTypes)
 
 // returns allows opcodes like `txn` to be specific about their return value
 // types, based on the field requested, rather than use Any as specified by
@@ -969,7 +980,19 @@ func asmDefault(ops *OpStream, spec *OpSpec, args []string) error {
 	return nil
 }
 
-func typeSwap(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
+// Interprets the arg at index argIndex as byte-long immediate
+func getByteImm(args []string, argIndex int) (byte, bool) {
+	if len(args) <= argIndex {
+		return 0, false
+	}
+	n, err := strconv.ParseUint(args[argIndex], 0, 8)
+	if err != nil {
+		return 0, false
+	}
+	return byte(n), true
+}
+
+func typeSwap(pgm *ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	topTwo := StackTypes{StackAny, StackAny}
 	top := len(pgm.stack) - 1
 	if top >= 0 {
@@ -982,12 +1005,9 @@ func typeSwap(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	return nil, reversed
 }
 
-func typeDig(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
-	if len(args) == 0 {
-		return nil, nil
-	}
-	n, err := strconv.ParseUint(args[0], 0, 64)
-	if err != nil {
+func typeDig(pgm *ProgramKnowledge, args []string) (StackTypes, StackTypes) {
+	n, ok := getByteImm(args, 0)
+	if !ok {
 		return nil, nil
 	}
 	depth := int(n) + 1
@@ -1008,7 +1028,7 @@ func typeDig(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	return anys, returns
 }
 
-func typeEquals(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
+func typeEquals(pgm *ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	top := len(pgm.stack) - 1
 	if top >= 0 {
 		//Require arg0 and arg1 to have same type
@@ -1017,15 +1037,15 @@ func typeEquals(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	return nil, nil
 }
 
-func typeDup(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
+func typeDup(pgm *ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	top := len(pgm.stack) - 1
 	if top >= 0 {
 		return StackTypes{pgm.stack[top]}, StackTypes{pgm.stack[top], pgm.stack[top]}
 	}
 	return nil, nil
 }
 
-func typeDupTwo(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
+func typeDupTwo(pgm *ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	topTwo := StackTypes{StackAny, StackAny}
 	top := len(pgm.stack) - 1
 	if top >= 0 {
@@ -1037,7 +1057,7 @@ func typeDupTwo(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	return nil, append(topTwo, topTwo...)
 }
 
-func typeSelect(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
+func typeSelect(pgm *ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	top := len(pgm.stack) - 1
 	if top >= 2 {
 		if pgm.stack[top-1] == pgm.stack[top-2] {
@@ -1047,20 +1067,17 @@ func typeSelect(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	return nil, nil
 }
 
-func typeSetBit(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
+func typeSetBit(pgm *ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	top := len(pgm.stack) - 1
 	if top >= 2 {
 		return nil, StackTypes{pgm.stack[top-2]}
 	}
 	return nil, nil
 }
 
-func typeCover(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
-	if len(args) == 0 {
-		return nil, nil
-	}
-	n, err := strconv.ParseUint(args[0], 0, 64)
-	if err != nil {
+func typeCover(pgm *ProgramKnowledge, args []string) (StackTypes, StackTypes) {
+	n, ok := getByteImm(args, 0)
+	if !ok {
 		return nil, nil
 	}
 	depth := int(n) + 1
@@ -1086,12 +1103,9 @@ func typeCover(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	return anys, returns
 }
 
-func typeUncover(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
-	if len(args) == 0 {
-		return nil, nil
-	}
-	n, err := strconv.ParseUint(args[0], 0, 64)
-	if err != nil {
+func typeUncover(pgm *ProgramKnowledge, args []string) (StackTypes, StackTypes) {
+	n, ok := getByteImm(args, 0)
+	if !ok {
 		return nil, nil
 	}
 	depth := int(n) + 1
@@ -1114,7 +1128,7 @@ func typeUncover(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	return anys, returns
 }
 
-func typeTxField(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
+func typeTxField(pgm *ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	if len(args) != 1 {
 		return nil, nil
 	}
@@ -1125,6 +1139,51 @@ func typeTxField(pgm ProgramKnowledge, args []string) (StackTypes, StackTypes) {
 	return StackTypes{fs.ftype}, nil
 }
 
+func typeStore(pgm *ProgramKnowledge, args []string) (StackTypes, StackTypes) {
+	scratchIndex, ok := getByteImm(args, 0)
+	if !ok {
+		return nil, nil
+	}
+	top := len(pgm.stack) - 1
+	if top >= 0 {
+		pgm.scratchSpace[scratchIndex] = pgm.stack[top]
+	}
+	return nil, nil
+}
+
+func typeStores(pgm *ProgramKnowledge, args []string) (StackTypes, StackTypes) {
+	top := len(pgm.stack) - 1
+	if top < 0 {
+		return nil, nil
+	}
+	for i := range pgm.scratchSpace {
+		// We can't know what slot stacktop is being stored in, but we can at least keep the slots that are the same type as stacktop
+		if pgm.scratchSpace[i] != pgm.stack[top] {
+			pgm.scratchSpace[i] = StackAny
+		}
+	}
+	return nil, nil
+}
+
+func typeLoad(pgm *ProgramKnowledge, args []string) (StackTypes, StackTypes) {
+	scratchIndex, ok := getByteImm(args, 0)
+	if !ok {
+		return nil, nil
+	}
+	return nil, StackTypes{pgm.scratchSpace[scratchIndex]}
+}
+
+func typeLoads(pgm *ProgramKnowledge, args []string) (StackTypes, StackTypes) {
+	scratchType := pgm.scratchSpace[0]
+	for _, item := range pgm.scratchSpace {
+		// If all the scratch slots are one type, then we can say we are loading that type
+		if item != scratchType {
+			return nil, nil
+		}
+	}
+	return nil, StackTypes{scratchType}
+}
+
 // keywords or "pseudo-ops" handle parsing and assembling special asm language
 // constructs like 'addr' We use an OpSpec here, but it's somewhat degenerate,
 // since they don't have opcodes or eval functions. But it does need a lot of
@@ -1366,7 +1425,7 @@ func (ops *OpStream) assemble(text string) error {
 			}
 			args, returns := spec.Arg.Types, spec.Return.Types
 			if spec.OpDetails.refine != nil {
-				nargs, nreturns := spec.OpDetails.refine(ops.known, fields[1:])
+				nargs, nreturns := spec.OpDetails.refine(&ops.known, fields[1:])
 				if nargs != nil {
 					args = nargs
 				}

data/transactions/logic/assembler_test.go

@@ -2400,6 +2400,29 @@ func TestSetBitTypeCheck(t *testing.T) {
 	testProg(t, "byte 0x1234; int 2; int 3; setbit; !", AssemblerMaxVersion, Expect{5, "! arg 0..."})
 }
 
+func TestScratchTypeCheck(t *testing.T) {
+	partitiontest.PartitionTest(t)
+	t.Parallel()
+	// All scratch slots should start as uint64
+	testProg(t, "load 0; int 1; +", AssemblerMaxVersion)
+	// Check load and store accurately using the scratch space
+	testProg(t, "byte 0x01; store 0; load 0; int 1; +", AssemblerMaxVersion, Expect{5, "+ arg 0..."})
+	// Loads should know the type it's loading if all the slots are the same type
+	testProg(t, "int 0; loads; btoi", AssemblerMaxVersion, Expect{3, "btoi arg 0..."})
+	// Loads doesn't know the type when slot types vary
+	testProg(t, "byte 0x01; store 0; int 1; loads; btoi", AssemblerMaxVersion)
+	// Stores should only set slots to StackAny if they are not the same type as what is being stored
+	testProg(t, "byte 0x01; store 0; int 3; byte 0x01; stores; load 0; int 1; +", AssemblerMaxVersion, Expect{8, "+ arg 0..."})
+	// ScratchSpace should reset after hitting label in deadcode
+	testProg(t, "byte 0x01; store 0; b label1; label1:; load 0; int 1; +", AssemblerMaxVersion)
+	// But it should reset to StackAny not uint64
+	testProg(t, "int 1; store 0; b label1; label1:; load 0; btoi", AssemblerMaxVersion)
+	// Callsubs should also reset the scratch space
+	testProg(t, "callsub A; load 0; btoi; return; A: byte 0x01; store 0; retsub", AssemblerMaxVersion)
+	// But the scratchspace should still be tracked after the callsub
+	testProg(t, "callsub A; int 1; store 0; load 0; btoi; return; A: retsub", AssemblerMaxVersion, Expect{5, "btoi arg 0..."})
+}
+
 func TestCoverAsm(t *testing.T) {
 	partitiontest.PartitionTest(t)
 	t.Parallel()

data/transactions/logic/eval_test.go

@@ -4534,6 +4534,8 @@ func TestLog(t *testing.T) {
 		source      string
 		runMode     runMode
 		errContains string
+		// For cases where assembly errors, we manually put in the bytes
+		assembledBytes []byte
 	}{
 		{
 			source:      fmt.Sprintf(`byte  "%s"; log; int 1`, strings.Repeat("a", maxLogSize+1)),
@@ -4561,9 +4563,10 @@ func TestLog(t *testing.T) {
 			runMode:     modeApp,
 		},
 		{
-			source:      `load 0; log`,
-			errContains: "log arg 0 wanted []byte but got uint64",
-			runMode:     modeApp,
+			source:         `load 0; log`,
+			errContains:    "log arg 0 wanted []byte but got uint64",
+			runMode:        modeApp,
+			assembledBytes: []byte{byte(ep.Proto.LogicSigVersion), 0x34, 0x00, 0xb0},
 		},
 		{
 			source:      `byte  "a logging message"; log; int 1`,
@@ -4575,7 +4578,11 @@ func TestLog(t *testing.T) {
 	for _, c := range failCases {
 		switch c.runMode {
 		case modeApp:
-			testApp(t, c.source, ep, c.errContains)
+			if c.assembledBytes == nil {
+				testApp(t, c.source, ep, c.errContains)
+			} else {
+				testAppBytes(t, c.assembledBytes, ep, c.errContains)
+			}
 		default:
 			testLogic(t, c.source, AssemblerMaxVersion, ep, c.errContains, c.errContains)
 		}

data/transactions/logic/opcodes.go

@@ -443,8 +443,8 @@ var OpSpecs = []OpSpec{
 	{0x32, "global", opGlobal, proto(":a"), 1, field("f", &GlobalFields)},
 	{0x33, "gtxn", opGtxn, proto(":a"), 1, immediates("t", "f").field("f", &TxnScalarFields)},
 	{0x33, "gtxn", opGtxn, proto(":a"), 2, immediates("t", "f").field("f", &TxnFields).assembler(asmGtxn2)},
-	{0x34, "load", opLoad, proto(":a"), 1, immediates("i")},
-	{0x35, "store", opStore, proto("a:"), 1, immediates("i")},
+	{0x34, "load", opLoad, proto(":a"), 1, stacky(typeLoad, "i")},
+	{0x35, "store", opStore, proto("a:"), 1, stacky(typeStore, "i")},
 	{0x36, "txna", opTxna, proto(":a"), 2, immediates("f", "i").field("f", &TxnArrayFields)},
 	{0x37, "gtxna", opGtxna, proto(":a"), 2, immediates("t", "f", "i").field("f", &TxnArrayFields)},
 	// Like gtxn, but gets txn index from stack, rather than immediate arg
@@ -458,8 +458,8 @@ var OpSpecs = []OpSpec{
 	{0x3d, "gaids", opGaids, proto("i:i"), 4, only(modeApp)},
 
 	// Like load/store, but scratch slot taken from TOS instead of immediate
-	{0x3e, "loads", opLoads, proto("i:a"), 5, opDefault()},
-	{0x3f, "stores", opStores, proto("ia:"), 5, opDefault()},
+	{0x3e, "loads", opLoads, proto("i:a"), 5, stacky(typeLoads)},
+	{0x3f, "stores", opStores, proto("ia:"), 5, stacky(typeStores)},
 
 	{0x40, "bnz", opBnz, proto("i:"), 1, opBranch()},
 	{0x41, "bz", opBz, proto("i:"), 2, opBranch()},