pkg/compiler: extend parent reference support in len

Earlier only len[parent, T] was supported and meant the size of the
whole structure.

Logically, len[parent:b, T] should be equivalent to just len[b, T].

Let len[parent:parent:a, T] refer to the structure that encloses the
current one.

Support len fields inside unions.

6 files changed, 180 insertions, 55 deletions

diff --git a/pkg/compiler/check.go b/pkg/compiler/check.go
index aa83f8c8d..3250dac31 100644
--- a/pkg/compiler/check.go
+++ b/pkg/compiler/check.go
@@ -411,7 +411,7 @@ func (comp *compiler) checkLenTargets() {
 		case *ast.Call:
 			for _, arg := range n.Args {
 				checked := make(map[string]bool)
-				parents := []parentDesc{{fields: n.Args}}
+				parents := []parentDesc{{fields: n.Args, call: n.Name.Name}}
 				comp.checkLenType(arg.Type, arg.Type, parents, checked, warned, true)
 			}
 		}
@@ -419,10 +419,18 @@ func (comp *compiler) checkLenTargets() {
 }
 
 type parentDesc struct {
+	call   string
 	name   string
 	fields []*ast.Field
 }
 
+func (pd *parentDesc) String() string {
+	if pd.call != "" {
+		return fmt.Sprintf("<%s>", pd.call)
+	}
+	return pd.name
+}
+
 // templateBase return the part before '[' for full template names.
 func templateBase(name string) string {
 	if pos := strings.IndexByte(name, '['); pos != -1 {
@@ -442,15 +450,22 @@ func (comp *compiler) checkLenType(t0, t *ast.Type, parents []parentDesc,
 	if desc == typeStruct {
 		s := comp.structs[t.Ident]
 		// Prune recursion, can happen even on correct tree via opt pointers.
-		if checked[s.Name.Name] {
+		// There may be several paths to the same type, let's at least look
+		// at the nearest parent.
+		checkName := s.Name.Name
+		if len(parents) > 1 {
+			checkName = parents[len(parents)-2].name + " " + checkName
+		}
+		if checked[checkName] {
 			return
 		}
-		checked[s.Name.Name] = true
+		checked[checkName] = true
 		fields := s.Fields
 		if s.IsUnion {
 			fields = nil
 		}
 		parentName := parentTargetName(s)
+		parents = append([]parentDesc{}, parents...)
 		parents = append(parents, parentDesc{name: parentName, fields: fields})
 		for _, fld := range s.Fields {
 			comp.checkLenType(fld.Type, fld.Type, parents, checked, warned, false)
@@ -471,9 +486,17 @@ func (comp *compiler) checkLenType(t0, t *ast.Type, parents []parentDesc,
 
 func (comp *compiler) checkLenTarget(arg, t0, t *ast.Type, parents []parentDesc, warned map[string]bool) {
 	targets := append([]*ast.Type{arg}, arg.Colon...)
+	const maxParents = 2
 	for i, target := range targets {
-		if target.Ident == prog.ParentRef && len(targets) != 1 {
-			comp.error(target.Pos, "%v can't be part of path expressions", prog.ParentRef)
+		if target.Ident == prog.ParentRef &&
+			(i >= maxParents || targets[0].Ident != prog.ParentRef) {
+			// It's not a fundamental limitation, but it helps prune recursion in checkLenType().
+			// If we need more, we need to adjust the key of the "checked" map.
+			if !warned[parents[len(parents)-1].name] {
+				comp.error(target.Pos, "%v may only stay at the beginning (max %d times)",
+					prog.ParentRef, maxParents)
+			}
+			warned[parents[len(parents)-1].name] = true
 			return
 		}
 		if target.Ident == prog.SyscallRef {
@@ -487,12 +510,36 @@ func (comp *compiler) checkLenTarget(arg, t0, t *ast.Type, parents []parentDesc,
 			}
 		}
 	}
+	// Drop parents from the prefix (it will simplify further code).
+	droppedParents := 0
+	for len(targets) > 0 && targets[0].Ident == prog.ParentRef {
+		target := targets[0]
+		if parents[len(parents)-1].call != "" {
+			comp.error(target.Pos, "%v reached the call (%v)",
+				prog.ParentRef, parents[0].call)
+			return
+		}
+		droppedParents++
+		targets = targets[1:]
+		// Ignore the first "parent" item.
+		if droppedParents > 1 {
+			if len(parents) < 2 {
+				comp.error(target.Pos, "too many %v elements", prog.ParentRef)
+				return
+			}
+			parents = parents[:len(parents)-1]
+		}
+	}
 	comp.checkLenTargetRec(t0, t, targets, parents, warned)
 }
 
 func (comp *compiler) checkLenTargetRec(t0, t *ast.Type, targets []*ast.Type,
 	parents []parentDesc, warned map[string]bool) {
 	if len(targets) == 0 {
+		if t.Ident == "offsetof" {
+			comp.error(t.Pos, "%v must refer to fields", t.Ident)
+			return
+		}
 		return
 	}
 	target := targets[0]
@@ -537,22 +584,20 @@ func (comp *compiler) checkLenTargetRec(t0, t *ast.Type, targets []*ast.Type,
 	}
 	for pi := len(parents) - 1; pi >= 0; pi-- {
 		parent := parents[pi]
-		if parent.name != "" && (parent.name == target.Ident || target.Ident == prog.ParentRef) ||
+		if parent.name != "" && parent.name == target.Ident ||
 			parent.name == "" && target.Ident == prog.SyscallRef {
-			if len(targets) == 0 {
-				if t.Ident == "offsetof" {
-					comp.error(target.Pos, "%v must refer to fields", t.Ident)
-					return
-				}
-			} else {
-				parents1 := make([]parentDesc, pi+1)
-				copy(parents1, parents[:pi+1])
-				comp.checkLenTargetRec(t0, t, targets, parents1, warned)
-			}
+			parents1 := make([]parentDesc, pi+1)
+			copy(parents1, parents[:pi+1])
+			comp.checkLenTargetRec(t0, t, targets, parents1, warned)
 			return
 		}
 	}
-	comp.error(target.Pos, "%v target %v does not exist", t.Ident, target.Ident)
+	warnKey := parents[len(parents)-1].name + " " + target.Pos.String()
+	if !warned[warnKey] {
+		comp.error(target.Pos, "%v target %v does not exist in %s",
+			t.Ident, target.Ident, parents[len(parents)-1].String())
+	}
+	warned[warnKey] = true
 }
 
 func CollectUnused(desc *ast.Description, target *targets.Target, eh ast.ErrorHandler) ([]ast.Node, error) {
diff --git a/pkg/compiler/testdata/all.txt b/pkg/compiler/testdata/all.txt
index edf4bd896..f7639c761 100644
--- a/pkg/compiler/testdata/all.txt
+++ b/pkg/compiler/testdata/all.txt
@@ -95,6 +95,7 @@ len_expr2 {
 	f23	ptr[in, len_expr4]
 	f24	ptr[in, ptr[in, len_expr4]]
 	f25	len[f21:f31, int32]
+	f26	len_union
 }
 
 len_expr3 {
@@ -112,6 +113,11 @@ len_expr4 {
 	f41	int32
 }
 
+len_union [
+	u1	len[parent:parent:f21:f31, int32]
+	u2	int32
+]
+
 foo_len_expr(a ptr[in, len_expr1], b ptr[in, array[int8, 3]])
 
 # Pointer type.
@@ -342,4 +348,4 @@ union_compressed [
 compressed$1(a compressed_image) (no_generate, no_minimize)
 compressed$2(a ptr[in, compressed_image]) (no_generate, no_minimize)
 compressed$3(a ptr[in, struct_compressed]) (no_generate, no_minimize)
-compressed$4(a ptr[in, union_compressed]) (no_generate, no_minimize)
\ No newline at end of file
+compressed$4(a ptr[in, union_compressed]) (no_generate, no_minimize)
diff --git a/pkg/compiler/testdata/errors2.txt b/pkg/compiler/testdata/errors2.txt
index 13e69ef8a..d08f3e01e 100644
--- a/pkg/compiler/testdata/errors2.txt
+++ b/pkg/compiler/testdata/errors2.txt
@@ -86,14 +86,14 @@ foo$sr0(a ptr[in, use_sr])
 foo$100(a int32, b len[a])
 foo$101(a len[a])		### len target a refers to itself
 foo$102(a ptr[in, len[a, int8]])	### len target a refers to itself
-foo$103(a int32, b len[c])	### len target c does not exist
-foo$104(a len[parent])		### len target parent does not exist
+foo$103(a int32, b len[c])	### len target c does not exist in <foo$103>
+foo$104(a len[parent])		### parent reached the call (foo$104)
 foo$105(a ptr[in, int32], b ptr[in, array[len[a, int32]]])
 foo$106(a int32, b ptr[in, csum[a, inet, int32]])
-foo$107(a int32, b ptr[in, csum[c, inet, int32]])	### csum target c does not exist
+foo$107(a int32, b ptr[in, csum[c, inet, int32]])	### csum target c does not exist in <foo$107>
 
 s1 {
-	f1	len[s2, int32]	### len target s2 does not exist
+	f1	len[s2, int32]	### len target s2 does not exist in s1
 }
 
 s2 {
@@ -123,12 +123,12 @@ s8 {
 } [align[7]]			###  bad struct s8 alignment 7 (must be a sane power of 2)
 
 u0 [
-	f	len[f1, int32]	### len target f1 does not exist
+	f	len[f1, int32]	### len target f1 does not exist in u0
 ]
 
 u1 [
 	f1	ptr[in, array[int8]]
-	f2	len[f1, int32]	### len target f1 does not exist
+	f2	len[f1, int32]	### len target f1 does not exist in u1
 ]
 
 u2 [
@@ -144,17 +144,17 @@ foo$201(a ptr[in, s1])
 foo$202(a ptr[in, s7])
 foo$203(a u1)
 
-foo$204(a ptr[in, slen1])
+foo$204(a ptr[in, slen1], b ptr[in, no_slen22])
 
 slen1 {
 	f0 len[parent, int32]
-	f1 len[parent:f0, int32]	### parent can't be part of path expressions
-	f2 len[slen21:a, int32]		### len target a does not exist
-	f3 len[f0:parent, int32]	### parent can't be part of path expressions
+	f1 len[parent:f0, int32]
+	f2 len[slen21:a, int32]		### len target a does not exist in slen2
+	f3 len[f0:parent, int32]	### parent may only stay at the beginning (max 2 times)
 	f4 len[slen2:f, int32]		### len path slen2 does not refer to a struct
 	f5 len[slen22:f, int32]		### len path slen22 does not refer to a struct
 	f6 len[syscall, int32]		### no argument name after syscall reference
-	f7 len[syscall:b, int32]	### len target b does not exist
+	f7 len[syscall:c, int32]	### len target c does not exist in <foo$204>
 	f8 offsetof[parent, int32]	### offsetof must refer to fields
 	f9 offsetof[slen1, int32]	### offsetof must refer to fields
 	f10 len[f0:syscall:b, int32]	### syscall can't be in the middle of path expressions
@@ -168,13 +168,22 @@ slen1 {
 slen2 {
 	f	int32
 	f1	slen3
+	f2	len[parent:parent:slen22, int32] ### len target slen22 does not exist in no_slen22
 }
 
 slen3 [
 	a int32
 	b int32
+	c len[parent:parent:parent:x, int32] ### parent may only stay at the beginning (max 2 times)
+	d len[parent:a, int32] ### len target a does not exist in slen3
 ]
 
+no_slen22 {
+	# This one violates the parent:parent:slen22 chain.
+	f slen2
+}
+
+
 # Resource ctor tests.
 
 resource r100[int32]		### resource r100 can't be created (never mentioned as a syscall return value or output argument/field)
diff --git a/prog/size.go b/prog/size.go
index 11f1ccaf8..70ec678b6 100644
--- a/prog/size.go
+++ b/prog/size.go
@@ -17,32 +17,44 @@ const (
 func (target *Target) assignSizes(args []Arg, fields []Field, parentsMap map[Arg]Arg,
 	syscallArgs []Arg, syscallFields []Field, autos map[Arg]bool, overlayField int) {
 	for _, arg := range args {
-		if arg = InnerArg(arg); arg == nil {
-			continue // Pointer to optional len field, no need to fill in value.
-		}
-		typ, ok := arg.Type().(*LenType)
-		if !ok {
-			continue
-		}
-		if autos != nil {
-			if !autos[arg] {
-				continue
-			}
-			delete(autos, arg)
-		}
-		a := arg.(*ConstArg)
-		if typ.Path[0] == SyscallRef {
-			target.assignSize(a, nil, typ.Path[1:], syscallArgs, syscallFields, parentsMap, 0)
-		} else {
-			target.assignSize(a, a, typ.Path, args, fields, parentsMap, overlayField)
+		target.assignArgSize(arg, args, fields, parentsMap, syscallArgs,
+			syscallFields, autos, overlayField)
+	}
+}
+
+func (target *Target) assignArgSize(arg Arg, args []Arg, fields []Field, parentsMap map[Arg]Arg,
+	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.
+	}
+	typ, ok := arg.Type().(*LenType)
+	if !ok {
+		return
+	}
+	if autos != nil {
+		if !autos[arg] {
+			return
 		}
+		delete(autos, arg)
+	}
+	a := arg.(*ConstArg)
+	if typ.Path[0] == SyscallRef {
+		target.assignSize(a, nil, typ.Path[1:], syscallArgs, syscallFields, parentsMap, 0)
+	} else {
+		target.assignSize(a, a, typ.Path, args, fields, parentsMap, overlayField)
 	}
 }
 
 func (target *Target) assignSizeStruct(dst *ConstArg, buf Arg, path []string, parentsMap map[Arg]Arg) {
-	arg := buf.(*GroupArg)
-	typ := arg.Type().(*StructType)
-	target.assignSize(dst, buf, path, arg.Inner, typ.Fields, parentsMap, typ.OverlayField)
+	switch arg := buf.(type) {
+	case *GroupArg:
+		typ := arg.Type().(*StructType)
+		target.assignSize(dst, buf, path, arg.Inner, typ.Fields, parentsMap, typ.OverlayField)
+	case *UnionArg:
+		target.assignSize(dst, buf, path, nil, nil, parentsMap, 0)
+	default:
+		panic(fmt.Sprintf("unexpected arg type %v", arg))
+	}
 }
 
 func (target *Target) assignSize(dst *ConstArg, pos Arg, path []string, args []Arg,
@@ -132,15 +144,22 @@ 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) {
-			if _, ok := arg.Type().(*StructType); ok {
-				for _, field := range arg.(*GroupArg).Inner {
-					parentsMap[InnerArg(field)] = arg
-				}
-			}
+			saveToParentsMap(arg, parentsMap)
 		})
 	}
 	target.assignSizes(args, fields, parentsMap, args, fields, autos, 0)
@@ -149,6 +168,9 @@ func (target *Target) assignSizesArray(args []Arg, fields []Field, autos map[Arg
 			if typ, ok := arg.Type().(*StructType); ok {
 				target.assignSizes(arg.(*GroupArg).Inner, typ.Fields, parentsMap, args, fields, autos, typ.OverlayField)
 			}
+			if v, ok := arg.(*UnionArg); ok {
+				target.assignArgSize(v.Option, nil, nil, parentsMap, args, fields, autos, 0)
+			}
 		})
 	}
 }
diff --git a/prog/size_test.go b/prog/size_test.go
index ea8b7a58a..2f7dc9ce2 100644
--- a/prog/size_test.go
+++ b/prog/size_test.go
@@ -166,5 +166,17 @@ test$length11(&(0x7f0000000000)=ANY=[@ANYBLOB="11"], 0x42)
 test$length30(&(0x7f0000000000)=ANY=[@ANYBLOB="11"], 0x42, &(0x7f0000000000)=0x43, 0x44)
 `,
 		},
+		{
+			In:  "test$length32(&(0x7f0000000000)={[0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0], {0x0}, &(0x7f0000000040)={0x0}})",
+			Out: "test$length32(&(0x7f0000000000)={[0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0], {0x8}, &(0x7f0000000040)={0x8}})",
+		},
+		{
+			In:  "test$length33(&(0x7f0000000000)={[0x0, 0x0, 0x0, 0x0], 0x0})",
+			Out: "test$length33(&(0x7f0000000000)={[0x0, 0x0, 0x0, 0x0], 0x4})",
+		},
+		{
+			In:  "test$length34(&(0x7f0000000000)={[0x0, 0x0, 0x0, 0x0], &(0x7f0000000040)=@u1=0x0})",
+			Out: "test$length34(&(0x7f0000000000)={[0x0, 0x0, 0x0, 0x0], &(0x7f0000000040)=@u1=0x4})",
+		},
 	})
 }
diff --git a/sys/test/test.txt b/sys/test/test.txt
index 1b63cd525..788eb8484 100644
--- a/sys/test/test.txt
+++ b/sys/test/test.txt
@@ -486,6 +486,37 @@ offsetof0 {
 	o7	offsetof[f7, int32]
 }
 
+parent_parent_helper {
+	f1	array[int32, 8]
+	f2	parent_parent_length
+	f3	ptr[in, parent_parent_length]
+}
+
+parent_parent_length {
+	f1	len[parent:parent:f1, int32]
+}
+
+test$length32(a0 ptr[in, parent_parent_helper])
+
+one_parent_length {
+	f1	array[int32, 4]
+	f2	len[parent:f1, int32]
+}
+
+test$length33(a0 ptr[in, one_parent_length])
+
+parent_union_struct {
+	f1	array[int32, 4]
+	f2	ptr[in, parent_union]
+}
+
+parent_union [
+	u1	len[parent:parent:f1, int32]
+	u2	int32
+]
+
+test$length34(a0 ptr[in, parent_union_struct])
+
 # Big endian
 
 test$end0(a0 ptr[in, syz_end_int_struct])