pkg/compiler: support size attribute for unions

5 files changed, 58 insertions, 29 deletions

diff --git a/pkg/compiler/check.go b/pkg/compiler/check.go
index 2a6afae3a..4dac37260 100644
--- a/pkg/compiler/check.go
+++ b/pkg/compiler/check.go
@@ -240,10 +240,11 @@ func (comp *compiler) checkAttributeValues() {
 		switch n := decl.(type) {
 		case *ast.Struct:
 			for _, attr := range n.Attrs {
-				if !n.IsUnion && attr.Ident == "size" {
+				if attr.Ident == "size" {
+					_, typ, name := n.Info()
 					if comp.structIsVarlen(n.Name.Name) {
-						comp.error(attr.Pos, "varlen struct %v has size attribute",
-							n.Name.Name)
+						comp.error(attr.Pos, "varlen %v %v has size attribute",
+							typ, name)
 					}
 					sz := attr.Args[0].Value
 					if sz == 0 || sz > 1<<20 {
@@ -906,7 +907,7 @@ func (comp *compiler) checkVarlen(n *ast.Struct) {
 	// Non-varlen unions can't have varlen fields.
 	// Non-packed structs can't have varlen fields in the middle.
 	if n.IsUnion {
-		if varlen := comp.parseUnionAttrs(n); varlen {
+		if varlen, _ := comp.parseUnionAttrs(n); varlen {
 			return
 		}
 	} else {
diff --git a/pkg/compiler/compiler.go b/pkg/compiler/compiler.go
index 314cbef41..f0d1176e7 100644
--- a/pkg/compiler/compiler.go
+++ b/pkg/compiler/compiler.go
@@ -145,9 +145,11 @@ func (comp *compiler) structIsVarlen(name string) bool {
 		return varlen
 	}
 	s := comp.structs[name]
-	if s.IsUnion && comp.parseUnionAttrs(s) {
-		comp.structVarlen[name] = true
-		return true
+	if s.IsUnion {
+		if varlen, _ := comp.parseUnionAttrs(s); varlen {
+			comp.structVarlen[name] = true
+			return true
+		}
 	}
 	comp.structVarlen[name] = false // to not hang on recursive types
 	varlen := false
@@ -161,18 +163,21 @@ func (comp *compiler) structIsVarlen(name string) bool {
 	return varlen
 }
 
-func (comp *compiler) parseUnionAttrs(n *ast.Struct) (varlen bool) {
+func (comp *compiler) parseUnionAttrs(n *ast.Struct) (varlen bool, size uint64) {
+	size = sizeUnassigned
 	for _, attr := range n.Attrs {
 		switch attr.Ident {
 		case "varlen":
+			if len(attr.Args) != 0 {
+				comp.error(attr.Pos, "%v attribute has args", attr.Ident)
+			}
 			varlen = true
+		case "size":
+			size = comp.parseSizeAttr(n, attr)
 		default:
 			comp.error(attr.Pos, "unknown union %v attribute %v",
 				n.Name.Name, attr.Ident)
 		}
-		if len(attr.Args) != 0 {
-			comp.error(attr.Pos, "%v attribute has args", attr.Ident)
-		}
 	}
 	return
 }
@@ -207,19 +212,7 @@ func (comp *compiler) parseStructAttrs(n *ast.Struct) (packed bool, size, align
 			}
 			align = a
 		case attr.Ident == "size":
-			if len(attr.Args) != 1 {
-				comp.error(attr.Pos, "%v attribute is expected to have 1 argument", attr.Ident)
-			}
-			sz := attr.Args[0]
-			if unexpected, _, ok := checkTypeKind(sz, kindInt); !ok {
-				comp.error(sz.Pos, "unexpected %v, expect int", unexpected)
-				return
-			}
-			if sz.HasColon || len(sz.Args) != 0 {
-				comp.error(sz.Pos, "size attribute has colon or args")
-				return
-			}
-			size = sz.Value
+			size = comp.parseSizeAttr(n, attr)
 		default:
 			comp.error(attr.Pos, "unknown struct %v attribute %v",
 				n.Name.Name, attr.Ident)
@@ -228,6 +221,23 @@ func (comp *compiler) parseStructAttrs(n *ast.Struct) (packed bool, size, align
 	return
 }
 
+func (comp *compiler) parseSizeAttr(n *ast.Struct, attr *ast.Type) uint64 {
+	if len(attr.Args) != 1 {
+		comp.error(attr.Pos, "%v attribute is expected to have 1 argument", attr.Ident)
+		return sizeUnassigned
+	}
+	sz := attr.Args[0]
+	if unexpected, _, ok := checkTypeKind(sz, kindInt); !ok {
+		comp.error(sz.Pos, "unexpected %v, expect int", unexpected)
+		return sizeUnassigned
+	}
+	if sz.HasColon || len(sz.Args) != 0 {
+		comp.error(sz.Pos, "size attribute has colon or args")
+		return sizeUnassigned
+	}
+	return sz.Value
+}
+
 func (comp *compiler) getTypeDesc(t *ast.Type) *typeDesc {
 	if desc := builtinTypes[t.Ident]; desc != nil {
 		return desc
diff --git a/pkg/compiler/consts.go b/pkg/compiler/consts.go
index 789674578..1b008bd4f 100644
--- a/pkg/compiler/consts.go
+++ b/pkg/compiler/consts.go
@@ -87,7 +87,7 @@ func (comp *compiler) extractConsts() map[string]*ConstInfo {
 		switch n := decl.(type) {
 		case *ast.Struct:
 			for _, attr := range n.Attrs {
-				if !n.IsUnion && attr.Ident == "size" {
+				if attr.Ident == "size" {
 					info := getConstInfo(infos, attr.Pos)
 					info.consts[attr.Args[0].Ident] = true
 				}
@@ -199,7 +199,7 @@ func (comp *compiler) patchConsts(consts map[string]uint64) {
 			}
 			if n, ok := decl.(*ast.Struct); ok {
 				for _, attr := range n.Attrs {
-					if !n.IsUnion && attr.Ident == "size" {
+					if attr.Ident == "size" {
 						sz := attr.Args[0]
 						comp.patchIntConst(sz.Pos, &sz.Value,
 							&sz.Ident, consts, &missing)
diff --git a/pkg/compiler/gen.go b/pkg/compiler/gen.go
index 74d1b5d78..5bf239e79 100644
--- a/pkg/compiler/gen.go
+++ b/pkg/compiler/gen.go
@@ -165,14 +165,24 @@ func (comp *compiler) genStructDescs(syscalls []*prog.Syscall) []*prog.KeyedStru
 			if !checkStruct(t.Key, &t.StructDesc) {
 				return
 			}
+			structNode := comp.structNodes[t.StructDesc]
+			varlen, sizeAttr := comp.parseUnionAttrs(structNode)
 			t.TypeSize = 0
-			varlen := comp.parseUnionAttrs(comp.structNodes[t.StructDesc])
 			if !varlen {
 				for _, fld := range t.Fields {
-					if t.TypeSize < fld.Size() {
-						t.TypeSize = fld.Size()
+					sz := fld.Size()
+					if sizeAttr != sizeUnassigned && sz > sizeAttr {
+						comp.error(structNode.Pos, "union %v has size attribute %v"+
+							" which is less than field %v size %v",
+							structNode.Name.Name, sizeAttr, fld.Name(), sz)
+					}
+					if t.TypeSize < sz {
+						t.TypeSize = sz
 					}
 				}
+				if sizeAttr != sizeUnassigned {
+					t.TypeSize = sizeAttr
+				}
 			}
 		}
 	}
diff --git a/pkg/compiler/testdata/errors2.txt b/pkg/compiler/testdata/errors2.txt
index a8d29789a..86148a563 100644
--- a/pkg/compiler/testdata/errors2.txt
+++ b/pkg/compiler/testdata/errors2.txt
@@ -91,6 +91,14 @@ u1 [
 	f2	len[f1, int32]	### len target f1 does not exist
 ]
 
+u2 [
+	f1	int8
+] [size[0]]			### size attribute has bad value 0, expect [1, 1<<20]
+
+u3 [
+	f1	int8
+] [varlen, size[8]]		### varlen union u3 has size attribute
+
 foo$200(a ptr[in, s2])
 foo$201(a ptr[in, s1])
 foo$202(a u1)