diff --git a/client/llb/llb.go b/client/llb/llb.go
index 461a85308..50696c792 100644
--- a/client/llb/llb.go
+++ b/client/llb/llb.go
@@ -132,7 +132,7 @@ func (eo *exec) marshalTo(list [][]byte, cache map[digest.Digest]struct{}) (dige
 		return eo.mounts[i].dest < eo.mounts[j].dest
 	})
 
-	var outputIndex int64 = 0
+	var outputIndex pb.OutputIndex
 
 	for _, m := range eo.mounts {
 		var dgst digest.Digest
@@ -145,21 +145,23 @@ func (eo *exec) marshalTo(list [][]byte, cache map[digest.Digest]struct{}) (dige
 		if err != nil {
 			return "", list, err
 		}
-		inputIndex := len(pop.Inputs)
-		for i := range pop.Inputs {
-			if pop.Inputs[i].Digest == dgst {
-				inputIndex = i
+
+		var mountIndex pb.OutputIndex
+		if m.parent != nil {
+			mountIndex = m.parent.outputIndex
+		}
+
+		inputIndex := pb.InputIndex(len(pop.Inputs))
+		for i, inp := range pop.Inputs {
+			if inp.Digest == dgst && inp.Index == mountIndex {
+				inputIndex = pb.InputIndex(i)
 				break
 			}
 		}
 		if dgst == "" {
 			inputIndex = pb.Empty
 		}
-		if inputIndex == len(pop.Inputs) {
-			var mountIndex int64
-			if m.parent != nil {
-				mountIndex = m.parent.outputIndex
-			}
+		if inputIndex == pb.InputIndex(len(pop.Inputs)) {
 			pop.Inputs = append(pop.Inputs, &pb.Input{
 				Digest: dgst,
 				Index:  mountIndex,
@@ -167,7 +169,7 @@ func (eo *exec) marshalTo(list [][]byte, cache map[digest.Digest]struct{}) (dige
 		}
 
 		pm := &pb.Mount{
-			Input:    int64(inputIndex),
+			Input:    inputIndex,
 			Dest:     m.dest,
 			Readonly: m.readonly,
 		}
@@ -177,7 +179,7 @@ func (eo *exec) marshalTo(list [][]byte, cache map[digest.Digest]struct{}) (dige
 		} else {
 			pm.Output = pb.SkipOutput
 		}
-		m.outputIndex = outputIndex - 1
+		m.outputIndex = pm.Output
 		peo.Mounts = append(peo.Mounts, pm)
 	}
 
@@ -191,9 +193,8 @@ type mount struct {
 	// either parent or source has to be set
 	parent      *mount
 	source      *source
-	hasOutput   bool  // TODO: remove
-	outputIndex int64 // filled in after marshal
-	state       *ExecState
+	hasOutput   bool           // TODO: remove
+	outputIndex pb.OutputIndex // filled in after marshal
 }
 
 func (m *mount) marshalTo(list [][]byte, cache map[digest.Digest]struct{}) (digest.Digest, [][]byte, error) {
@@ -210,7 +211,7 @@ func (m *mount) marshalTo(list [][]byte, cache map[digest.Digest]struct{}) (dige
 			po := &pb.Op{}
 			po.Inputs = append(po.Inputs, &pb.Input{
 				Digest: dgst,
-				Index:  int64(m.outputIndex),
+				Index:  m.outputIndex,
 			})
 			return appendResult(po, list, cache)
 		}
diff --git a/solver/pb/const.go b/solver/pb/const.go
index 2000ce0e5..7bfdf43bd 100644
--- a/solver/pb/const.go
+++ b/solver/pb/const.go
@@ -1,5 +1,8 @@
 package pb
 
+type InputIndex int64
+type OutputIndex int64
+
 const RootMount = "/"
-const SkipOutput = -1 // TODO: custom type
-const Empty = -1      // TODO: custom type
+const SkipOutput OutputIndex = -1
+const Empty InputIndex = -1
diff --git a/solver/pb/ops.pb.go b/solver/pb/ops.pb.go
index 147be7daa..9c986215d 100644
--- a/solver/pb/ops.pb.go
+++ b/solver/pb/ops.pb.go
@@ -171,7 +171,7 @@ func _Op_OneofUnmarshaler(msg proto.Message, tag, wire int, b *proto.Buffer) (bo
 
 type Input struct {
 	Digest github_com_opencontainers_go_digest.Digest `protobuf:"bytes,1,opt,name=digest,proto3,customtype=github.com/opencontainers/go-digest.Digest" json:"digest"`
-	Index  int64                                      `protobuf:"varint,2,opt,name=index,proto3" json:"index,omitempty"`
+	Index  OutputIndex                                `protobuf:"varint,2,opt,name=index,proto3,customtype=OutputIndex" json:"index"`
 }
 
 func (m *Input) Reset()         { *m = Input{} }
@@ -212,11 +212,11 @@ func (m *Meta) String() string { return proto.CompactTextString(m) }
 func (*Meta) ProtoMessage()    {}
 
 type Mount struct {
-	Input    int64  `protobuf:"varint,1,opt,name=input,proto3" json:"input,omitempty"`
-	Selector string `protobuf:"bytes,2,opt,name=selector,proto3" json:"selector,omitempty"`
-	Dest     string `protobuf:"bytes,3,opt,name=dest,proto3" json:"dest,omitempty"`
-	Output   int64  `protobuf:"varint,4,opt,name=output,proto3" json:"output,omitempty"`
-	Readonly bool   `protobuf:"varint,5,opt,name=readonly,proto3" json:"readonly,omitempty"`
+	Input    InputIndex  `protobuf:"varint,1,opt,name=input,proto3,customtype=InputIndex" json:"input"`
+	Selector string      `protobuf:"bytes,2,opt,name=selector,proto3" json:"selector,omitempty"`
+	Dest     string      `protobuf:"bytes,3,opt,name=dest,proto3" json:"dest,omitempty"`
+	Output   OutputIndex `protobuf:"varint,4,opt,name=output,proto3,customtype=OutputIndex" json:"output"`
+	Readonly bool        `protobuf:"varint,5,opt,name=readonly,proto3" json:"readonly,omitempty"`
 }
 
 func (m *Mount) Reset()         { *m = Mount{} }
@@ -240,8 +240,8 @@ func (m *CopyOp) GetSrc() []*CopySource {
 }
 
 type CopySource struct {
-	Input    int64  `protobuf:"varint,1,opt,name=input,proto3" json:"input,omitempty"`
-	Selector string `protobuf:"bytes,2,opt,name=selector,proto3" json:"selector,omitempty"`
+	Input    InputIndex `protobuf:"varint,1,opt,name=input,proto3,customtype=InputIndex" json:"input"`
+	Selector string     `protobuf:"bytes,2,opt,name=selector,proto3" json:"selector,omitempty"`
 }
 
 func (m *CopySource) Reset()         { *m = CopySource{} }
@@ -1085,7 +1085,7 @@ func (m *Input) Unmarshal(data []byte) error {
 				}
 				b := data[iNdEx]
 				iNdEx++
-				m.Index |= (int64(b) & 0x7F) << shift
+				m.Index |= (OutputIndex(b) & 0x7F) << shift
 				if b < 0x80 {
 					break
 				}
@@ -1405,7 +1405,7 @@ func (m *Mount) Unmarshal(data []byte) error {
 				}
 				b := data[iNdEx]
 				iNdEx++
-				m.Input |= (int64(b) & 0x7F) << shift
+				m.Input |= (InputIndex(b) & 0x7F) << shift
 				if b < 0x80 {
 					break
 				}
@@ -1482,7 +1482,7 @@ func (m *Mount) Unmarshal(data []byte) error {
 				}
 				b := data[iNdEx]
 				iNdEx++
-				m.Output |= (int64(b) & 0x7F) << shift
+				m.Output |= (OutputIndex(b) & 0x7F) << shift
 				if b < 0x80 {
 					break
 				}
@@ -1681,7 +1681,7 @@ func (m *CopySource) Unmarshal(data []byte) error {
 				}
 				b := data[iNdEx]
 				iNdEx++
-				m.Input |= (int64(b) & 0x7F) << shift
+				m.Input |= (InputIndex(b) & 0x7F) << shift
 				if b < 0x80 {
 					break
 				}
diff --git a/solver/pb/ops.proto b/solver/pb/ops.proto
index 9fc18ca7d..7bbc46bc3 100644
--- a/solver/pb/ops.proto
+++ b/solver/pb/ops.proto
@@ -15,7 +15,7 @@ message Op {
 
 message Input {
 	string digest = 1 [(gogoproto.customtype) = "github.com/opencontainers/go-digest.Digest", (gogoproto.nullable) = false];
-	int64 index = 2;
+	int64 index = 2 [(gogoproto.customtype) = "OutputIndex", (gogoproto.nullable) = false];
 }
 
 message ExecOp {
@@ -30,10 +30,10 @@ message Meta {
 }
 
 message Mount {
-	int64 input = 1;
+	int64 input = 1 [(gogoproto.customtype) = "InputIndex", (gogoproto.nullable) = false];
 	string selector = 2;
 	string dest = 3;
-	int64 output = 4;
+	int64 output = 4 [(gogoproto.customtype) = "OutputIndex", (gogoproto.nullable) = false];
 	bool readonly = 5;
 }
 
@@ -43,7 +43,7 @@ message CopyOp {
 }
 
 message CopySource {
-	int64 input = 1;
+	int64 input = 1 [(gogoproto.customtype) = "InputIndex", (gogoproto.nullable) = false];
 	string selector = 2;
 }