prog: collect parents during arg traversal

This spares the need to construct a parents map for len[A, T] and
conditional fields calculations.

3 files changed, 101 insertions, 102 deletions

diff --git a/prog/analysis.go b/prog/analysis.go
index 600f9bee4..7202a7d92 100644
--- a/prog/analysis.go
+++ b/prog/analysis.go
@@ -104,18 +104,41 @@ func (s *state) analyzeImpl(c *Call, resources bool) {
 	})
 }
 
+type parentStack []Arg
+
+func allocStack() parentStack {
+	// Let's save some allocations during stack traversal.
+	return make([]Arg, 0, 4)
+}
+
+func pushStack(ps parentStack, a Arg) parentStack {
+	return append(ps, a)
+}
+
+func popStack(ps parentStack) (parentStack, Arg) {
+	if len(ps) > 0 {
+		return ps[:len(ps)-1], ps[len(ps)-1]
+	}
+	return ps, nil
+}
+
 type ArgCtx struct {
-	Parent *[]Arg      // GroupArg.Inner (for structs) or Call.Args containing this arg.
-	Fields []Field     // Fields of the parent struct/syscall.
-	Base   *PointerArg // Pointer to the base of the heap object containing this arg.
-	Offset uint64      // Offset of this arg from the base.
-	Stop   bool        // If set by the callback, subargs of this arg are not visited.
+	Parent      *[]Arg      // GroupArg.Inner (for structs) or Call.Args containing this arg.
+	Fields      []Field     // Fields of the parent struct/syscall.
+	Base        *PointerArg // Pointer to the base of the heap object containing this arg.
+	Offset      uint64      // Offset of this arg from the base.
+	Stop        bool        // If set by the callback, subargs of this arg are not visited.
+	parentStack parentStack // Struct and union arguments by which the argument can be reached.
 }
 
 func ForeachSubArg(arg Arg, f func(Arg, *ArgCtx)) {
 	foreachArgImpl(arg, &ArgCtx{}, f)
 }
 
