Make computeLoEMaxExtra return also a list of blocks

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

@@ -176,7 +176,7 @@ initialChainSelection immutableDB volatileDB lgrDB tracer cfg varInvalid
         suffixesAfterI :: [NonEmpty (HeaderHash blk)]
         suffixesAfterI = Paths.maximalCandidates succsOf limit (AF.anchorToPoint i)
           where
-            limit = k <$ loE
+            limit = ([], k) <$ loE
 
         constructChain ::
              NonEmpty (HeaderHash blk)
@@ -648,7 +648,7 @@ chainSelectionForBlock cdb@CDB{..} blockCache hdr punish = electric $ do
          -- ^ The current chain and ledger
       -> LoE (AnchoredFragment (Header blk))
          -- ^ LoE fragment
-      -> LoE Word64
+      -> LoE ([Header blk], Word64)
          -- ^ How many extra blocks to select after @b@ at most.
       -> m (Point blk)
     addToCurrentChain succsOf curChainAndLedger loeFrag maxExtra = do
@@ -738,7 +738,7 @@ chainSelectionForBlock cdb@CDB{..} blockCache hdr punish = electric $ do
          -- ^ The current chain (anchored at @i@) and ledger
       -> LoE (AnchoredFragment (Header blk))
          -- ^ LoE fragment
-      -> LoE Word64
+      -> LoE ([Header blk], Word64)
          -- ^ How many extra blocks to select after @b@ at most.
       -> ChainDiff (HeaderFields blk)
          -- ^ Header fields for @(x,b]@
@@ -793,16 +793,10 @@ chainSelectionForBlock cdb@CDB{..} blockCache hdr punish = electric $ do
     -- | How many extra blocks to select at most after the tip of @newBlockFrag@
     -- according to the LoE.
     --
-    -- There are two cases to consider:
-    --
-    -- 1. If @newBlockFrag@ and @loeFrag@ are on the same chain, then we cannot
-    --    select more than @loeLimit@ blocks after @loeFrag@.
-    --
-    -- 2. If @newBlockFrag@ and @loeFrag@ are on different chains, then we
-    --   cannot select more than @loeLimit@ blocks after their intersection.
-    --
-    -- In any case, 'Nothing' is returned if @newBlockFrag@ extends beyond
-    -- what LoE allows.
+    -- The result is the pair of a list of blocks (coming from the LoE) that are
+    -- the only one that may be selected and a number of additional blocks after
+    -- that, or 'Nothing' if @newBlockFrag@ forks of a strict prefix of the LoE,
+    -- or if it extends the LoE beyond @k@.
     computeLoEMaxExtra ::
          (HasHeader x, HeaderHash x ~ HeaderHash blk)
       => LoE (AnchoredFragment (Header blk))
@@ -811,26 +805,32 @@ chainSelectionForBlock cdb@CDB{..} blockCache hdr punish = electric $ do
       -> AnchoredFragment x
          -- ^ The fragment with the new block @b@ as its tip, with the same
          -- anchor as @curChain@.
-      -> Maybe (LoE Word64)
+      -> Maybe (LoE ([Header blk], Word64))
+    computeLoEMaxExtra LoEDisabled _ = Just LoEDisabled
     computeLoEMaxExtra (LoEEnabled loeFrag) newBlockFrag =
-        -- Both fragments are on the same chain
-        if loeSuffixLength == 0 || rollback == 0 then
-          if rollback <= k + loeSuffixLength
-            then Just $ LoEEnabled $ k + loeSuffixLength - rollback
-            else Nothing
-        else
-          -- REVIEW: Since we filter out candidates that fork off the LoE later,
-          -- it does not really matter what we return here and we could just
-          -- return @Nothing@ or merge it in the previous case.
-          if rollback <= k
-            then Just $ LoEEnabled $ k - rollback
-            else Nothing
+      if
+        -- The LoE suffix is empty, which means that the candidate is either
+        -- exactly the LoE fragment or extends it. We can select @k@ blocks from
+        -- the LoE tip.
+        | AF.null loeSuffix
+        , rollback <= k
+        , extra <- k - rollback
+        -> Just $ LoEEnabled ([], extra)
+
+        -- The rollback is 0 and the LoE suffix is not empty. We can select the
+        -- blocks of the LoE (and only those) and then @k@ more.
+        | rollback == 0
+        -> Just $ LoEEnabled (AF.toOldestFirst loeSuffix, k)
+
+        -- The LoE suffix is not empty and the rollback is not 0, which means
+        -- that the candidate forks off the LoE fragment at some point, but not
+        -- at the tip, which is forbidden.
+        | otherwise
+        -> Nothing
       where
         d = Diff.diff newBlockFrag loeFrag
         rollback = Diff.getRollback d
-        loeSuffixLength = fromIntegral $ AF.length (Diff.getSuffix d)
-    computeLoEMaxExtra LoEDisabled _ =
-        Just LoEDisabled
+        loeSuffix = Diff.getSuffix d
 
     mkSelectionChangedInfo ::
          AnchoredFragment (Header blk) -- ^ old chain

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

@@ -70,7 +70,7 @@ maximalCandidates ::
      forall blk.
      (ChainHash blk -> Set (HeaderHash blk))
      -- ^ @filterByPredecessor@
-  -> LoE Word64 -- ^ Max length of any candidate
+  -> LoE ([Header blk], Word64) -- ^ Max length of any candidate
   -> Point blk -- ^ @B@
   -> [NonEmpty (HeaderHash blk)]
      -- ^ Each element in the list is a list of hashes from which we can
@@ -85,8 +85,8 @@ maximalCandidates succsOf loeLimit b = mapMaybe (NE.nonEmpty . applyLoE) $ go (p
                , candidate <- go (BlockHash next)
                ]
     applyLoE
-      | LoEEnabled limit <- loeLimit
-      = take (fromIntegral limit)
+      | LoEEnabled (blocks, limit) <- loeLimit
+      = take (length blocks + fromIntegral limit)
       | otherwise
       = id
 
@@ -104,7 +104,7 @@ extendWithSuccessors ::
      forall blk. HasHeader blk
   => (ChainHash blk -> Set (HeaderHash blk))
   -> LookupBlockInfo blk
-  -> LoE Word64 -- ^ Max extra length for any suffix
+  -> LoE ([Header blk], Word64) -- ^ Max extra length for any suffix
   -> ChainDiff (HeaderFields blk)
   -> NonEmpty (ChainDiff (HeaderFields blk))
 extendWithSuccessors succsOf lookupBlockInfo loeLimit diff =