Improve TxSeq decoders

eras/alonzo/impl/src/Cardano/Ledger/Alonzo/TxSeq/Internal.hs

@@ -46,15 +46,14 @@ import Cardano.Ledger.Binary (
   withSlice,
  )
 import Cardano.Ledger.Core
-import Cardano.Ledger.Shelley.BlockChain (constructMetadata, indexLookupSeq)
+import Cardano.Ledger.Shelley.BlockChain (auxDataSeqDecoder)
 import Control.Monad (unless)
 import Data.ByteString (ByteString)
 import Data.ByteString.Builder (shortByteString, toLazyByteString)
 import qualified Data.ByteString.Lazy as BSL
 import Data.Coerce (coerce)
-import qualified Data.Map.Strict as Map
+import Data.Functor.Identity (Identity (..))
 import Data.Maybe.Strict (maybeToStrictMaybe, strictMaybeToMaybe)
-import Data.Monoid (All (..))
 import Data.Proxy (Proxy (..))
 import qualified Data.Sequence as Seq
 import Data.Sequence.Strict (StrictSeq)
@@ -189,17 +188,10 @@ instance AlonzoEraTx era => DecCBOR (Annotator (AlonzoTxSeq era)) where
     let bodiesLength = length bodies
         inRange x = (0 <= x) && (x <= (bodiesLength - 1))
         witsLength = length wits
-    (auxData, auxDataAnn) <- withSlice $
-      do
-        auxDataMap <- decCBOR
-        unless
-          (getAll (Map.foldMapWithKey (\k _ -> All (inRange k)) auxDataMap))
-          ( fail
-              ( "Some Auxiliarydata index is not in the range: 0 .. "
-                  ++ show (bodiesLength - 1)
-              )
-          )
-        pure (constructMetadata bodiesLength auxDataMap)
+    (auxData, auxDataAnn) <- withSlice $ do
+      auxDataMap <- decCBOR
+      auxDataSeqDecoder bodiesLength auxDataMap False
+
     (isValIdxs, isValAnn) <- withSlice decCBOR
     let validFlags = alignedValidFlags bodiesLength isValIdxs
     unless
@@ -248,14 +240,9 @@ instance
     let bodiesLength = length bodies
         inRange x = (0 <= x) && (x <= (bodiesLength - 1))
         witsLength = length wits
-    unless
-      (getAll (Map.foldMapWithKey (\k _ -> All (inRange k)) auxDataMap))
-      ( fail
-          ( "Some Auxiliarydata index is not in the range: 0 .. "
-              ++ show (bodiesLength - 1)
-          )
-      )
-    let auxData = indexLookupSeq bodiesLength auxDataMap
+    auxData <-
+      fmap (fmap runIdentity)
+        <$> auxDataSeqDecoder bodiesLength (fmap pure auxDataMap) False
     Annotated isValidIdxs isValidBytes <- decodeAnnotated decCBOR
     let validFlags = alignedValidFlags bodiesLength isValidIdxs
     unless

eras/shelley/impl/CHANGELOG.md

@@ -2,13 +2,14 @@
 
 ## 1.16.0.0
 
+* Add `auxDataSeqDecoder`
+* Remove `constructMetadata`
 * Add `DecCBOR` instances for:
   * `ShelleyTxWits`
   * `ShelleyTxAuxData`
   * `ShelleyTxBody`
   * `ShelleyTx`
   * `ShelleyTxSeq`
-* Add `indexLookupSeq`
 * Add `segWitTx`
 * Rename `segwitTx` to `segWitAnnTx`
 * Converted `CertState` to a type family

eras/shelley/impl/src/Cardano/Ledger/Shelley/BlockChain.hs

