haskell-opentimestamps-0.5.5.0: src/OpenTimestamps/Prune.hs

{-# LANGUAGE OverloadedRecordDot #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

{-# HLINT ignore "Use tuple-section" #-}

{- | Timestamp pruning functionality for OpenTimestamps.

This module provides functions to prune timestamp trees by removing
suboptimal attestations and empty branches while preserving
the most valuable verification paths.
-}
module OpenTimestamps.Prune
  ( pruneTimestamp
  ) where

import qualified Data.ByteString as BS
import Data.List (foldl1')
import qualified Data.Map.Strict as Map
import Data.Maybe (catMaybes)
import qualified Data.Set as Set
import OpenTimestamps.Attestation
  ( Attestation (..)
  , isBitcoinAttestation
  )
import OpenTimestamps.Op (Op (Append, Prepend))
import OpenTimestamps.Timestamp (Timestamp (..))

{- | Compare Bitcoin attestations to select the best one.

Criteria: lowest height, then lowest depth.
-}
compareBitcoinAttestations ::
  (Attestation, Int, BS.ByteString) ->
  (Attestation, Int, BS.ByteString) ->
  (Attestation, Int, BS.ByteString)
compareBitcoinAttestations
  currentBest@(Bitcoin h1, depth1, _)
  (Bitcoin h2, depth2, msg2) =
    case compare h1 h2 of
      LT -> currentBest
      GT -> (Bitcoin h2, depth2, msg2)
      EQ ->
        if depth1 < depth2
          then currentBest
          else (Bitcoin h2, depth2, msg2)
compareBitcoinAttestations _ _ = error "Should only compare Bitcoin attestations"

-- | Calculates the length added by an operation for depth adjustment.
opLength :: Op -> Int
opLength (Append bs) = BS.length bs
opLength (Prepend bs) = BS.length bs
opLength _ = 0 -- Other ops don't have arguments that add to length

{- | Recursively prunes a sub-timestamp and adjusts the depth of its
best Bitcoin attestation.
-}
processSubTimestampAndAdjustDepth ::
  Op -> Timestamp -> (Timestamp, Maybe (Attestation, Int, BS.ByteString))
processSubTimestampAndAdjustDepth op subTs =
  let (prunedSubTs, mSubBestAttWithMsg) = pruneAttestations subTs
      mAdjustedSubBestAtt =
        fmap
          ( \(att, depth, msg) ->
              (att, depth + 1 + opLength op, msg)
          )
          mSubBestAttWithMsg
   in (prunedSubTs, mAdjustedSubBestAtt)

{- | Collects all Bitcoin attestations from the current timestamp and
its sub-timestamps.
-}
collectAllBitcoinAttestations ::
  Timestamp ->
  [Maybe (Attestation, Int, BS.ByteString)] ->
  Set.Set (Attestation, Int, BS.ByteString)
collectAllBitcoinAttestations ts subBestAttestations =
  let currentBitcoinAttestations =
        Set.map (\att -> (att, 0, ts.timestampMsg)) $
          Set.filter isBitcoinAttestation ts.attestations
   in currentBitcoinAttestations
        `Set.union` Set.fromList (catMaybes subBestAttestations)

-- | Selects the single best Bitcoin attestation from a set, if any exist.
selectBestBitcoinAttestation ::
  Set.Set (Attestation, Int, BS.ByteString) ->
  Maybe (Attestation, Int, BS.ByteString)
selectBestBitcoinAttestation allBitcoinAttestations =
  if Set.null allBitcoinAttestations
    then Nothing
    else
      Just $
        foldl1'
          compareBitcoinAttestations
          (Set.toList allBitcoinAttestations)

{- | Recursively prunes attestations, keeping only the optimal Bitcoin
attestation (lowest height) and all non-Bitcoin attestations. Returns
the pruned timestamp, the optimal Bitcoin attestation found in its
subtree, and its depth. The depth is calculated as the number of
operations from the root of the subtree to the attestation.
-}
pruneAttestations :: Timestamp -> (Timestamp, Maybe (Attestation, Int, BS.ByteString))
pruneAttestations ts = (newTs, mBestAttestationWithMsg)
  where
    -- Recursively prune sub-timestamps and collect their best attestations
    -- and adjusted depths.
    (prunedOpsList, subBestAttestations) =
      unzip $
        Map.elems $
          Map.mapWithKey processSubTimestampAndAdjustDepth ts.ops

    prunedOpsMap = Map.fromList $ zip (Map.keys ts.ops) prunedOpsList

    -- Collect all Bitcoin attestations from the current node and sub-nodes.
    allBitcoinAttestations = collectAllBitcoinAttestations ts subBestAttestations

    -- Determine the single best Bitcoin attestation for this subtree.
    mBestAttestationWithMsg = selectBestBitcoinAttestation allBitcoinAttestations

    -- All Bitcoin attestations at this node are removed.
    newAttestations = Set.filter (not . isBitcoinAttestation) ts.attestations

    newTs = ts {attestations = newAttestations, ops = prunedOpsMap}

-- | Prune empty branches from the timestamp tree. Returns (pruned timestamp, changed flag).
pruneTree :: Timestamp -> (Timestamp, Bool)
pruneTree ts = (ts {ops = newOps}, changed)
  where
    (newOps, changedOps) = Map.foldlWithKey' foldFn (Map.empty, False) ts.ops
    foldFn (accOps, accChanged) op subTs =
      let (prunedSubTs, subChanged) = pruneTree subTs
       in if Set.null prunedSubTs.attestations && Map.null prunedSubTs.ops
            then (accOps, accChanged || subChanged || True) -- branch is empty, remove it
            else (Map.insert op prunedSubTs accOps, accChanged || subChanged)
    changed =
      changedOps
        || ( Set.null ts.attestations
               && Map.null newOps
               && not (Set.null ts.attestations && Map.null ts.ops)
           )

{- | Reinserts the globally best Bitcoin attestation into the pruned
timestamp at the correct message.
-}
reinsertBestBitcoinAttestation :: Attestation -> BS.ByteString -> Timestamp -> Timestamp
reinsertBestBitcoinAttestation att targetMsg currentTs
  | timestampMsg currentTs == targetMsg =
      currentTs {attestations = Set.insert att currentTs.attestations}
  | otherwise =
      currentTs
        { ops =
            Map.map
              (reinsertBestBitcoinAttestation att targetMsg)
              currentTs.ops
        }

{- | Prune a timestamp by discarding suboptimal attestations and
removing empty branches.
-}
pruneTimestamp :: Timestamp -> Timestamp
pruneTimestamp initialTs =
  let (tsWithoutBitcoins, mGlobalBestAttestationWithMsg) = pruneAttestations initialTs
      tsWithBestAttReinserted = case mGlobalBestAttestationWithMsg of
        Nothing -> tsWithoutBitcoins
        Just (globalBestAtt, _, globalBestMsg) ->
          reinsertBestBitcoinAttestation globalBestAtt globalBestMsg tsWithoutBitcoins
   in fst (pruneTree tsWithBestAttReinserted)