Refactor out TypeCastSteps to reduce code duplication

Author: ethanfrey

go/ops.go

@@ -43,24 +43,6 @@ func asInner(op *ProofOp) (*InnerOp, error) {
 	return inner.Inner, nil
 }
 
-func (op *ProofOp) Apply(args ...[]byte) ([]byte, error) {
-	o := op.Op
-	switch o.(type) {
-	case *ProofOp_Leaf:
-		if len(args) != 2 {
-			return nil, fmt.Errorf("Need key and value args, got %d", len(args))
-		}
-		return op.GetLeaf().Apply(args[0], args[1])
-	case *ProofOp_Inner:
-		if len(args) != 1 {
-			return nil, fmt.Errorf("Need one child hash, got %d", len(args))
-		}
-		return op.GetInner().Apply(args[0])
-	default:
-		panic("Unknown proof op")
-	}
-}
-
 func (op *LeafOp) Apply(key []byte, value []byte) ([]byte, error) {
 	if len(key) == 0 {
 		return nil, fmt.Errorf("Leaf op needs key")

go/proof.go

@@ -18,11 +18,11 @@ var IavlSpec = &ProofSpec{
 // You must validate the result is what you have in a header.
 // Returns error if the calculations cannot be performed.
 func (p *ExistenceProof) Calculate() ([]byte, error) {
-	if len(p.Steps) == 0 {
-		return nil, fmt.Errorf("Existence Proof needs at least one step")
+	leaf, inners, err := p.TypeCastSteps()
+	if err != nil {
+		return nil, err
 	}
 
-	leaf, inners := p.Steps[0], p.Steps[1:]
 	// leaf step takes the key and value as input
 	res, err := leaf.Apply(p.Key, p.Value)
 	if err != nil {
@@ -40,31 +40,42 @@ func (p *ExistenceProof) Calculate() ([]byte, error) {
 }
 
 func (p *ExistenceProof) CheckAgainstSpec(spec *ProofSpec) error {
-	if len(p.Steps) == 0 {
-		return fmt.Errorf("Existence Proof needs at least one step")
-	}
-
-	leafOp, inners := p.Steps[0], p.Steps[1:]
-	leaf, err := asLeaf(leafOp)
+	leaf, inners, err := p.TypeCastSteps()
 	if err != nil {
 		return err
 	}
 	err = checkLeaf(leaf, spec.LeafSpec)
 	if err != nil {
 		return err
 	}
-	for _, innerOp := range inners {
-		inner, err := asInner(innerOp)
-		if err != nil {
-			return err
-		}
+	for _, inner := range inners {
 		if err := checkInner(inner, spec.LeafSpec.Prefix); err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
+func (p *ExistenceProof) TypeCastSteps() (*LeafOp, []*InnerOp, error) {
+	if len(p.Steps) == 0 {
+		return nil, nil, fmt.Errorf("Existence Proof needs at least one step")
+	}
+	leafOp, inners := p.Steps[0], p.Steps[1:]
+	leaf, err := asLeaf(leafOp)
+	if err != nil {
+		return nil, nil, err
+	}
+	var result []*InnerOp
+	for _, innerOp := range inners {
+		inner, err := asInner(innerOp)
+		if err != nil {
+			return nil, nil, err
+		}
+		result = append(result, inner)
+	}
+	return leaf, result, nil
+}
+
 func checkLeaf(leaf *LeafOp, spec *LeafOp) error {
 	if leaf.Hash != spec.Hash {
 		return fmt.Errorf("Unexpected HashOp: %d", leaf.Hash)