+func foreachSubArgWithStack(arg Arg, f func(Arg, *ArgCtx)) {
+	foreachArgImpl(arg, &ArgCtx{parentStack: allocStack()}, f)
+}
+
 func ForeachArg(c *Call, f func(Arg, *ArgCtx)) {
 	ctx := &ArgCtx{}
 	if c.Ret != nil {
@@ -131,6 +154,13 @@ func ForeachArg(c *Call, f func(Arg, *ArgCtx)) {
 func foreachArgImpl(arg Arg, ctx *ArgCtx, f func(Arg, *ArgCtx)) {
 	ctx0 := *ctx
 	defer func() { *ctx = ctx0 }()
+
+	if ctx.parentStack != nil {
+		switch arg.Type().(type) {
+		case *StructType, *UnionType:
+			ctx.parentStack = pushStack(ctx.parentStack, arg)
+		}
+	}
 	f(arg, ctx)
 	if ctx.Stop {
 		return
diff --git a/prog/expr.go b/prog/expr.go
index 38358b214..cfe6eb70b 100644
--- a/prog/expr.go
+++ b/prog/expr.go
@@ -53,30 +53,19 @@ func (v *Value) Evaluate(finder ArgFinder) (uint64, bool) {
 	return constArg.Val, true
 }
 
-func argFinderConstructor(t *Target, c *Call) func(*UnionArg) ArgFinder {
-	parentsMap := callParentsMap(c)
-	return func(unionArg *UnionArg) ArgFinder {
-		return func(path []string) Arg {
-			f := t.findArg(unionArg.Option, path, nil, nil, parentsMap, 0)
-			if f == nil {
-				return nil
-			}
-			if f.isAnyPtr {
-				return SquashedArgFound
-			}
-			return f.arg
+func makeArgFinder(t *Target, c *Call, unionArg *UnionArg, parents parentStack) ArgFinder {
+	return func(path []string) Arg {
+		f := t.findArg(unionArg.Option, path, nil, nil, parents, 0)
+		if f == nil {
+			return nil
 		}
+		if f.isAnyPtr {
+			return SquashedArgFound
+		}
+		return f.arg
 	}
 }
 
-func callParentsMap(c *Call) map[Arg]Arg {
-	parentsMap := map[Arg]Arg{}
-	ForeachArg(c, func(arg Arg, _ *ArgCtx) {
-		saveToParentsMap(arg, parentsMap)
-	})
-	return parentsMap
-}
-
 func (r *randGen) patchConditionalFields(c *Call, s *state) (extra []*Call, changed bool) {
 	if r.inPatchConditional {
 		return nil, false
@@ -88,8 +77,7 @@ func (r *randGen) patchConditionalFields(c *Call, s *state) (extra []*Call, chan
 	var anyPatched bool
 	for {
 		replace := map[Arg]Arg{}
-		makeArgFinder := argFinderConstructor(r.target, c)
-		forEachStaleUnion(r.target, c, makeArgFinder,
+		forEachStaleUnion(r.target, c,
 			func(unionArg *UnionArg, unionType *UnionType, okIndices []int) {
 				idx := okIndices[r.Intn(len(okIndices))]
 				newType, newDir := unionType.Fields[idx].Type,
@@ -112,39 +100,40 @@ func (r *randGen) patchConditionalFields(c *Call, s *state) (extra []*Call, chan
 	return extraCalls, anyPatched
 }
 
-func forEachStaleUnion(target *Target, c *Call, makeArgFinder func(*UnionArg) ArgFinder,
-	cb func(*UnionArg, *UnionType, []int)) {
-	ForeachArg(c, func(arg Arg, argCtx *ArgCtx) {
-		if target.isAnyPtr(arg.Type()) {
-			argCtx.Stop = true
-			return
-		}
-		unionArg, ok := arg.(*UnionArg)
-		if !ok {
-			return
-		}
-		unionType, ok := arg.Type().(*UnionType)
-		if !ok || !unionType.isConditional() {
-			return
-		}
-		argFinder := makeArgFinder(unionArg)
-		ok, calculated := checkUnionArg(unionArg.Index, unionType, argFinder)
-		if !calculated {
-			// Let it stay as is.
-			return
-		}
-		if !unionArg.transient && ok {
-			return
-		}
-		matchingIndices := matchingUnionArgs(unionType, argFinder)
-		if len(matchingIndices) == 0 {
-			// Conditional fields are transformed in such a way
-			// that one field always matches.
-			// For unions we demand that there's a field w/o conditions.
-			panic(fmt.Sprintf("no matching union fields: %#v", unionType))
-		}
-		cb(unionArg, unionType, matchingIndices)
-	})
+func forEachStaleUnion(target *Target, c *Call, cb func(*UnionArg, *UnionType, []int)) {
+	for _, callArg := range c.Args {
+		foreachSubArgWithStack(callArg, func(arg Arg, argCtx *ArgCtx) {
+			if target.isAnyPtr(arg.Type()) {
+				argCtx.Stop = true
+				return
+			}
+			unionArg, ok := arg.(*UnionArg)
+			if !ok {
+				return
+			}
+			unionType, ok := arg.Type().(*UnionType)
+			if !ok || !unionType.isConditional() {
+				return
+			}
+			argFinder := makeArgFinder(target, c, unionArg, argCtx.parentStack)
+			ok, calculated := checkUnionArg(unionArg.Index, unionType, argFinder)
+			if !calculated {
+				// Let it stay as is.
+				return
+			}
+			if !unionArg.transient && ok {
+				return
+			}
+			matchingIndices := matchingUnionArgs(unionType, argFinder)
+			if len(matchingIndices) == 0 {
+				// Conditional fields are transformed in such a way
+				// that one field always matches.
+				// For unions we demand that there's a field w/o conditions.
+				panic(fmt.Sprintf("no matching union fields: %#v", unionType))
+			}
+			cb(unionArg, unionType, matchingIndices)
+		})
+	}
 }
 
 func checkUnionArg(idx int, typ *UnionType, finder ArgFinder) (ok, calculated bool) {
@@ -186,9 +175,7 @@ var ErrViolatedConditions = errors.New("conditional fields rules violation")
 
 func (c *Call) checkConditions(target *Target, ignoreTransient bool) error {
 	var ret error
-
-	makeArgFinder := argFinderConstructor(target, c)
-	forEachStaleUnion(target, c, makeArgFinder,
+	forEachStaleUnion(target, c,
 		func(a *UnionArg, t *UnionType, okIndices []int) {
 			if ignoreTransient && a.transient {
 				return
@@ -205,8 +192,7 @@ func (c *Call) setDefaultConditions(target *Target, transientOnly bool) bool {
 	// Replace stale conditions with the default values of their correct types.
 	for {
 		replace := map[Arg]Arg{}
-		makeArgFinder := argFinderConstructor(target, c)
-		forEachStaleUnion(target, c, makeArgFinder,
+		forEachStaleUnion(target, c,
 			func(unionArg *UnionArg, unionType *UnionType, okIndices []int) {
 				if transientOnly && !unionArg.transient {
 					return
diff --git a/prog/size.go b/prog/size.go
index 6463821e0..d354aadbe 100644
--- a/prog/size.go
+++ b/prog/size.go
@@ -14,15 +14,15 @@ const (
 	SyscallRef = "syscall"
 )
 
-func (target *Target) assignSizes(args []Arg, fields []Field, parentsMap map[Arg]Arg,
+func (target *Target) assignSizes(args []Arg, fields []Field, parents parentStack,
 	syscallArgs []Arg, syscallFields []Field, autos map[Arg]bool, overlayField int) {
 	for _, arg := range args {
-		target.assignArgSize(arg, args, fields, parentsMap, syscallArgs,
+		target.assignArgSize(arg, args, fields, parents, syscallArgs,
 			syscallFields, autos, overlayField)
 	}
 }
 
-func (target *Target) assignArgSize(arg Arg, args []Arg, fields []Field, parentsMap map[Arg]Arg,
+func (target *Target) assignArgSize(arg Arg, args []Arg, fields []Field, parents parentStack,
 	syscallArgs []Arg, syscallFields []Field, autos map[Arg]bool, overlayField int) {
 	if arg = InnerArg(arg); arg == nil {
 		return // Pointer to optional len field, no need to fill in value.
@@ -39,15 +39,15 @@ func (target *Target) assignArgSize(arg Arg, args []Arg, fields []Field, parents
 	}
 	a := arg.(*ConstArg)
 	if typ.Path[0] == SyscallRef {
-		target.assignSize(a, nil, typ.Path[1:], syscallArgs, syscallFields, parentsMap, 0)
+		target.assignSize(a, nil, typ.Path[1:], syscallArgs, syscallFields, parents, 0)
 	} else {
-		target.assignSize(a, a, typ.Path, args, fields, parentsMap, overlayField)
+		target.assignSize(a, a, typ.Path, args, fields, parents, overlayField)
 	}
 }
 
 func (target *Target) assignSize(dst *ConstArg, pos Arg, path []string, args []Arg,
-	fields []Field, parentsMap map[Arg]Arg, overlayField int) {
-	found := target.findArg(pos, path, args, fields, parentsMap, overlayField)
+	fields []Field, parents parentStack, overlayField int) {
+	found := target.findArg(pos, path, args, fields, parents, overlayField)
 	if found != nil && !found.isAnyPtr {
 		dst.Val = target.computeSize(found.arg, found.offset, dst.Type().(*LenType))
 	}
@@ -59,20 +59,20 @@ type foundArg struct {
 	isAnyPtr bool
 }
 
-func (target *Target) findFieldStruct(buf Arg, path []string, parentsMap map[Arg]Arg) *foundArg {
+func (target *Target) findFieldStruct(buf Arg, path []string, parents parentStack) *foundArg {
 	switch arg := buf.(type) {
 	case *GroupArg:
 		typ := arg.Type().(*StructType)
-		return target.findArg(buf, path, arg.Inner, typ.Fields, parentsMap, typ.OverlayField)
+		return target.findArg(buf, path, arg.Inner, typ.Fields, parents, typ.OverlayField)
 	case *UnionArg:
-		return target.findArg(buf, path, nil, nil, parentsMap, 0)
+		return target.findArg(buf, path, nil, nil, parents, 0)
 	default:
 		panic(fmt.Sprintf("unexpected arg type %#v", arg))
 	}
 }
 
 func (target *Target) findArg(pos Arg, path []string, args []Arg, fields []Field,
-	parentsMap map[Arg]Arg, overlayField int) *foundArg {
+	parents parentStack, overlayField int) *foundArg {
 	elem := path[0]
 	path = path[1:]
 	var offset uint64
@@ -98,23 +98,23 @@ func (target *Target) findArg(pos Arg, path []string, args []Arg, fields []Field
 			return &foundArg{nil, offset, false}
 		}
 		if len(path) != 0 {
-			return target.findFieldStruct(buf, path, parentsMap)
+			return target.findFieldStruct(buf, path, parents)
 		}
 		return &foundArg{buf, offset, false}
 	}
 	if elem == ParentRef {
-		buf := parentsMap[pos]
+		parents, buf := popStack(parents)
 		if len(path) != 0 {
-			return target.findFieldStruct(buf, path, parentsMap)
+			return target.findFieldStruct(buf, path, parents)
 		}
 		return &foundArg{buf, noOffset, false}
 	}
-	for buf := parentsMap[pos]; buf != nil; buf = parentsMap[buf] {
+	for parents, buf := popStack(parents); buf != nil; parents, buf = popStack(parents) {
 		if elem != buf.Type().TemplateName() {
 			continue
 		}
 		if len(path) != 0 {
-			return target.findFieldStruct(buf, path, parentsMap)
+			return target.findFieldStruct(buf, path, parents)
 		}
 		return &foundArg{buf, noOffset, false}
 	}
@@ -162,32 +162,15 @@ func (target *Target) computeSize(arg Arg, offset uint64, lenType *LenType) uint
 	}
 }
 
-func saveToParentsMap(arg Arg, parentsMap map[Arg]Arg) {
-	if _, ok := arg.Type().(*StructType); ok {
-		for _, field := range arg.(*GroupArg).Inner {
-			parentsMap[InnerArg(field)] = arg
-		}
-	}
-	if unionArg, ok := arg.(*UnionArg); ok {
-		parentsMap[InnerArg(unionArg.Option)] = arg
-	}
-}
-
 func (target *Target) assignSizesArray(args []Arg, fields []Field, autos map[Arg]bool) {
-	parentsMap := make(map[Arg]Arg)
-	for _, arg := range args {
-		ForeachSubArg(arg, func(arg Arg, _ *ArgCtx) {
-			saveToParentsMap(arg, parentsMap)
-		})
-	}
-	target.assignSizes(args, fields, parentsMap, args, fields, autos, 0)
+	target.assignSizes(args, fields, nil, args, fields, autos, 0)
 	for _, arg := range args {
-		ForeachSubArg(arg, func(arg Arg, _ *ArgCtx) {
+		foreachSubArgWithStack(arg, func(arg Arg, ctx *ArgCtx) {
 			if typ, ok := arg.Type().(*StructType); ok {
-				target.assignSizes(arg.(*GroupArg).Inner, typ.Fields, parentsMap, args, fields, autos, typ.OverlayField)
+				target.assignSizes(arg.(*GroupArg).Inner, typ.Fields, ctx.parentStack, args, fields, autos, typ.OverlayField)
 			}
 			if v, ok := arg.(*UnionArg); ok {
-				target.assignArgSize(v.Option, nil, nil, parentsMap, args, fields, autos, 0)
+				target.assignArgSize(v.Option, nil, nil, ctx.parentStack, args, fields, autos, 0)
 			}
 		})
 	}