Also follow the loeFrag when switching to a fork

fixup to 8b3f49a

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Storage/ChainDB/Impl/ChainSel.hs

@@ -595,7 +595,7 @@ chainSelectionForBlock cdb@CDB{..} blockCache hdr punish = electric $ do
       , Just maxExtra <- computeLoEMaxExtra loeFrag newBlockFrag -> do
         -- ### Switch to a fork
         traceWith addBlockTracer (TrySwitchToAFork p diff)
-        switchToAFork succsOf' lookupBlockInfo' curChainAndLedger maxExtra diff
+        switchToAFork succsOf' lookupBlockInfo' curChainAndLedger loeFrag maxExtra diff
 
       -- We cannot reach the block from the current selection
       | otherwise -> do
@@ -707,16 +707,16 @@ chainSelectionForBlock cdb@CDB{..} blockCache hdr punish = electric $ do
         curTip      = castPoint $ AF.headPoint curChain
         curHead     = AF.headAnchor curChain
 
-        -- Either frag extends loe or loe extends frag
-        followsLoEFrag :: LoE (AnchoredFragment (Header blk))
-                      -> AnchoredFragment (Header blk)
-                      -> Bool
-        followsLoEFrag LoEDisabled _ = True
-        followsLoEFrag (LoEEnabled loe) frag =
-          case AF.intersect frag loe of
-            Just (_, _, AF.Empty{}, _) -> True
-            Just (_, _, _, AF.Empty{}) -> True
-            _                          -> False
+    -- Either frag extends loe or loe extends frag
+    followsLoEFrag :: LoE (AnchoredFragment (Header blk))
+                   -> AnchoredFragment (Header blk)
+                   -> Bool
+    followsLoEFrag LoEDisabled _ = True
+    followsLoEFrag (LoEEnabled loe) frag =
+      case AF.intersect frag loe of
+        Just (_, _, AF.Empty{}, _) -> True
+        Just (_, _, _, AF.Empty{}) -> True
+        _                          -> False
 
     -- | We have found a 'ChainDiff' through the VolatileDB connecting the new
     -- block to the current chain. We'll call the intersection/anchor @x@.
@@ -730,22 +730,26 @@ chainSelectionForBlock cdb@CDB{..} blockCache hdr punish = electric $ do
       -> LookupBlockInfo blk
       -> ChainAndLedger blk
          -- ^ The current chain (anchored at @i@) and ledger
+      -> LoE (AnchoredFragment (Header blk))
+         -- ^ LoE fragment
       -> LoELimit
          -- ^ How many extra blocks to select after @b@ at most.
       -> ChainDiff (HeaderFields blk)
          -- ^ Header fields for @(x,b]@
       -> m (Point blk)
-    switchToAFork succsOf lookupBlockInfo curChainAndLedger maxExtra diff = do
+    switchToAFork succsOf lookupBlockInfo curChainAndLedger loeFrag maxExtra diff = do
         -- We use a cache to avoid reading the headers from disk multiple
         -- times in case they're part of multiple forks that go through @b@.
         let initCache = Map.singleton (headerHash hdr) hdr
         chainDiffs <-
+            fmap (filter (followsLoEFrag loeFrag . Diff.getSuffix))
+
           -- 4. Filter out candidates that are not preferred over the current
           -- chain.
           --
           -- The suffixes all fork off from the current chain within @k@
           -- blocks, so it satisfies the precondition of 'preferCandidate'.
-            fmap
+          . fmap
               ( filter
                   ( preferAnchoredCandidate (bcfg chainSelEnv) curChain
                   . Diff.getSuffix