fix statedata fields not being updated correctly (#255)

beacon_chain/beacon_node.nim

@@ -202,7 +202,7 @@ template withState(
   ## TODO async transformations will lead to a race where cache gets updated
   ##      while waiting for future to complete - catch this here somehow?
 
-  updateState(pool, cache, blockSlot)
+  updateStateData(pool, cache, blockSlot)
 
   template state(): BeaconState {.inject.} = cache.data
   template blck(): BlockRef {.inject.} = cache.blck
@@ -369,11 +369,12 @@ proc proposeBlock(node: BeaconNode,
         signature: ValidatorSig(), # we need the rest of the block first!
       )
 
-    let ok = updateState(state, newBlock, {skipValidation})
-    doAssert ok # TODO: err, could this fail somehow?
-    root = hash_tree_root(state)
+    var tmpState = state
 
-    newBlock.state_root = root
+    let ok = updateState(tmpState, newBlock, {skipValidation})
+    doAssert ok # TODO: err, could this fail somehow?
+
+    newBlock.state_root = hash_tree_root(state)
 
     let blockRoot = signed_root(newBlock)
 

beacon_chain/block_pool.nim

@@ -137,7 +137,7 @@ proc addSlotMapping(pool: BlockPool, slot: uint64, br: BlockRef) =
       s.add(v)
   pool.blocksBySlot.mgetOrPut(slot, @[]).addIfMissing(br)
 
-proc updateState*(
+proc updateStateData*(
   pool: BlockPool, state: var StateData, bs: BlockSlot) {.gcsafe.}
 
 proc add*(
@@ -184,7 +184,7 @@ proc add*(
 
     # TODO if the block is from the future, we should not be resolving it (yet),
     #      but maybe we should use it as a hint that our clock is wrong?
-    updateState(pool, state, BlockSlot(blck: parent, slot: blck.slot - 1))
+    updateStateData(pool, state, BlockSlot(blck: parent, slot: blck.slot - 1))
 
     if not updateState(state.data, blck, {}):
       # TODO find a better way to log all this block data
@@ -204,6 +204,9 @@ proc add*(
     # Resolved blocks should be stored in database
     pool.db.putBlock(blockRoot, blck)
 
+    state.root = blck.state_root
+    state.blck = blockRef
+
     # This block *might* have caused a justification - make sure we stow away
     # that information:
     let
@@ -417,7 +420,7 @@ proc rewindState(pool: BlockPool, state: var StateData, bs: BlockSlot):
 
   ancestors
 
-proc updateState*(pool: BlockPool, state: var StateData, bs: BlockSlot) =
+proc updateStateData*(pool: BlockPool, state: var StateData, bs: BlockSlot) =
   ## Rewind or advance state such that it matches the given block and slot -
   ## this may include replaying from an earlier snapshot if blck is on a
   ## different branch or has advanced to a higher slot number than slot
@@ -427,9 +430,12 @@ proc updateState*(pool: BlockPool, state: var StateData, bs: BlockSlot) =
   # We need to check the slot because the state might have moved forwards
   # without blocks
   if state.blck.root == bs.blck.root and state.data.slot <= bs.slot:
-    # Might be that we're moving to the same block but later slot
-    skipSlots(state.data, bs.slot) do (state: BeaconState):
-      pool.maybePutState(state, bs.blck)
+    if state.data.slot != bs.slot:
+      # Might be that we're moving to the same block but later slot
+      skipSlots(state.data, bs.slot) do (state: BeaconState):
+        pool.maybePutState(state, bs.blck)
+
+      state.root = hash_tree_root(state.data)
 
     return # State already at the right spot
 
@@ -488,7 +494,7 @@ proc updateHead*(pool: BlockPool, state: var StateData, blck: BlockRef) =
   pool.db.putHeadBlock(blck.root)
 
   # Start off by making sure we have the right state
-  updateState(pool, state, BlockSlot(blck: blck, slot: blck.slot))
+  updateStateData(pool, state, BlockSlot(blck: blck, slot: blck.slot))
 
   if lastHead != blck.parent:
     notice "Updated head with new parent",

tests/test_block_pool.nim

@@ -46,6 +46,7 @@ suite "Block pool processing":
     check:
       b1Ref.isSome()
       b1Ref.get().refs.root == b1Root
+      hash_tree_root(state.data) == state.root
 
   test "Reverse order block add & get":
     var
@@ -67,6 +68,8 @@ suite "Block pool processing":
 
     discard pool.add(state, b1Root, b1)
 
+    check: hash_tree_root(state.data) == state.root
+
     let
       b1r = pool.get(b1Root)
       b2r = pool.get(b2Root)
@@ -87,5 +90,6 @@ suite "Block pool processing":
       pool2 = BlockPool.init(db)
 
     check:
+      hash_tree_root(state.data) == state.root
       pool2.get(b1Root).isSome()
       pool2.get(b2Root).isSome()