beacon/light: add LightChain unit tests and fix bugs

beacon/light/light_chain.go

@@ -107,11 +107,12 @@ func (lc *LightChain) loadChainRange() {
 			return
 		}
 		if cr.ChainInit {
-			if lc.chainHead, err = lc.getHeaderBySlot(cr.ChainHead); err != nil {
+			// cannot use getHeaderBySlot until chainHead and chainTail are initialized
+			if lc.chainHead, err = lc.getHeader(cr.ChainHead, lc.getCanonicalHash(cr.ChainHead)); err != nil {
 				log.Error("Chain head not found")
 				return
 			}
-			if lc.chainTail, err = lc.getHeaderBySlot(cr.ChainTail); err != nil {
+			if lc.chainTail, err = lc.getHeader(cr.ChainTail, lc.getCanonicalHash(cr.ChainTail)); err != nil {
 				log.Error("Chain tail not found")
 				return
 			}
@@ -144,6 +145,7 @@ func (lc *LightChain) storeChainRange(batch ethdb.Batch) {
 	if cr == lc.lastStoredRange {
 		return
 	}
+	lc.lastStoredRange = cr
 	rangeEnc, err := rlp.EncodeToBytes(&cr)
 	if err != nil {
 		log.Error("Failed to encode chain range data", "error", err)
@@ -152,10 +154,10 @@ func (lc *LightChain) storeChainRange(batch ethdb.Batch) {
 	batch.Put(chainRangeKey, rangeEnc)
 }
 
-// SetChainHead sets the canonical chain head and also finds the new tail if it
+// SetHead sets the canonical chain head and also finds the new tail if it
 // does not share a common ancestor with the old head. The state range is also
 // automatically updated so that it applies to the new canonical chain.
-func (lc *LightChain) SetChainHead(head types.Header) {
+func (lc *LightChain) SetHead(head types.Header) {
 	lc.lock.Lock()
 	defer lc.lock.Unlock()
 
@@ -176,16 +178,29 @@ func (lc *LightChain) SetChainHead(head types.Header) {
 		lc.deleteCanonicalHash(batch, slot)
 	}
 	lc.chainHead = head
+	var (
+		hasStateRange, lastHasState bool //applies to the new chain section after the common ancestor
+		firstState, lastState       types.Header
+	)
 	for !lc.IsCanonical(head) {
 		lc.storeCanonicalHash(batch, head.Slot, head.Hash())
+		if lc.HasStateProof(head) {
+			if !lastHasState {
+				hasStateRange, lastHasState = true, true
+				lastState = head
+			}
+			firstState = head
+		} else {
+			lastHasState = false
+		}
 		parent, err := lc.GetParent(head)
 		if err != nil {
 			for slot := lc.chainTail.Slot; slot < head.Slot; slot++ {
 				lc.deleteCanonicalHash(batch, slot)
 			}
 			lc.chainTail = head
-			lc.stateInit = false
+			// set state range to the new section if there is one
-			lc.reinitStateChain(batch, head)
+			lc.stateInit, lc.stateTail, lc.stateHead = hasStateRange, firstState, lastState
 			return
 		}
 		for slot := parent.Slot + 1; slot < head.Slot; slot++ {
@@ -193,7 +208,9 @@ func (lc *LightChain) SetChainHead(head types.Header) {
 		}
 		head = parent
 	}
-	if lc.stateInit && lc.stateHead.Slot >= head.Slot {
+	// head is now at the common ancestor
+	if lc.stateInit && lc.stateHead.Slot > head.Slot {
+		// first revert state range to common ancestor
 		if head.Slot >= lc.stateTail.Slot {
 			lc.stateHead = head
 		} else {
@@ -201,9 +218,13 @@ func (lc *LightChain) SetChainHead(head types.Header) {
 		}
 	}
 	if lc.stateInit {
-		lc.extendStateHead(batch)
+		// extend with new state range if they are adjacent (otherwise leave old continuous range)
+		if lastHasState && lc.stateHead == head {
+			lc.stateHead = lastState
+		}
 	} else {
-		lc.reinitStateChain(batch, head)
+		// set state range to the new section if there is one
+		lc.stateInit, lc.stateTail, lc.stateHead = hasStateRange, firstState, lastState
 	}
 }
 
@@ -554,9 +575,14 @@ func (lc *LightChain) GetStateProof(header types.Header) (merkle.MultiProof, err
 	return merkle.MultiProof{Format: lc.stateProofFormat, Values: state.Values}, nil
 }
 
+// StateProofFormat returns the expected state proof format for the given header.
+func (lc *LightChain) StateProofFormat(header types.Header) merkle.ProofFormat {
+	return lc.stateProofFormat
+}
+
 // AddStateProof adds a state proof. If it belongs to a canonical header then
 // the state range is also updated.
-func (lc *LightChain) AddStateProof(header types.Header, proof merkle.MultiProof) error {
+func (lc *LightChain) AddStateProof(header types.Header, proof merkle.MultiProof) (err error) {
 	lc.lock.Lock()
 	defer lc.lock.Unlock()
 
@@ -567,6 +593,11 @@ func (lc *LightChain) AddStateProof(header types.Header, proof merkle.MultiProof
 		return ErrInvalidStateRoot
 	}
 	batch := lc.db.NewBatch()
+	defer func() {
+		if err = batch.Write(); err != nil {
+			log.Error("Failed to write batch to database", "error", err)
+		}
+	}()
 
 	stateEnc, err := rlp.EncodeToBytes(&stateProofData{Values: proof.Values})
 	if err != nil {
@@ -589,20 +620,10 @@ func (lc *LightChain) AddStateProof(header types.Header, proof merkle.MultiProof
 	} else if header.Slot < lc.stateTail.Slot && header.Slot >= lc.chainTail.Slot {
 		lc.extendStateTail(batch)
 	}
-
 	lc.storeChainRange(batch)
-	if err := batch.Write(); err != nil {
-		log.Error("Failed to write batch to database", "error", err)
-		return err
-	}
 	return nil
 }
 
-// StateProofFormat returns the expected state proof format for the given header.
-func (lc *LightChain) StateProofFormat(header types.Header) merkle.ProofFormat {
-	return lc.stateProofFormat
-}
-
 func (lc *LightChain) extendStateHead(batch ethdb.Batch) {
 	for slot := lc.stateHead.Slot + 1; slot <= lc.chainHead.Slot; slot++ {
 		if header, err := lc.getHeaderBySlot(slot); err == nil {
@@ -619,8 +640,9 @@ func (lc *LightChain) extendStateTail(batch ethdb.Batch) {
 	if lc.stateTail.Slot == 0 {
 		return
 	}
-	for slot := lc.stateTail.Slot - 1; slot >= lc.chainTail.Slot; slot-- {
+	for slotP1 := lc.stateTail.Slot; slotP1 > lc.chainTail.Slot; slotP1-- {
-		if header, err := lc.getHeaderBySlot(slot); err == nil {
+		// slotP1 == slot+1 to avoid uint64 underflow
+		if header, err := lc.getHeaderBySlot(slotP1 - 1); err == nil {
 			if lc.HasStateProof(header) {
 				lc.stateTail = header
 			} else {
@@ -629,15 +651,3 @@ func (lc *LightChain) extendStateTail(batch ethdb.Batch) {
 		}
 	}
 }
-
-func (lc *LightChain) reinitStateChain(batch ethdb.Batch, header types.Header) {
-	for slot := header.Slot; slot <= lc.chainHead.Slot; slot++ {
-		if header, err := lc.getHeaderBySlot(slot); err == nil && lc.HasStateProof(header) {
-			lc.stateInit = true
-			lc.stateHead = header
-			lc.stateTail = header
-			lc.extendStateHead(batch)
-			return
-		}
-	}
-}

beacon/light/light_chain_test.go

@@ -23,41 +23,240 @@ import (
 	"github.com/ethereum/go-ethereum/beacon/light/types"
 	"github.com/ethereum/go-ethereum/beacon/merkle"
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/ethdb/memorydb"
 )
 
-func makeChain(tail types.Header, headSlot uint64, format merkle.ProofFormat) (headers []types.Header, stateProofs []merkle.MultiProof) {
+func TestLightChainSetHead(t *testing.T) {
-	valueCount := merkle.ValueCount(format)
+	for _, reload := range []bool{false, true} {
-	for tail.Slot < headSlot {
+		c := newChainTest(t)
+		a1, a2 := c.makeChain(types.Header{}, 100, false, false)
+		b1, b2 := c.makeChain(a2, 150, true, false)
+		c1, c2 := c.makeChain(b2, 200, true, true)
+		d1, d2 := c.makeChain(c2, 250, true, true)
+		_, e2 := c.makeChain(d2, 300, true, false)
+		f1, f2 := c.makeChain(c2, 270, true, true)
+		c.checkCanonical(a1, false)
+		c.checkCanonical(f2, false)
+		c.checkTail(e2, b1)
+		c.checkTail(f2, b1)
+		c.checkRange(false, types.Header{}, types.Header{}, false, types.Header{}, types.Header{})
+		c.chain.SetHead(f2)
+		c.checkCanonical(a2, false)
+		c.checkCanonical(d1, false)
+		c.checkCanonical(e2, false)
+		if reload {
+			c.reloadChain()
+		}
+		c.checkRange(true, b1, f2, true, c1, f2)
+		c.chain.SetHead(e2)
+		c.checkCanonical(f1, false)
+		c.checkCanonical(f2, false)
+		c.checkRange(true, b1, e2, true, c1, d2)
+		_, g2 := c.makeChain(b2, 220, true, false)
+		if reload {
+			c.reloadChain()
+		}
+		c.chain.SetHead(g2)
+		c.checkCanonical(c1, false)
+		c.checkCanonical(f2, false)
+		c.checkRange(true, b1, g2, false, types.Header{}, types.Header{})
+		_, h2 := c.makeChain(types.Header{}, 100, false, false)
+		if reload {
+			c.reloadChain()
+		}
+		i1, i2 := c.makeChain(h2, 150, true, false)
+		j1, j2 := c.makeChain(i2, 200, true, true)
+		c.chain.SetHead(i2)
+		c.checkTail(j2, i1)
+		c.checkCanonical(b1, false)
+		c.checkCanonical(j1, false)
+		c.checkRange(true, i1, i2, false, types.Header{}, types.Header{})
+		if reload {
+			c.reloadChain()
+		}
+		c.chain.SetHead(j2)
+		c.checkRange(true, i1, j2, true, j1, j2)
+		c.chain.SetHead(i2)
+		c.checkCanonical(j1, false)
+		if reload {
+			c.reloadChain()
+		}
+		c.checkRange(true, i1, i2, false, types.Header{}, types.Header{})
+	}
+}
+
+func TestLightChainExtendHeaderTail(t *testing.T) {
+	for _, reload := range []bool{false, true} {
+		for _, reverse := range []bool{false, true} {
+			c := newChainTest(t)
+			a1, a2 := c.makeChain(types.Header{}, 50, false, false)
+			b1, b2 := c.makeChain(a2, 100, true, false)
+			c.chain.SetHead(b2)
+			c.checkTail(b2, b1)
+			c.checkRange(true, b1, b2, false, types.Header{}, types.Header{})
+			if reload {
+				c.reloadChain()
+			}
+			if reverse {
+				for i := len(c.headers) - 1; i >= 0; i-- {
+					c.chain.AddHeader(c.headers[i])
+				}
+			} else {
+				for _, header := range c.headers {
+					c.chain.AddHeader(header)
+				}
+			}
+			if reload {
+				c.reloadChain()
+			}
+			c.checkTail(b2, a1)
+			c.checkRange(true, a1, b2, false, types.Header{}, types.Header{})
+		}
+	}
+}
+
+func TestLightChainExtendStateRange(t *testing.T) {
+	for _, reload := range []bool{false /*, true*/} {
+		for _, reverse := range []bool{false, true} {
+			c := newChainTest(t)
+			a1, a2 := c.makeChain(types.Header{}, 50, true, false)
+			b1, b2 := c.makeChain(a2, 100, true, true)
+			_, c2 := c.makeChain(b2, 150, true, false)
+			c.chain.SetHead(c2)
+			c.checkRange(true, a1, c2, true, b1, b2)
+			if reload {
+				c.reloadChain()
+			}
+			if reverse {
+				for i := len(c.stateProofs) - 1; i >= 0; i-- {
+					sp := c.stateProofs[i]
+					c.chain.AddStateProof(sp.header, sp.proof)
+				}
+			} else {
+				for _, sp := range c.stateProofs {
+					c.chain.AddStateProof(sp.header, sp.proof)
+				}
+			}
+			if reload {
+				c.reloadChain()
+			}
+			c.checkRange(true, a1, c2, true, a1, c2)
+		}
+	}
+}
+
+type chainTest struct {
+	t           *testing.T
+	db          *memorydb.Database
+	proofFormat merkle.ProofFormat
+	chain       *LightChain
+	headers     []types.Header // not added to the chain yet
+	stateProofs []testProof    // not added to the chain yet
+}
+
+type testProof struct {
+	header types.Header
+	proof  merkle.MultiProof
+}
+
+func newChainTest(t *testing.T) *chainTest {
+	c := &chainTest{
+		t:           t,
+		db:          memorydb.New(),
+		proofFormat: merkle.NewIndexMapFormat().AddLeaf(42, nil).AddLeaf(67, nil),
+	}
+	c.chain = NewLightChain(c.db, c.proofFormat)
+	return c
+}
+
+func (c *chainTest) checkRange(chainInit bool, chainTail, chainHead types.Header, stateInit bool, stateTail, stateHead types.Header) {
+	ch, ct, ci := c.chain.HeaderRange()
+	if ci != chainInit || (ci && (ct != chainTail || ch != chainHead)) {
+		c.t.Errorf("Incorrect header chain range (expected: %v %d %d, got: %v %d %d)", chainInit, chainTail.Slot, chainHead.Slot, ci, ct.Slot, ch.Slot)
+	}
+	if chainInit {
+		c.checkCanonical(chainTail, true)
+		c.checkCanonical(chainHead, true)
+	}
+	sh, st, si := c.chain.StateProofRange()
+	if si != stateInit || (si && (st != stateTail || sh != stateHead)) {
+		c.t.Errorf("Incorrect state proof range (expected: %v %d %d, got: %v %d %d)", stateInit, stateTail.Slot, stateHead.Slot, si, st.Slot, sh.Slot)
+	}
+	if stateInit {
+		c.checkCanonical(stateTail, true)
+		c.checkCanonical(stateHead, true)
+	}
+}
+
+func (c *chainTest) checkCanonical(header types.Header, expected bool) {
+	if canonical := c.chain.IsCanonical(header); canonical != expected {
+		c.t.Errorf("Canonical status of header at slot %d is incorrect (expected: %v, got: %v)", header.Slot, expected, canonical)
+	}
+}
+
+func (c *chainTest) checkTail(header, expTail types.Header) {
+	for {
+		if parent, err := c.chain.GetParent(header); err == nil {
+			header = parent
+		} else {
+			break
+		}
+	}
+	if header != expTail {
+		c.t.Errorf("Incorrect chain tail found by repeated GetParent (expected slot: %d, got: %d)", expTail.Slot, header.Slot)
+	}
+}
+
+func (c *chainTest) reloadChain() {
+	c.chain = NewLightChain(c.db, c.proofFormat)
+}
+
+func (c *chainTest) makeChain(from types.Header, targetHeadSlot uint64, addHeaders, addStateProofs bool) (tail, head types.Header) {
+	head = from
+	valueCount := merkle.ValueCount(c.proofFormat)
+	for head.Slot < targetHeadSlot {
 		var (
 			slot       uint64
 			parentRoot common.Hash
 		)
-		if tail != (types.Header{}) {
+		if head != (types.Header{}) {
-			slot = tail.Slot + 1
+			slot = head.Slot + 1
-			parentRoot = tail.Hash()
+			parentRoot = head.Hash()
 		}
-		for slot < headSlot && rand.Intn(5) == 0 {
+		for slot < targetHeadSlot && rand.Intn(5) == 0 {
 			slot++
 		}
 		stateProof := merkle.MultiProof{
-			Format: format,
+			Format: c.proofFormat,
 			Values: make(merkle.Values, valueCount),
 		}
 		for i, _ := range stateProof.Values {
 			stateProof.Values[i] = merkle.Value(randomHash())
 		}
-		header := types.Header{
+		head = types.Header{
 			Slot:          slot,
 			ProposerIndex: uint64(rand.Intn(10000)),
 			BodyRoot:      randomHash(),
 			StateRoot:     stateProof.RootHash(),
 			ParentRoot:    parentRoot,
 		}
-
+		if tail == (types.Header{}) {
-		headers = append(headers, header)
+			tail = head
-		stateProofs = append(stateProofs, stateProof)
+		}
-		tail = header
+		if addHeaders {
+			c.chain.AddHeader(head)
+		} else {
+			c.headers = append(c.headers, head)
+		}
+		if addStateProofs {
+			if err := c.chain.AddStateProof(head, stateProof); err != nil {
+				c.t.Fatalf("AddStateProof failed (error: %v)", err)
+			}
+		} else {
+			c.stateProofs = append(c.stateProofs, testProof{head, stateProof})
+		}
 	}
+	return
 }
 
 func randomHash() (hash common.Hash) {

beacon/light/sync/header_sync.go

@@ -82,7 +82,7 @@ func (s *HeaderSync) Process(env *request.Environment) {
 	chainHead, chainTail, chainInit := s.chain.HeaderRange()
 	if !chainInit {
 		s.chain.AddHeader(s.targetHead)
-		s.chain.SetChainHead(s.targetHead)
+		s.chain.SetHead(s.targetHead)
 		s.selfTrigger.Trigger()
 		s.chainTrigger.Trigger()
 	}
@@ -104,7 +104,7 @@ func (s *HeaderSync) Process(env *request.Environment) {
 func (s *HeaderSync) trySyncHead(env *request.Environment, chainTailSlot uint64) bool {
 	for {
 		if s.syncPtr.Slot <= chainTailSlot || s.chain.IsCanonical(s.syncPtr) {
-			s.chain.SetChainHead(s.targetHead)
+			s.chain.SetHead(s.targetHead)
 			s.chainTrigger.Trigger()
 			return true
 		}

beacon/merkle/binary_merkle.go

@@ -82,7 +82,7 @@ func IsEqual(a, b ProofFormat) bool {
 // ValueCount returns the number of merkle values required for this proof format
 func ValueCount(f ProofFormat) int {
 	if f == nil {
-		return 0
+		return 1
 	}
 	l, r := f.Children()
 	return ValueCount(l) + ValueCount(r)