pkg/compiler: add size attribute for structs

The size attribute allows to pad a struct up to the specified size.

10 files changed, 162 insertions, 33 deletions

diff --git a/pkg/compiler/check.go b/pkg/compiler/check.go
index 58738fa3e..2a6afae3a 100644
--- a/pkg/compiler/check.go
+++ b/pkg/compiler/check.go
@@ -21,7 +21,8 @@ func (comp *compiler) typecheck() {
 }
 
 func (comp *compiler) check() {
-	comp.checkConsts()
+	comp.checkTypeValues()
+	comp.checkAttributeValues()
 	comp.checkUsed()
 	comp.checkRecursion()
 	comp.checkLenTargets()
@@ -215,7 +216,7 @@ func (comp *compiler) checkTypes() {
 	}
 }
 
-func (comp *compiler) checkConsts() {
+func (comp *compiler) checkTypeValues() {
 	for _, decl := range comp.desc.Nodes {
 		switch decl.(type) {
 		case *ast.Call, *ast.Struct, *ast.Resource, *ast.TypeDef:
@@ -234,6 +235,27 @@ func (comp *compiler) checkConsts() {
 	}
 }
 
+func (comp *compiler) checkAttributeValues() {
+	for _, decl := range comp.desc.Nodes {
+		switch n := decl.(type) {
+		case *ast.Struct:
+			for _, attr := range n.Attrs {
+				if !n.IsUnion && attr.Ident == "size" {
+					if comp.structIsVarlen(n.Name.Name) {
+						comp.error(attr.Pos, "varlen struct %v has size attribute",
+							n.Name.Name)
+					}
+					sz := attr.Args[0].Value
+					if sz == 0 || sz > 1<<20 {
+						comp.error(attr.Args[0].Pos, "size attribute has bad value %v"+
+							", expect [1, 1<<20]", sz)
+					}
+				}
+			}
+		}
+	}
+}
+
 func (comp *compiler) checkLenTargets() {
 	for _, decl := range comp.desc.Nodes {
 		switch n := decl.(type) {
@@ -536,8 +558,12 @@ func (comp *compiler) checkStruct(ctx checkCtx, n *ast.Struct) {
 		comp.checkType(ctx, f.Type, flags)
 	}
 	for _, attr := range n.Attrs {
-		if attr.Ident == "" || attr.HasColon {
-			comp.error(attr.Pos, "bad struct/union attribute")
+		if unexpected, _, ok := checkTypeKind(attr, kindIdent); !ok {
+			comp.error(attr.Pos, "unexpected %v, expect attribute", unexpected)
+			return
+		}
+		if attr.HasColon {
+			comp.error(attr.Pos2, "unexpected ':'")
 			return
 		}
 	}
@@ -884,7 +910,7 @@ func (comp *compiler) checkVarlen(n *ast.Struct) {
 			return
 		}
 	} else {
-		if packed, _ := comp.parseStructAttrs(n); packed {
+		if packed, _, _ := comp.parseStructAttrs(n); packed {
 			return
 		}
 	}
diff --git a/pkg/compiler/compiler.go b/pkg/compiler/compiler.go
index 8e20a4e67..314cbef41 100644
--- a/pkg/compiler/compiler.go
+++ b/pkg/compiler/compiler.go
@@ -93,12 +93,16 @@ func Compile(desc *ast.Description, consts map[string]uint64, target *targets.Ta
 		eh(w.pos, w.msg)
 	}
 	syscalls := comp.genSyscalls()
-	return &Prog{
+	prg := &Prog{
 		Resources:   comp.genResources(),
 		Syscalls:    syscalls,
 		StructDescs: comp.genStructDescs(syscalls),
 		Unsupported: comp.unsupported,
 	}
+	if comp.errors != 0 {
+		return nil
+	}
+	return prg
 }
 
 type compiler struct {
@@ -136,6 +140,27 @@ func (comp *compiler) warning(pos ast.Pos, msg string, args ...interface{}) {
 	comp.warnings = append(comp.warnings, warn{pos, fmt.Sprintf(msg, args...)})
 }
 
+func (comp *compiler) structIsVarlen(name string) bool {
+	if varlen, ok := comp.structVarlen[name]; ok {
+		return varlen
+	}
+	s := comp.structs[name]
+	if s.IsUnion && comp.parseUnionAttrs(s) {
+		comp.structVarlen[name] = true
+		return true
+	}
+	comp.structVarlen[name] = false // to not hang on recursive types
+	varlen := false
+	for _, fld := range s.Fields {
+		if comp.isVarlen(fld.Type) {
+			varlen = true
+			break
+		}
+	}
+	comp.structVarlen[name] = varlen
+	return varlen
+}
+
 func (comp *compiler) parseUnionAttrs(n *ast.Struct) (varlen bool) {
 	for _, attr := range n.Attrs {
 		switch attr.Ident {
@@ -146,20 +171,30 @@ func (comp *compiler) parseUnionAttrs(n *ast.Struct) (varlen bool) {
 				n.Name.Name, attr.Ident)
 		}
 		if len(attr.Args) != 0 {
-			comp.error(attr.Pos, "%v attribute had args", attr.Ident)
+			comp.error(attr.Pos, "%v attribute has args", attr.Ident)
 		}
 	}
 	return
 }
 
-func (comp *compiler) parseStructAttrs(n *ast.Struct) (packed bool, align uint64) {
+func (comp *compiler) parseStructAttrs(n *ast.Struct) (packed bool, size, align uint64) {
+	size = sizeUnassigned
 	for _, attr := range n.Attrs {
 		switch {
 		case attr.Ident == "packed":
+			if len(attr.Args) != 0 {
+				comp.error(attr.Pos, "%v attribute has args", attr.Ident)
+			}
 			packed = true
 		case attr.Ident == "align_ptr":
+			if len(attr.Args) != 0 {
+				comp.error(attr.Pos, "%v attribute has args", attr.Ident)
+			}
 			align = comp.ptrSize
 		case strings.HasPrefix(attr.Ident, "align_"):
+			if len(attr.Args) != 0 {
+				comp.error(attr.Pos, "%v attribute has args", attr.Ident)
+			}
 			a, err := strconv.ParseUint(attr.Ident[6:], 10, 64)
 			if err != nil {
 				comp.error(attr.Pos, "bad struct %v alignment %v",
@@ -171,13 +206,24 @@ func (comp *compiler) parseStructAttrs(n *ast.Struct) (packed bool, align uint64
 					n.Name.Name, a)
 			}
 			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
 		default:
 			comp.error(attr.Pos, "unknown struct %v attribute %v",
 				n.Name.Name, attr.Ident)
 		}
-		if len(attr.Args) != 0 {
-			comp.error(attr.Pos, "%v attribute had args", attr.Ident)
-		}
 	}
 	return
 }
diff --git a/pkg/compiler/consts.go b/pkg/compiler/consts.go
index b4f674124..789674578 100644
--- a/pkg/compiler/consts.go
+++ b/pkg/compiler/consts.go
@@ -83,6 +83,18 @@ func (comp *compiler) extractConsts() map[string]*ConstInfo {
 		}
 	}
 
+	for _, decl := range comp.desc.Nodes {
+		switch n := decl.(type) {
+		case *ast.Struct:
+			for _, attr := range n.Attrs {
+				if !n.IsUnion && attr.Ident == "size" {
+					info := getConstInfo(infos, attr.Pos)
+					info.consts[attr.Args[0].Ident] = true
+				}
+			}
+		}
+	}
+
 	comp.desc.Walk(ast.Recursive(func(n0 ast.Node) {
 		if n, ok := n0.(*ast.Int); ok {
 			info := getConstInfo(infos, n.Pos)
@@ -185,6 +197,15 @@ func (comp *compiler) patchConsts(consts map[string]uint64) {
 						&v.Ident, consts, &missing)
 				}
 			}
+			if n, ok := decl.(*ast.Struct); ok {
+				for _, attr := range n.Attrs {
+					if !n.IsUnion && attr.Ident == "size" {
+						sz := attr.Args[0]
+						comp.patchIntConst(sz.Pos, &sz.Value,
+							&sz.Ident, consts, &missing)
+					}
+				}
+			}
 			if missing == "" {
 				continue
 			}
diff --git a/pkg/compiler/consts_test.go b/pkg/compiler/consts_test.go
index 13780a128..522062991 100644
--- a/pkg/compiler/consts_test.go
+++ b/pkg/compiler/consts_test.go
@@ -37,6 +37,7 @@ func TestExtractConsts(t *testing.T) {
 		"CONST6", "CONST7", "CONST8", "CONST9", "CONST10",
 		"CONST11", "CONST12", "CONST13", "CONST14", "CONST15",
 		"CONST16", "CONST17", "CONST18", "CONST19", "CONST20",
+		"CONST21",
 	}
 	sort.Strings(wantConsts)
 	if !reflect.DeepEqual(info.Consts, wantConsts) {
diff --git a/pkg/compiler/gen.go b/pkg/compiler/gen.go
index ddcec2252..74d1b5d78 100644
--- a/pkg/compiler/gen.go
+++ b/pkg/compiler/gen.go
@@ -130,6 +130,7 @@ func (comp *compiler) genStructDescs(syscalls []*prog.Syscall) []*prog.KeyedStru
 			if !checkStruct(t.Key, &t.StructDesc) {
 				return
 			}
+			structNode := comp.structNodes[t.StructDesc]
 			// Add paddings, calculate size, mark bitfields.
 			varlen := false
 			for _, f := range t.Fields {
@@ -138,7 +139,7 @@ func (comp *compiler) genStructDescs(syscalls []*prog.Syscall) []*prog.KeyedStru
 				}
 			}
 			comp.markBitfields(t.Fields)
-			packed, alignAttr := comp.parseStructAttrs(comp.structNodes[t.StructDesc])
+			packed, sizeAttr, alignAttr := comp.parseStructAttrs(structNode)
 			t.Fields = comp.addAlignment(t.Fields, varlen, packed, alignAttr)
 			t.AlignAttr = alignAttr
 			t.TypeSize = 0
@@ -148,6 +149,17 @@ func (comp *compiler) genStructDescs(syscalls []*prog.Syscall) []*prog.KeyedStru
 						t.TypeSize += f.Size()
 					}
 				}
+				if sizeAttr != sizeUnassigned {
+					if t.TypeSize > sizeAttr {
+						comp.error(structNode.Pos, "struct %v has size attribute %v"+
+							" which is less than struct size %v",
+							structNode.Name.Name, sizeAttr, t.TypeSize)
+					}
+					if pad := sizeAttr - t.TypeSize; pad != 0 {
+						t.Fields = append(t.Fields, genPad(pad))
+					}
+					t.TypeSize = sizeAttr
+				}
 			}
 		case *prog.UnionType:
 			if !checkStruct(t.Key, &t.StructDesc) {
@@ -308,7 +320,7 @@ func (comp *compiler) typeAlign(t0 prog.Type) uint64 {
 	case *prog.ArrayType:
 		return comp.typeAlign(t.Type)
 	case *prog.StructType:
-		packed, alignAttr := comp.parseStructAttrs(comp.structNodes[t.StructDesc])
+		packed, _, alignAttr := comp.parseStructAttrs(comp.structNodes[t.StructDesc])
 		if alignAttr != 0 {
 			return alignAttr // overrided by user attribute
 		}
diff --git a/pkg/compiler/testdata/all.txt b/pkg/compiler/testdata/all.txt
index 442f08774..081bb1aab 100644
--- a/pkg/compiler/testdata/all.txt
+++ b/pkg/compiler/testdata/all.txt
@@ -183,6 +183,17 @@ foo$templ4(a ptr[in, templ_struct1[3]])
 foo$templ5(a ptr[in, templ_struct1[3]])
 foo$templ6(a ptr[in, templ_struct4])
 
+
+# Structs.
+
+s0 {
+	f1	int8
+} [size[64]]
+
+s1 {
+	f1	int8
+} [size[C2]]
+
 # Unions.
 
 u0 [
diff --git a/pkg/compiler/testdata/consts.txt b/pkg/compiler/testdata/consts.txt
index b0bf8d4f1..f7c8daa35 100644
--- a/pkg/compiler/testdata/consts.txt
+++ b/pkg/compiler/testdata/consts.txt
@@ -31,3 +31,7 @@ type templ1[C] {
 	f2	const[C, int8]
 }
 foo$1(a ptr[in, templ1[CONST20]])
+
+str2 {
+	f1	int8
+} [size[CONST21]]
diff --git a/pkg/compiler/testdata/errors.txt b/pkg/compiler/testdata/errors.txt
index 2b3a7539f..2d61d9e76 100644
--- a/pkg/compiler/testdata/errors.txt
+++ b/pkg/compiler/testdata/errors.txt
@@ -166,11 +166,23 @@ s8 {
 
 s9 {
 	f1	int8
-} ["foo"[0]]			### bad struct/union attribute
+} ["foo"[0]]			### unexpected string "foo", expect attribute
 
 s10 {
 	f1	int8
-} [packed[0]]			### packed attribute had args
+} [packed[0]]			### packed attribute has args
+
+s11 {
+	f1	int8
+} [size["foo"]]			### unexpected string "foo", expect int
+
+s12 {
+	f1	int8
+} [size[0:1]]			### size attribute has colon or args
+
+s13 {
+	f1	int8
+} [size[0[0]]]			### size attribute has colon or args
 
 u3 [
 	f1	int8
diff --git a/pkg/compiler/testdata/errors2.txt b/pkg/compiler/testdata/errors2.txt
index f00a6afb4..a8d29789a 100644
--- a/pkg/compiler/testdata/errors2.txt
+++ b/pkg/compiler/testdata/errors2.txt
@@ -73,6 +73,19 @@ s2 {
 	f1	s1
 }
 
+s3 {
+	f1	int8
+} [size[0]]			### size attribute has bad value 0, expect [1, 1<<20]
+
+s4 {
+	f1	int8
+} [size[1000000000]]		### size attribute has bad value 1000000000, expect [1, 1<<20]
+
+s6 {
+	f1	int8
+	f2	array[int8]
+} [size[10]]			### varlen struct s6 has size attribute
+
 u1 [
 	f1	ptr[in, array[int8]]
 	f2	len[f1, int32]	### len target f1 does not exist
diff --git a/pkg/compiler/types.go b/pkg/compiler/types.go
index 4c1123521..a7734262a 100644
--- a/pkg/compiler/types.go
+++ b/pkg/compiler/types.go
@@ -611,24 +611,7 @@ func init() {
 		return canBeArg, false
 	}
 	typeStruct.Varlen = func(comp *compiler, t *ast.Type, args []*ast.Type) bool {
-		if varlen, ok := comp.structVarlen[t.Ident]; ok {
-			return varlen
-		}
-		s := comp.structs[t.Ident]
-		if s.IsUnion && comp.parseUnionAttrs(s) {
-			comp.structVarlen[t.Ident] = true
-			return true
-		}
-		comp.structVarlen[t.Ident] = false // to not hang on recursive types
-		varlen := false
-		for _, fld := range s.Fields {
-			if comp.isVarlen(fld.Type) {
-				varlen = true
-				break
-			}
-		}
-		comp.structVarlen[t.Ident] = varlen
-		return varlen
+		return comp.structIsVarlen(t.Ident)
 	}
 	typeStruct.ZeroSize = func(comp *compiler, t *ast.Type, args []*ast.Type) bool {
 		for _, fld := range comp.structs[t.Ident].Fields {