1 files changed, 31 insertions, 27 deletions

diff --git a/prog/size.go b/prog/size.go
index 60cc7bd57..f00fd56b1 100644
--- a/prog/size.go
+++ b/prog/size.go
@@ -9,20 +9,22 @@ import (
 	"github.com/google/syzkaller/sys"
 )
 
-func generateSize(arg *Arg, lenType *sys.LenType) *Arg {
+func generateSize(arg Arg, lenType *sys.LenType) Arg {
 	if arg == nil {
 		// Arg is an optional pointer, set size to 0.
 		return constArg(lenType, 0)
 	}
 
-	switch arg.Type.(type) {
+	switch arg.Type().(type) {
 	case *sys.VmaType:
-		return pageSizeArg(lenType, arg.AddrPagesNum, 0)
+		a := arg.(*PointerArg)
+		return constArg(lenType, a.PagesNum*pageSize)
 	case *sys.ArrayType:
+		a := arg.(*GroupArg)
 		if lenType.ByteSize != 0 {
-			return constArg(lenType, arg.Size()/lenType.ByteSize)
+			return constArg(lenType, a.Size()/lenType.ByteSize)
 		} else {
-			return constArg(lenType, uintptr(len(arg.Inner)))
+			return constArg(lenType, uintptr(len(a.Inner)))
 		}
 	default:
 		if lenType.ByteSize != 0 {
@@ -33,42 +35,44 @@ func generateSize(arg *Arg, lenType *sys.LenType) *Arg {
 	}
 }
 
-func assignSizes(args []*Arg, parentsMap map[*Arg]*Arg) {
-	// Create a map of args and calculate size of the whole struct.
-	argsMap := make(map[string]*Arg)
+func assignSizes(args []Arg, parentsMap map[Arg]Arg) {
+	// Create a map from field names to args.
+	argsMap := make(map[string]Arg)
 	for _, arg := range args {
-		if sys.IsPad(arg.Type) {
+		if sys.IsPad(arg.Type()) {
 			continue
 		}
-		argsMap[arg.Type.FieldName()] = arg
+		argsMap[arg.Type().FieldName()] = arg
 	}
 
 	// Fill in size arguments.
 	for _, arg := range args {
-		if arg = arg.InnerArg(); arg == nil {
+		if arg = InnerArg(arg); arg == nil {
 			continue // Pointer to optional len field, no need to fill in value.
 		}
-		if typ, ok := arg.Type.(*sys.LenType); ok {
+		if typ, ok := arg.Type().(*sys.LenType); ok {
+			a := arg.(*ConstArg)
+
 			buf, ok := argsMap[typ.Buf]
 			if ok {
-				*arg = *generateSize(buf.InnerArg(), typ)
+				*a = *generateSize(InnerArg(buf), typ).(*ConstArg)
 				continue
 			}
 
 			if typ.Buf == "parent" {
-				arg.Val = parentsMap[arg].Size()
+				a.Val = parentsMap[arg].Size()
 				if typ.ByteSize != 0 {
-					arg.Val /= typ.ByteSize
+					a.Val /= typ.ByteSize
 				}
 				continue
 			}
 
 			sizeAssigned := false
 			for parent := parentsMap[arg]; parent != nil; parent = parentsMap[parent] {
-				if typ.Buf == parent.Type.Name() {
-					arg.Val = parent.Size()
+				if typ.Buf == parent.Type().Name() {
+					a.Val = parent.Size()
 					if typ.ByteSize != 0 {
-						arg.Val /= typ.ByteSize
+						a.Val /= typ.ByteSize
 					}
 					sizeAssigned = true
 					break
@@ -84,19 +88,19 @@ func assignSizes(args []*Arg, parentsMap map[*Arg]*Arg) {
 	}
 }
 
-func assignSizesArray(args []*Arg) {
-	parentsMap := make(map[*Arg]*Arg)
-	foreachArgArray(&args, nil, func(arg, base *Arg, _ *[]*Arg) {
-		if _, ok := arg.Type.(*sys.StructType); ok {
-			for _, field := range arg.Inner {
-				parentsMap[field.InnerArg()] = arg
+func assignSizesArray(args []Arg) {
+	parentsMap := make(map[Arg]Arg)
+	foreachArgArray(&args, nil, func(arg, base Arg, _ *[]Arg) {
+		if _, ok := arg.Type().(*sys.StructType); ok {
+			for _, field := range arg.(*GroupArg).Inner {
+				parentsMap[InnerArg(field)] = arg
 			}
 		}
 	})
 	assignSizes(args, parentsMap)
-	foreachArgArray(&args, nil, func(arg, base *Arg, _ *[]*Arg) {
-		if _, ok := arg.Type.(*sys.StructType); ok {
-			assignSizes(arg.Inner, parentsMap)
+	foreachArgArray(&args, nil, func(arg, base Arg, _ *[]Arg) {
+		if _, ok := arg.Type().(*sys.StructType); ok {
+			assignSizes(arg.(*GroupArg).Inner, parentsMap)
 		}
 	})
 }