prog: merge foreachSubargOffset into foreachArgImpl

2 files changed, 31 insertions, 55 deletions

diff --git a/prog/analysis.go b/prog/analysis.go
index 0c4102a47..2d691934a 100644
--- a/prog/analysis.go
+++ b/prog/analysis.go
@@ -112,12 +112,23 @@ func foreachArgImpl(arg Arg, ctx ArgCtx, f func(Arg, *ArgCtx)) {
 		if _, ok := a.Type().(*StructType); ok {
 			ctx.Parent = &a.Inner
 		}
+		var totalSize uint64
 		for _, arg1 := range a.Inner {
 			foreachArgImpl(arg1, ctx, f)
+			if !arg1.Type().BitfieldMiddle() {
+				size := arg1.Size()
+				ctx.Offset += size
+				totalSize += size
+			}
+		}
+		if totalSize > a.Size() {
+			panic(fmt.Sprintf("bad group arg size %v, should be <= %v for %+v",
+				totalSize, a.Size(), a))
 		}
 	case *PointerArg:
 		if a.Res != nil {
 			ctx.Base = a
+			ctx.Offset = 0
 			foreachArgImpl(a.Res, ctx, f)
 		}
 	case *UnionArg:
@@ -125,38 +136,6 @@ func foreachArgImpl(arg Arg, ctx ArgCtx, f func(Arg, *ArgCtx)) {
 	}
 }
 
-func foreachSubargOffset(arg Arg, f func(arg Arg, offset uint64)) {
-	var rec func(Arg, uint64) uint64
-	rec = func(arg1 Arg, offset uint64) uint64 {
-		switch a := arg1.(type) {
-		case *GroupArg:
-			f(arg1, offset)
-			var totalSize uint64
-			for _, arg2 := range a.Inner {
-				size := rec(arg2, offset)
-				if !arg2.Type().BitfieldMiddle() {
-					offset += size
-					totalSize += size
-				}
-			}
-			if totalSize > arg1.Size() {
-				panic(fmt.Sprintf("bad group arg size %v, should be <= %v for %+v", totalSize, arg1.Size(), arg1))
-			}
-		case *UnionArg:
-			f(arg1, offset)
-			size := rec(a.Option, offset)
-			offset += size
-			if size > arg1.Size() {
-				panic(fmt.Sprintf("bad union arg size %v, should be <= %v for arg %+v with type %+v", size, arg1.Size(), arg1, arg1.Type()))
-			}
-		default:
-			f(arg1, offset)
-		}
-		return arg1.Size()
-	}
-	rec(arg, 0)
-}
-
 // TODO(dvyukov): combine RequiresBitmasks and RequiresChecksums into a single function
 // to not walk the tree twice. They are always used together anyway.
 func RequiresBitmasks(p *Prog) bool {
diff --git a/prog/encodingexec.go b/prog/encodingexec.go
index d0c8f6c80..b5573f60f 100644
--- a/prog/encodingexec.go
+++ b/prog/encodingexec.go
@@ -86,30 +86,27 @@ func (p *Prog) SerializeForExec(buffer []byte) (int, error) {
 		}
 		// Calculate arg offsets within structs.
 		// Generate copyin instructions that fill in data into pointer arguments.
-		ForeachArg(c, func(arg Arg, _ *ArgCtx) {
-			if a, ok := arg.(*PointerArg); ok && a.Res != nil {
-				foreachSubargOffset(a.Res, func(arg1 Arg, offset uint64) {
-					addr := p.Target.PhysicalAddr(arg) + offset
-					if isUsed(arg1) || csumUses[arg1] {
-						w.args[arg1] = argInfo{Addr: addr}
-					}
-					if _, ok := arg1.(*GroupArg); ok {
-						return
-					}
-					if _, ok := arg1.(*UnionArg); ok {
-						return
-					}
-					if a1, ok := arg1.(*DataArg); ok &&
-						(a1.Type().Dir() == DirOut || len(a1.Data()) == 0) {
-						return
-					}
-					if !IsPad(arg1.Type()) && arg1.Type().Dir() != DirOut {
-						w.write(execInstrCopyin)
-						w.write(addr)
-						w.writeArg(arg1)
-					}
-				})
+		ForeachArg(c, func(arg Arg, ctx *ArgCtx) {
+			if ctx.Base == nil {
+				return
+			}
+			addr := p.Target.PhysicalAddr(ctx.Base) + ctx.Offset
+			if isUsed(arg) || csumUses[arg] {
+				w.args[arg] = argInfo{Addr: addr}
 			}
+			if _, ok := arg.(*GroupArg); ok {
+				return
+			}
+			if _, ok := arg.(*UnionArg); ok {
+				return
+			}
+			typ := arg.Type()
+			if typ.Dir() == DirOut || IsPad(typ) || arg.Size() == 0 {
+				return
+			}
+			w.write(execInstrCopyin)
+			w.write(addr)
+			w.writeArg(arg)
 		})
 		// Generate checksum calculation instructions starting from the last one,
 		// since checksum values can depend on values of the latter ones