{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}

-- | Allocation-free cross-era comparison of Cardano transaction ids.
--
-- Both 'Ouroboros.Consensus.HardFork.Combinator.AcrossEras.OneEraGenTxId'
-- instances order by the txid hash, ignoring the era. We read each id's 32-byte
-- hash (Blake2b-256) as four big-endian 'Word64#' and compare those in
-- registers. The words are unboxed, so no hash is boxed on the heap on any
-- era's path:
--
--   * Shelley-based eras store the hash as 'PackedBytes32' (four words
--     already); we unbox its fields.
--   * Byron stores it as a 'ShortByteString'; we read four big-endian words
--     out.
--
-- The extraction walks the era sum with a single dictionary used on every
-- branch, so it is strict in the dictionary and allocates no per-level thunk.
module Ouroboros.Consensus.Cardano.CanHardFork.OptimizedTxIdComparison
  ( ToTxIdWords (..)
  , compareCardanoGenTxId
  ) where

import Cardano.Crypto (abstractHashToShort)
import qualified Cardano.Crypto.Hash as Hash
import Cardano.Crypto.Hash.Class (PackedBytes (PackedBytes32))
import Cardano.Crypto.PackedBytes (unpackBytes)
import qualified Cardano.Ledger.Core as SL
import qualified Cardano.Ledger.Shelley.API as SL
import Data.Bits (unsafeShiftL, (.|.))
import Data.ByteString.Short (ShortByteString)
import qualified Data.ByteString.Short as SBS
import Data.SOP.Constraint (All)
import qualified Data.SOP.Strict as SOP
import GHC.Exts (Word64#, gtWord64#, ltWord64#)
import GHC.Word (Word64 (W64#))
import Ouroboros.Consensus.Byron.Ledger
import Ouroboros.Consensus.Ledger.SupportsMempool (GenTxId)
import Ouroboros.Consensus.Shelley.Ledger
import Ouroboros.Consensus.TypeFamilyWrappers (WrapGenTxId, unwrapGenTxId)

-- | Order two Cardano transaction ids by their txid hash, ignoring the era.
--
-- 'txIdWords' reads each hash as four unboxed words and 'compareW64' orders
-- them in registers. Two requirements keep this allocation-free:
--
-- * Both branches of 'txIdWords' force the @All ToTxIdWords ys@ dictionary
--   ('Z' reads its head, 'S' its tail), so GHC compiles the tail-dictionary
--   read as a strict field access rather than a per-step thunk.
--
-- * The compared words are unboxed ('Word64#'), so no hash is boxed on the
--   heap.
compareCardanoGenTxId ::
  All ToTxIdWords xs =>
  SOP.NS WrapGenTxId xs -> SOP.NS WrapGenTxId xs -> Ordering
compareCardanoGenTxId :: forall (xs :: [*]).
All ToTxIdWords xs =>
NS WrapGenTxId xs -> NS WrapGenTxId xs -> Ordering
compareCardanoGenTxId NS WrapGenTxId xs
l NS WrapGenTxId xs
r =
  case NS WrapGenTxId xs -> (# Word64#, Word64#, Word64#, Word64# #)
forall (ys :: [*]).
All ToTxIdWords ys =>
NS WrapGenTxId ys -> (# Word64#, Word64#, Word64#, Word64# #)
txIdWords NS WrapGenTxId xs
l of
    (# Word64#
a0, Word64#
a1, Word64#
a2, Word64#
a3 #) -> case NS WrapGenTxId xs -> (# Word64#, Word64#, Word64#, Word64# #)
forall (ys :: [*]).
All ToTxIdWords ys =>
NS WrapGenTxId ys -> (# Word64#, Word64#, Word64#, Word64# #)
txIdWords NS WrapGenTxId xs
r of
      (# Word64#
b0, Word64#
b1, Word64#
b2, Word64#
b3 #) ->
        Word64# -> Word64# -> Ordering
compareW64 Word64#
a0 Word64#
b0 Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Word64# -> Word64# -> Ordering
compareW64 Word64#
a1 Word64#
b1 Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Word64# -> Word64# -> Ordering
compareW64 Word64#
a2 Word64#
b2 Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Word64# -> Word64# -> Ordering
compareW64 Word64#
a3 Word64#
b3

-- | The four big-endian 64-bit words of an era's 32-byte txid hash
-- (Blake2b-256), unboxed so no hash is materialised on the heap. An era whose
-- txid hash is not four words is rejected by the 'Shelley' instance's
-- 'PackedBytes32' match rather than miscompared.
class ToTxIdWords blk where
  toTxIdWords :: GenTxId blk -> (# Word64#, Word64#, Word64#, Word64# #)

instance ToTxIdWords ByronBlock where
  toTxIdWords :: GenTxId ByronBlock -> (# Word64#, Word64#, Word64#, Word64# #)
toTxIdWords (ByronTxId TxId
i) = ShortByteString -> (# Word64#, Word64#, Word64#, Word64# #)
sbsWords (TxId -> ShortByteString
forall algo a. AbstractHash algo a -> ShortByteString
abstractHashToShort TxId
i)
  toTxIdWords (ByronDlgId CertificateId
i) = ShortByteString -> (# Word64#, Word64#, Word64#, Word64# #)
sbsWords (CertificateId -> ShortByteString
forall algo a. AbstractHash algo a -> ShortByteString
abstractHashToShort CertificateId
i)
  toTxIdWords (ByronUpdateProposalId UpId
i) = ShortByteString -> (# Word64#, Word64#, Word64#, Word64# #)
sbsWords (UpId -> ShortByteString
forall algo a. AbstractHash algo a -> ShortByteString
abstractHashToShort UpId
i)
  toTxIdWords (ByronUpdateVoteId VoteId
i) = ShortByteString -> (# Word64#, Word64#, Word64#, Word64# #)
sbsWords (VoteId -> ShortByteString
forall algo a. AbstractHash algo a -> ShortByteString
abstractHashToShort VoteId
i)

instance ShelleyBasedEra era => ToTxIdWords (ShelleyBlock proto era) where
  toTxIdWords :: GenTxId (ShelleyBlock proto era)
-> (# Word64#, Word64#, Word64#, Word64# #)
toTxIdWords (ShelleyTxId TxId
i) =
    case Hash HASH EraIndependentTxBody -> PackedBytes (HashSize HASH)
forall h a. Hash h a -> PackedBytes (HashSize h)
Hash.hashToPackedBytes (SafeHash EraIndependentTxBody -> Hash HASH EraIndependentTxBody
forall i. SafeHash i -> Hash HASH i
SL.extractHash (TxId -> SafeHash EraIndependentTxBody
SL.unTxId TxId
i)) of
      PackedBytes32 (W64# Word64#
w0) (W64# Word64#
w1) (W64# Word64#
w2) (W64# Word64#
w3) -> (# Word64#
w0, Word64#
w1, Word64#
w2, Word64#
w3 #)
      -- The ledger hash is Blake2b-256, always 'PackedBytes32'; this arm is
      -- unreachable and pays a copy only if it ever runs.
      PackedBytes (HashSize HASH)
pb -> ShortByteString -> (# Word64#, Word64#, Word64#, Word64# #)
sbsWords (PackedBytes 32 -> ShortByteString
forall (n :: Nat). PackedBytes n -> ShortByteString
unpackBytes PackedBytes 32
PackedBytes (HashSize HASH)
pb)

-- | The four big-endian 'Word64#' of a 32-byte 'ShortByteString', assembled by
-- shifting its bytes. Portable (no byte swap) and allocation-free: the words go
-- straight into an unboxed tuple. This is the byte order the raw-hash reference
-- compares by, so it agrees with the oracle (checked by the property test across
-- word boundaries).
sbsWords :: ShortByteString -> (# Word64#, Word64#, Word64#, Word64# #)
sbsWords :: ShortByteString -> (# Word64#, Word64#, Word64#, Word64# #)
sbsWords ShortByteString
sbs =
  (#
    Word64 -> Word64#
unbox (Int -> Word64
word64BigEndianAt Int
0)
    , Word64 -> Word64#
unbox (Int -> Word64
word64BigEndianAt Int
8)
    , Word64 -> Word64#
unbox (Int -> Word64
word64BigEndianAt Int
16)
    , Word64 -> Word64#
unbox (Int -> Word64
word64BigEndianAt Int
24)
  #)
 where
  -- Inline so the four calls don't share one heap-allocated closure.
  {-# INLINE word64BigEndianAt #-}
  word64BigEndianAt :: Int -> Word64
word64BigEndianAt Int
byteOffset =
    -- with bytes b0 b1 … b7 starting at 'byteOffset' the result looks like:
    --    (b0 << 56) | (b1 << 48) | (b2 << 40) | (b3 << 32)
    --  | (b4 << 24) | (b5 << 16) | (b6 << 8)  | b7
    (Int -> Word64
byte (Int
byteOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
56)
      Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Int -> Word64
byte (Int
byteOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
48)
      Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Int -> Word64
byte (Int
byteOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
40)
      Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Int -> Word64
byte (Int
byteOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
32)
      Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Int -> Word64
byte (Int
byteOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
24)
      Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Int -> Word64
byte (Int
byteOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
5) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
16)
      Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Int -> Word64
byte (Int
byteOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
6) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
8)
      Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Int -> Word64
byte (Int
byteOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7)
  byte :: Int -> Word64
byte Int
k = Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (HasCallStack => ShortByteString -> Int -> Word8
ShortByteString -> Int -> Word8
SBS.index ShortByteString
sbs Int
k) :: Word64
  unbox :: Word64 -> Word64#
unbox (W64# Word64#
w) = Word64#
w

-- | The txid hash of whichever era the id sits in, as four big-endian 'Word64#'.
-- Direct recursion down the sum, building no intermediate 'NS'. Uses its single
-- dictionary on every branch, so it is strict in it and allocates no per-level
-- thunk.
txIdWords ::
  All ToTxIdWords ys =>
  SOP.NS WrapGenTxId ys -> (# Word64#, Word64#, Word64#, Word64# #)
txIdWords :: forall (ys :: [*]).
All ToTxIdWords ys =>
NS WrapGenTxId ys -> (# Word64#, Word64#, Word64#, Word64# #)
txIdWords (SOP.Z WrapGenTxId x
x) = GenTxId x -> (# Word64#, Word64#, Word64#, Word64# #)
forall blk.
ToTxIdWords blk =>
GenTxId blk -> (# Word64#, Word64#, Word64#, Word64# #)
toTxIdWords (WrapGenTxId x -> GenTxId x
forall blk. WrapGenTxId blk -> GenTxId blk
unwrapGenTxId WrapGenTxId x
x)
txIdWords (SOP.S NS WrapGenTxId xs1
y) = NS WrapGenTxId xs1 -> (# Word64#, Word64#, Word64#, Word64# #)
forall (ys :: [*]).
All ToTxIdWords ys =>
NS WrapGenTxId ys -> (# Word64#, Word64#, Word64#, Word64# #)
txIdWords NS WrapGenTxId xs1
y

-- | Order two 64-bit words as unsigned, in registers.
compareW64 :: Word64# -> Word64# -> Ordering
compareW64 :: Word64# -> Word64# -> Ordering
compareW64 Word64#
a Word64#
b = case Word64#
a Word64# -> Word64# -> Int#
`ltWord64#` Word64#
b of
  Int#
1# -> Ordering
LT
  Int#
_ -> case Word64#
a Word64# -> Word64# -> Int#
`gtWord64#` Word64#
b of
    Int#
1# -> Ordering
GT
    Int#
_ -> Ordering
EQ