[Storehouse] Improve InMemoryRegisterStore's IsBlockExecuted method

engine/execution/storehouse/in_memory_register_store.go

@@ -246,10 +246,18 @@ func (s *InMemoryRegisterStore) IsBlockExecuted(height uint64, blockID flow.Iden
 	defer s.RUnlock()
 
 	// finalized and executed blocks are pruned
-	if height <= s.prunedHeight {
+	// so if the height is below pruned height, in memory register store is not sure if
+	// it's executed or not
+	if height < s.prunedHeight {
 		return false, fmt.Errorf("below pruned height")
 	}
 
+	// if the height is the pruned height, then it's executed only if the blockID is the prunedID
+	// since the pruned block must be finalized and executed.
+	if height == s.prunedHeight {
+		return blockID == s.prunedID, nil
+	}
+
 	_, ok := s.registersByBlockID[blockID]
 	return ok, nil
 }
@@ -260,9 +268,18 @@ func (s *InMemoryRegisterStore) findFinalizedFork(height uint64, blockID flow.Id
 	s.RLock()
 	defer s.RUnlock()
 
-	if height <= s.prunedHeight {
+	if height < s.prunedHeight {
 		return nil, fmt.Errorf("cannot find finalized fork at height %d, it is pruned (prunedHeight: %v)", height, s.prunedHeight)
 	}
+
+	if height == s.prunedHeight {
+		if blockID != s.prunedID {
+			return nil, fmt.Errorf("cannot find finalized fork at height %d, it is pruned (prunedHeight: %v, prunedID: %v)", height, s.prunedHeight, s.prunedID)
+		}
+
+		return nil, nil
+	}
+
 	prunedHeight := height
 	block := blockID
 

engine/execution/storehouse/in_memory_register_store_test.go

@@ -183,6 +183,68 @@ func TestInMemoryRegisterStore(t *testing.T) {
 		require.Equal(t, regB.Value, valB)
 	})
 
+	t.Run("IsBlockExecuted", func(t *testing.T) {
+		t.Parallel()
+		pruned := uint64(10)
+		lastID := unittest.IdentifierFixture()
+		store := NewInMemoryRegisterStore(pruned, lastID)
+
+		height := pruned + 1 // above the pruned pruned
+		blockID := unittest.IdentifierFixture()
+		reg := unittest.RegisterEntryFixture()
+		err := store.SaveRegisters(
+			height,
+			blockID,
+			lastID,
+			flow.RegisterEntries{reg},
+		)
+		require.NoError(t, err)
+
+		// above the pruned height and is executed
+		executed, err := store.IsBlockExecuted(height, blockID)
+		require.NoError(t, err)
+		require.True(t, executed)
+
+		// above the pruned height, and is not executed
+		executed, err = store.IsBlockExecuted(pruned+1, unittest.IdentifierFixture())
+		require.NoError(t, err)
+		require.False(t, executed)
+
+		executed, err = store.IsBlockExecuted(pruned+2, unittest.IdentifierFixture())
+		require.NoError(t, err)
+		require.False(t, executed)
+
+		// below the pruned height
+		_, err = store.IsBlockExecuted(pruned-1, unittest.IdentifierFixture())
+		require.Error(t, err)
+
+		// equal to the pruned height and is the pruned block
+		executed, err = store.IsBlockExecuted(pruned, lastID)
+		require.NoError(t, err)
+		require.True(t, executed)
+
+		// equal to the pruned height, but is not the pruned block
+		executed, err = store.IsBlockExecuted(pruned, unittest.IdentifierFixture())
+		require.NoError(t, err)
+		require.False(t, executed)
+
+		// prune a new block
+		require.NoError(t, store.Prune(height, blockID))
+		// equal to the pruned height and is the pruned block
+		executed, err = store.IsBlockExecuted(height, blockID)
+		require.NoError(t, err)
+		require.True(t, executed)
+
+		// equal to the pruned height, but is not the pruned block
+		executed, err = store.IsBlockExecuted(height, unittest.IdentifierFixture())
+		require.NoError(t, err)
+		require.False(t, executed)
+
+		// below the pruned height
+		_, err = store.IsBlockExecuted(pruned, lastID)
+		require.Error(t, err)
+	})
+
 	// 5. Given A(X: 1, Y: 2), GetRegister(A, X) should return 1, GetRegister(A, X) should return 2
 	t.Run("GetRegistersOK", func(t *testing.T) {
 		t.Parallel()