From 5b274c4ac7e6cf3cb1e1ce78ddfedb4ac17b3e27 Mon Sep 17 00:00:00 2001
From: Thomas Winant <thomas@well-typed.com>
Date: Fri, 17 Apr 2020 16:16:23 +0200
Subject: [PATCH] More compact encoding of the overlay schedule

Fixes #1374.
---
 .../src/Shelley/Spec/Ledger/LedgerState.hs    | 32 +++++++++++++++++--
 .../Spec/Ledger/Examples/Serialization.hs     |  4 ++-
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/shelley/chain-and-ledger/executable-spec/src/Shelley/Spec/Ledger/LedgerState.hs b/shelley/chain-and-ledger/executable-spec/src/Shelley/Spec/Ledger/LedgerState.hs
index b3c94f10ad0..cecf8cae99d 100644
--- a/shelley/chain-and-ledger/executable-spec/src/Shelley/Spec/Ledger/LedgerState.hs
+++ b/shelley/chain-and-ledger/executable-spec/src/Shelley/Spec/Ledger/LedgerState.hs
@@ -83,6 +83,8 @@ import           Cardano.Crypto.Hash (byteCount)
 import           Cardano.Prelude (NoUnexpectedThunks (..))
 import           Control.Monad.Trans.Reader (ReaderT (..), asks)
 import           Data.Foldable (toList)
+import           Data.List.NonEmpty (NonEmpty)
+import qualified Data.List.NonEmpty as NonEmpty
 import           Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
 import           Data.Proxy (Proxy (..))
@@ -395,7 +397,7 @@ instance Crypto crypto => FromCBOR (UTxOState crypto)
 data OBftSlot crypto =
     NonActiveSlot
   | ActiveSlot !(GenKeyHash crypto)
-  deriving (Show, Eq, Generic)
+  deriving (Show, Eq, Ord, Generic)
 
 instance Crypto crypto
   => ToCBOR (OBftSlot crypto)
@@ -433,7 +435,7 @@ instance Crypto crypto => ToCBOR (NewEpochState crypto)
  where
   toCBOR (NewEpochState e bp bc es ru pd os) =
     encodeListLen 7 <> toCBOR e <> toCBOR bp <> toCBOR bc <> toCBOR es
-      <> toCBOR ru <> toCBOR pd <> toCBOR os
+      <> toCBOR ru <> toCBOR pd <> toCBOR (compactOverlaySchedule os)
 
 instance Crypto crypto => FromCBOR (NewEpochState crypto)
  where
@@ -445,9 +447,33 @@ instance Crypto crypto => FromCBOR (NewEpochState crypto)
     es <- fromCBOR
     ru <- fromCBOR
     pd <- fromCBOR
-    os <- fromCBOR
+    os <- decompactOverlaySchedule <$> fromCBOR
     pure $ NewEpochState e bp bc es ru pd os
 
+-- | Convert the overlay schedule to a representation that is more compact
+-- when serialised to a bytestring, but less efficient for lookups.
+--
+-- Each genesis key hash will only be stored once, instead of each time it is
+-- assigned to a slot.
+compactOverlaySchedule
+  :: Map SlotNo (OBftSlot crypto)
+  -> Map (OBftSlot crypto) (NonEmpty SlotNo)
+compactOverlaySchedule =
+    Map.foldrWithKey'
+      (\slot obftSlot ->
+        Map.insertWith (<>) obftSlot (pure slot))
+      Map.empty
+
+-- | Inverse of 'compactOverlaySchedule'
+decompactOverlaySchedule
+  :: Map (OBftSlot crypto) (NonEmpty SlotNo)
+  -> Map SlotNo (OBftSlot crypto)
+decompactOverlaySchedule compact = Map.fromList
+    [ (slot, obftSlot)
+    | (obftSlot, slots) <- Map.toList compact
+    , slot <- NonEmpty.toList slots
+    ]
+
 getGKeys
   :: NewEpochState crypto
   -> Set (GenKeyHash crypto)
diff --git a/shelley/chain-and-ledger/executable-spec/test/Test/Shelley/Spec/Ledger/Examples/Serialization.hs b/shelley/chain-and-ledger/executable-spec/test/Test/Shelley/Spec/Ledger/Examples/Serialization.hs
index cd3dd8db2d5..65d12de6e72 100644
--- a/shelley/chain-and-ledger/executable-spec/test/Test/Shelley/Spec/Ledger/Examples/Serialization.hs
+++ b/shelley/chain-and-ledger/executable-spec/test/Test/Shelley/Spec/Ledger/Examples/Serialization.hs
@@ -7,6 +7,7 @@
 
 module Test.Shelley.Spec.Ledger.Examples.Serialization where
 
+import           Data.List.NonEmpty (NonEmpty ((:|)))
 import qualified Data.Maybe as Maybe (fromJust)
 import           Data.String (fromString)
 import qualified Shelley.Spec.Ledger.MetaData as MD
@@ -1109,6 +1110,7 @@ serializationTests = testGroup "Serialization Tests"
              }) :: StrictMaybe RewardUpdate
       pd = (PoolDistr Map.empty) :: PoolDistr
       os = Map.singleton (SlotNo 1) (ActiveSlot testGKeyHash)
+      compactOs = Map.singleton (ActiveSlot testGKeyHash) (SlotNo 1 :| [])
       nes = NewEpochState
               e
               (BlocksMade bs)
@@ -1127,6 +1129,6 @@ serializationTests = testGroup "Serialization Tests"
       <> S es
       <> S ru
       <> S pd
-      <> S os
+      <> S compactOs
     )
  ]