prog: minor refactoring around arguments

Introduce isUsed(arg) helper, use it in several places.
Move method definitions closer to their types.
Simplify presence check for ArgUsed.Used() in several places.

4 files changed, 76 insertions, 59 deletions

diff --git a/prog/encoding.go b/prog/encoding.go
index 2b7999efc..7f0844dea 100644
--- a/prog/encoding.go
+++ b/prog/encoding.go
@@ -33,7 +33,7 @@ func (p *Prog) Serialize() []byte {
 	vars := make(map[Arg]int)
 	varSeq := 0
 	for _, c := range p.Calls {
-		if len(*c.Ret.(ArgUsed).Used()) != 0 {
+		if isUsed(c.Ret) {
 			fmt.Fprintf(buf, "r%v = ", varSeq)
 			vars[c.Ret] = varSeq
 			varSeq++
@@ -58,7 +58,7 @@ func serialize(arg Arg, buf *bytes.Buffer, vars map[Arg]int, varSeq *int) {
 		fmt.Fprintf(buf, "nil")
 		return
 	}
-	if used, ok := arg.(ArgUsed); ok && len(*used.Used()) != 0 {
+	if isUsed(arg) {
 		fmt.Fprintf(buf, "<r%v=>", *varSeq)
 		vars[arg] = *varSeq
 		*varSeq++
diff --git a/prog/encodingexec.go b/prog/encodingexec.go
index 937f3c322..c8855ebd6 100644
--- a/prog/encodingexec.go
+++ b/prog/encodingexec.go
@@ -94,8 +94,7 @@ func (p *Prog) SerializeForExec(buffer []byte, pid int) (int, error) {
 		foreachArg(c, func(arg, _ Arg, _ *[]Arg) {
 			if a, ok := arg.(*PointerArg); ok && a.Res != nil {
 				foreachSubargOffset(a.Res, func(arg1 Arg, offset uint64) {
-					used, ok := arg1.(ArgUsed)
-					if (ok && len(*used.Used()) != 0) || csumUses[arg1] {
+					if isUsed(arg1) || csumUses[arg1] {
 						w.args[arg1] = argInfo{Addr: p.Target.physicalAddr(arg) + offset}
 					}
 					if _, ok := arg1.(*GroupArg); ok {
@@ -164,13 +163,13 @@ func (p *Prog) SerializeForExec(buffer []byte, pid int) (int, error) {
 		for _, arg := range c.Args {
 			w.writeArg(arg, pid, csumMap)
 		}
-		if len(*c.Ret.(ArgUsed).Used()) != 0 {
+		if isUsed(c.Ret) {
 			w.args[c.Ret] = argInfo{Idx: instrSeq}
 		}
 		instrSeq++
 		// Generate copyout instructions that persist interesting return values.
 		foreachArg(c, func(arg, base Arg, _ *[]Arg) {
-			if used, ok := arg.(ArgUsed); !ok || len(*used.Used()) == 0 {
+			if !isUsed(arg) {
 				return
 			}
 			switch arg.(type) {
diff --git a/prog/prog.go b/prog/prog.go
index 66de938b3..f78aac582 100644
--- a/prog/prog.go
+++ b/prog/prog.go
@@ -23,6 +23,24 @@ type Arg interface {
 	Size() uint64
 }
 
+// ArgUser is interface of an argument that uses value of another output argument.
+type ArgUser interface {
+	Uses() *Arg
+}
+
+// ArgUsed is interface of an argument that can be used by other arguments.
+type ArgUsed interface {
+	Used() *map[Arg]bool
+}
+
+func isUsed(arg Arg) bool {
+	used, ok := arg.(ArgUsed)
+	if !ok {
+		return false
+	}
+	return len(*used.Used()) != 0
+}
+
 type ArgCommon struct {
 	typ Type
 }
@@ -37,6 +55,10 @@ type ConstArg struct {
 	Val uint64
 }
 
+func MakeConstArg(t Type, v uint64) Arg {
+	return &ConstArg{ArgCommon: ArgCommon{typ: t}, Val: v}
+}
+
 func (arg *ConstArg) Size() uint64 {
 	return arg.typ.Size()
 }
@@ -79,6 +101,16 @@ type PointerArg struct {
 	Res        Arg    // pointee
 }
 
+func MakePointerArg(t Type, page uint64, off int, npages uint64, obj Arg) Arg {
+	return &PointerArg{
+		ArgCommon:  ArgCommon{typ: t},
+		PageIndex:  page,
+		PageOffset: off,
+		PagesNum:   npages,
+		Res:        obj,
+	}
+}
+
 func (arg *PointerArg) Size() uint64 {
 	return arg.typ.Size()
 }
@@ -125,6 +157,10 @@ type GroupArg struct {
 	Inner []Arg
 }
 
+func MakeGroupArg(t Type, inner []Arg) Arg {
+	return &GroupArg{ArgCommon: ArgCommon{typ: t}, Inner: inner}
+}
+
 func (arg *GroupArg) Size() uint64 {
 	typ0 := arg.Type()
 	if !typ0.Varlen() {
@@ -160,6 +196,10 @@ type UnionArg struct {
 	OptionType Type
 }
 
+func MakeUnionArg(t Type, opt Arg, typ Type) Arg {
+	return &UnionArg{ArgCommon: ArgCommon{typ: t}, Option: opt, OptionType: typ}
+}
+
 func (arg *UnionArg) Size() uint64 {
 	if !arg.Type().Varlen() {
 		return arg.Type().Size()
@@ -179,10 +219,31 @@ type ResultArg struct {
 	uses  map[Arg]bool // ArgResult args that use this arg
 }
 
+func MakeResultArg(t Type, r Arg, v uint64) Arg {
+	arg := &ResultArg{ArgCommon: ArgCommon{typ: t}, Res: r, Val: v}
+	if r == nil {
+		return arg
+	}
+	used := r.(ArgUsed)
+	if *used.Used() == nil {
+		*used.Used() = make(map[Arg]bool)
+	}
+	(*used.Used())[arg] = true
+	return arg
+}
+
 func (arg *ResultArg) Size() uint64 {
 	return arg.typ.Size()
 }
 
+func (arg *ResultArg) Used() *map[Arg]bool {
+	return &arg.uses
+}
+
+func (arg *ResultArg) Uses() *Arg {
+	return &arg.Res
+}
+
 // Used for ResourceType and VmaType.
 // This argument denotes syscall return value.
 type ReturnArg struct {
@@ -190,30 +251,18 @@ type ReturnArg struct {
 	uses map[Arg]bool // ArgResult args that use this arg
 }
 
-func (arg *ReturnArg) Size() uint64 {
-	panic("not called")
-}
-
-type ArgUsed interface {
-	Used() *map[Arg]bool
+func MakeReturnArg(t Type) Arg {
+	return &ReturnArg{ArgCommon: ArgCommon{typ: t}}
 }
 
-func (arg *ResultArg) Used() *map[Arg]bool {
-	return &arg.uses
+func (arg *ReturnArg) Size() uint64 {
+	panic("not called")
 }
 
 func (arg *ReturnArg) Used() *map[Arg]bool {
 	return &arg.uses
 }
 
-type ArgUser interface {
-	Uses() *Arg
-}
-
-func (arg *ResultArg) Uses() *Arg {
-	return &arg.Res
-}
-
 // Returns inner arg for pointer args.
 func InnerArg(arg Arg) Arg {
 	if t, ok := arg.Type().(*PtrType); ok {
@@ -248,39 +297,6 @@ func encodeValue(value uint64, size uint64, bigEndian bool) uint64 {
 	}
 }
 
-func MakeConstArg(t Type, v uint64) Arg {
-	return &ConstArg{ArgCommon: ArgCommon{typ: t}, Val: v}
-}
-
-func MakeResultArg(t Type, r Arg, v uint64) Arg {
-	arg := &ResultArg{ArgCommon: ArgCommon{typ: t}, Res: r, Val: v}
-	if r == nil {
-		return arg
-	}
-	used := r.(ArgUsed)
-	if *used.Used() == nil {
-		*used.Used() = make(map[Arg]bool)
-	}
-	(*used.Used())[arg] = true
-	return arg
-}
-
-func MakePointerArg(t Type, page uint64, off int, npages uint64, obj Arg) Arg {
-	return &PointerArg{ArgCommon: ArgCommon{typ: t}, PageIndex: page, PageOffset: off, PagesNum: npages, Res: obj}
-}
-
-func MakeGroupArg(t Type, inner []Arg) Arg {
-	return &GroupArg{ArgCommon: ArgCommon{typ: t}, Inner: inner}
-}
-
-func MakeUnionArg(t Type, opt Arg, typ Type) Arg {
-	return &UnionArg{ArgCommon: ArgCommon{typ: t}, Option: opt, OptionType: typ}
-}
-
-func MakeReturnArg(t Type) Arg {
-	return &ReturnArg{ArgCommon: ArgCommon{typ: t}}
-}
-
 func defaultArg(t Type) Arg {
 	switch typ := t.(type) {
 	case *IntType, *ConstType, *FlagsType, *LenType, *ProcType, *CsumType:
@@ -430,7 +446,7 @@ func (p *Prog) replaceArgCheck(c *Call, arg, arg1 Arg, calls []*Call) {
 func (p *Prog) removeArg(c *Call, arg0 Arg) {
 	foreachSubarg(arg0, func(arg, _ Arg, _ *[]Arg) {
 		if a, ok := arg.(*ResultArg); ok && a.Res != nil {
-			if _, ok := (*a.Res.(ArgUsed).Used())[arg]; !ok {
+			if !(*a.Res.(ArgUsed).Used())[arg] {
 				panic("broken tree")
 			}
 			delete(*a.Res.(ArgUsed).Used(), arg)
diff --git a/prog/validation.go b/prog/validation.go
index d0f6d0bd2..288a8172b 100644
--- a/prog/validation.go
+++ b/prog/validation.go
@@ -48,7 +48,8 @@ func (c *Call) validate(ctx *validCtx) error {
 		if used, ok := arg.(ArgUsed); ok {
 			for u := range *used.Used() {
 				if u == nil {
-					return fmt.Errorf("syscall %v: nil reference in uses for arg %+v", c.Meta.Name, arg)
+					return fmt.Errorf("syscall %v: nil reference in uses for arg %+v",
+						c.Meta.Name, arg)
 				}
 				ctx.uses[u] = arg
 			}
@@ -246,8 +247,9 @@ func (c *Call) validate(ctx *validCtx) error {
 				return fmt.Errorf("syscall %v: result arg '%v' references out-of-tree result: %p%+v -> %p%+v",
 					c.Meta.Name, a.Type().Name(), arg, arg, a.Res, a.Res)
 			}
-			if _, ok := (*a.Res.(ArgUsed).Used())[arg]; !ok {
-				return fmt.Errorf("syscall %v: result arg '%v' has broken link (%+v)", c.Meta.Name, a.Type().Name(), *a.Res.(ArgUsed).Used())
+			if !(*a.Res.(ArgUsed).Used())[arg] {
+				return fmt.Errorf("syscall %v: result arg '%v' has broken link (%+v)",
+					c.Meta.Name, a.Type().Name(), *a.Res.(ArgUsed).Used())
 			}
 		case *ReturnArg:
 			switch a.Type().(type) {