@@ -18,8 +18,7 @@
 
 module Cardano.Ledger.Shelley.BlockChain (
   ShelleyTxSeq (ShelleyTxSeq, txSeqTxns', TxSeq'),
-  constructMetadata,
-  indexLookupSeq,
+  auxDataSeqDecoder,
   txSeqTxns,
   bbHash,
   bBodySize,
@@ -61,7 +60,9 @@ import Control.Monad (unless)
 import Data.ByteString (ByteString)
 import qualified Data.ByteString.Lazy as BSL
 import Data.Coerce (coerce)
-import Data.Map.Strict (Map)
+import Data.Functor.Identity (Identity (..))
+import Data.IntMap (IntMap)
+import qualified Data.IntMap as IntMap
 import qualified Data.Map.Strict as Map
 import Data.Monoid (All (..))
 import Data.Sequence (Seq)
@@ -191,18 +192,6 @@ bbHash (TxSeq' _ bodies wits md) =
     hashStrict = Hash.hashWith id
     hashPart = Hash.hashToBytes . hashStrict . BSL.toStrict
 
--- | Given a size and a mapping from indices to maybe metadata,
---  return a sequence whose size is the size parameter and
---  whose non-Nothing values correspond to the values in the mapping.
-constructMetadata ::
-  Int ->
-  Map Int (Annotator (TxAuxData era)) ->
-  Seq (Maybe (Annotator (TxAuxData era)))
-constructMetadata = indexLookupSeq
-
-indexLookupSeq :: Int -> Map Int a -> Seq (Maybe a)
-indexLookupSeq n md = fmap (`Map.lookup` md) (Seq.fromList [0 .. n - 1])
-
 -- | The parts of the Tx in Blocks that have to have DecCBOR(Annotator x) instances.
 --   These are exactly the parts that are SafeToHash.
 -- | Decode a TxSeq, used in decoding a Block.
@@ -215,24 +204,19 @@ txSeqDecoder ::
 txSeqDecoder lax = do
   (bodies, bodiesAnn) <- withSlice decCBOR
   (wits, witsAnn) <- withSlice decCBOR
-  let b = length bodies
-      inRange x = (0 <= x) && (x <= (b - 1))
-      w = length wits
-  (metadata, metadataAnn) <- withSlice $
-    do
-      m <- decCBOR
-      unless -- TODO this PR introduces this new test, That didn't used to run in the Shelley
-        (lax || all inRange (Map.keysSet m)) -- Era,  Is it possible there might be some blocks, that should have been caught on the chain?
-        (fail ("Some Auxiliarydata index is not in the range: 0 .. " ++ show (b - 1)))
-      pure (constructMetadata @era b m)
+  let bodiesLength = length bodies
+      witsLength = length wits
+  (metadata, metadataAnn) <- withSlice $ do
+    auxDataMap <- decCBOR
+    auxDataSeqDecoder bodiesLength auxDataMap lax
 
   unless
-    (lax || b == w)
+    (lax || bodiesLength == witsLength)
     ( fail $
         "different number of transaction bodies ("
-          <> show b
+          <> show bodiesLength
           <> ") and witness sets ("
-          <> show w
+          <> show witsLength
           <> ")"
     )
 
@@ -242,6 +226,21 @@ txSeqDecoder lax = do
             Seq.zipWith3 segWitAnnTx bodies wits metadata
   pure $ TxSeq' <$> txns <*> bodiesAnn <*> witsAnn <*> metadataAnn
 
+auxDataSeqDecoder ::
+  Int -> IntMap (m (TxAuxData era)) -> Bool -> Decoder s (Seq (Maybe (m (TxAuxData era))))
+auxDataSeqDecoder bodiesLength auxDataMap lax = do
+  unless
+    (lax || getAll (IntMap.foldMapWithKey (\k _ -> All (inRange k)) auxDataMap))
+    (fail ("Some Auxiliarydata index is not in the range: 0 .. " ++ show (bodiesLength - 1)))
+  pure (indexLookupSeq bodiesLength auxDataMap)
+  where
+    inRange x = (0 <= x) && (x <= (bodiesLength - 1))
+    -- Given a size and a mapping from indices to maybe values,
+    -- return a sequence whose size is the size parameter and
+    -- whose non-Nothing values correspond to the values in the mapping.
+    indexLookupSeq :: Int -> IntMap a -> Seq (Maybe a)
+    indexLookupSeq n ixMap = Seq.fromList [IntMap.lookup ix ixMap | ix <- [0 .. n - 1]]
+
 instance EraTx era => DecCBOR (Annotator (ShelleyTxSeq era)) where
   decCBOR = txSeqDecoder False
 
@@ -256,13 +255,12 @@ instance
   decCBOR = do
     Annotated bodies bodiesBytes <- decodeAnnotated decCBOR
     Annotated wits witsBytes <- decodeAnnotated decCBOR
-    Annotated auxDataMap auxDataBytes <- decodeAnnotated decCBOR
+    Annotated (auxDataMap :: IntMap (TxAuxData era)) auxDataBytes <- decodeAnnotated decCBOR
     let bodiesLength = length bodies
-    let inRange x = (0 <= x) && (x <= (bodiesLength - 1))
-    unless
-      (getAll (Map.foldMapWithKey (\k _ -> All (inRange k)) auxDataMap))
-      (fail ("Some Auxiliarydata index is not in the range: 0 .. " ++ show (bodiesLength - 1)))
-    let auxData = indexLookupSeq bodiesLength auxDataMap
+    auxData <-
+      fmap (fmap runIdentity)
+        <$> auxDataSeqDecoder bodiesLength (fmap pure auxDataMap) False
+
     let witsLength = length wits
     unless
       (bodiesLength == witsLength)

libs/cardano-ledger-binary/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## 1.6.0.0
 
+* Add `DecCBOR` instance for `Data.IntMap`
+* Add `decodeIntMap`
 * Add `ToCBOR` instance for `PV1.Data`
 * Add `DecCBOR` instance for `Annotated a ByteString`
 * Add `originalBytesExpectedFailureMessage` needed for testing

libs/cardano-ledger-binary/src/Cardano/Ledger/Binary/Decoding/DecCBOR.hs

@@ -66,6 +66,7 @@ import Data.ByteString.Short.Internal (ShortByteString(SBS))
 import Data.Fixed (Fixed (..))
 import Data.IP (IPv4, IPv6)
 import Data.Int (Int16, Int32, Int64, Int8)
+import qualified Data.IntMap as IntMap
 import Data.List.NonEmpty (NonEmpty)
 import qualified Data.Map.Strict as Map
 import qualified Data.Maybe.Strict as SMaybe
@@ -406,6 +407,10 @@ instance (Ord k, DecCBOR k, DecCBOR v) => DecCBOR (Map.Map k v) where
   decCBOR = decodeMap decCBOR decCBOR
   {-# INLINE decCBOR #-}
 
+instance DecCBOR v => DecCBOR (IntMap.IntMap v) where
+  decCBOR = decodeIntMap decCBOR
+  {-# INLINE decCBOR #-}
+
 instance
   ( Ord k
   , DecCBOR k

libs/cardano-ledger-binary/src/Cardano/Ledger/Binary/Decoding/Decoder.hs

@@ -67,6 +67,7 @@ module Cardano.Ledger.Binary.Decoding.Decoder (
   decodeVector,
   decodeSet,
   setTag,
+  decodeIntMap,
   decodeMap,
   decodeMapByKey,
   decodeMapLikeEnforceNoDuplicates,
@@ -259,6 +260,7 @@ import qualified Data.ByteString.Lazy as BSL
 import Data.Functor.Compose (Compose (..))
 import Data.IP (IPv4, IPv6, fromHostAddress, fromHostAddress6)
 import Data.Int (Int16, Int32, Int64, Int8)
+import qualified Data.IntMap as IntMap
 import qualified Data.List.NonEmpty as NE (NonEmpty, nonEmpty)
 import qualified Data.Map.Strict as Map
 import Data.Maybe.Strict (StrictMaybe (..), maybeToStrictMaybe)
@@ -794,15 +796,46 @@ decodeMapLikeEnforceNoDuplicates ::
   Decoder s (Maybe Int) ->
   Decoder s (k, v) ->
   Decoder s (Map.Map k v)
-decodeMapLikeEnforceNoDuplicates decodeLenOrIndef =
+decodeMapLikeEnforceNoDuplicates =
+  decodeMapLikeEnforceNoDuplicatesInternal Map.fromList Map.size
+{-# INLINE decodeMapLikeEnforceNoDuplicates #-}
+
+decodeMapLikeEnforceNoDuplicatesInternal ::
+  ([(k, v)] -> m) ->
+  (m -> Int) ->
+  Decoder s (Maybe Int) ->
+  Decoder s (k, v) ->
+  Decoder s m
+decodeMapLikeEnforceNoDuplicatesInternal fromPairs size decodeLenOrIndef =
   -- We first decode into a list because most of the time the encoded Map will be in sorted
   -- order and there is a nice optimization on the `Map.fromList` that can take advantage of
   -- that fact. In case when encoded data is not sorted the penalty of going through a list
   -- is insignificant.
   decodeListLikeEnforceNoDuplicates decodeLenOrIndef (:) $ \xs ->
-    let result = Map.fromList (reverse xs)
-     in (Map.size result, result)
-{-# INLINE decodeMapLikeEnforceNoDuplicates #-}
+    let result = fromPairs (reverse xs)
+     in (size result, result)
+
+decodeIntMap :: Decoder s v -> Decoder s (IntMap.IntMap v)
+decodeIntMap decodeValue =
+  ifDecoderVersionAtLeast
+    (natVersion @2)
+    ( ifDecoderVersionAtLeast
+        (natVersion @9)
+        ( decodeMapLikeEnforceNoDuplicatesInternal
+            IntMap.fromList
+            IntMap.size
+            decodeMapLenOrIndef
+            decodeKeyValue
+        )
+        (IntMap.fromList <$> decodeCollection decodeMapLenOrIndef decodeKeyValue)
+    )
+    ( decodeMapSkel
+        (IntMap.fromDistinctAscList . reverse)
+        decodeKeyValue
+    )
+  where
+    decodeKeyValue = (,) <$> decodeInt <*> decodeValue
+{-# INLINE decodeIntMap #-}
 
 -- | Decode `VMap`. Unlike `decodeMap` it does not behavee differently for
 -- version prior to 2.