1 files changed, 48 insertions, 26 deletions
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)
+			}
 		})
 	}
 }