More precise type for Index (#1435)

jasagredo · web-flow · commit 9129bb056101 · 2025-03-31T14:32:00.000Z
Taken from [here](https://github.com/IntersectMBO/ouroboros-consensus/pull/1399/files#r1961871210) we also make the payload strict by using the strict `NS` (#1331).
diff --git a/ouroboros-consensus-cardano/changelog.d/js-index.md b/ouroboros-consensus-cardano/changelog.d/js-index.md
@@ -0,0 +1 @@
+<!-- empty -->
diff --git a/ouroboros-consensus-cardano/src/ouroboros-consensus-cardano/Ouroboros/Consensus/Cardano/Node.hs b/ouroboros-consensus-cardano/src/ouroboros-consensus-cardano/Ouroboros/Consensus/Cardano/Node.hs
@@ -69,7 +69,7 @@ import           Data.Functor.These (These1 (..))
 import qualified Data.Map.Strict as Map
 import           Data.SOP.BasicFunctors
 import           Data.SOP.Counting
-import           Data.SOP.Index (Index (..))
+import           Data.SOP.Index
 import           Data.SOP.OptNP (NonEmptyOptNP, OptNP (OptSkip))
 import qualified Data.SOP.OptNP as OptNP
 import           Data.SOP.Strict
diff --git a/ouroboros-consensus/changelog.d/js-index.md b/ouroboros-consensus/changelog.d/js-index.md
@@ -0,0 +1,4 @@
+### Non-Breaking
+
+- Define some functions in terms of SOP combinators leveraging the recent
+  simplification of `Index` in `sop-extras` package.
diff --git a/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Abstract/SingleEraBlock.hs b/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Abstract/SingleEraBlock.hs
@@ -159,7 +159,7 @@ eraIndexEmpty (EraIndex ns) = case ns of {}
 eraIndexFromNS :: SListI xs => NS f xs -> EraIndex xs
 eraIndexFromNS = EraIndex . hmap (const (K ()))
 
-eraIndexFromIndex :: Index xs blk -> EraIndex xs
+eraIndexFromIndex :: SListI xs => Index xs blk -> EraIndex xs
 eraIndexFromIndex index = EraIndex $ injectNS index (K ())
 
 eraIndexZero :: EraIndex (x ': xs)
diff --git a/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Block.hs b/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Block.hs
@@ -205,7 +205,7 @@ distribAnnTip AnnTip{..} =
     distrib (WrapTipInfo info) =
         AnnTip annTipSlotNo annTipBlockNo info
 
-undistribAnnTip :: SListI xs => NS AnnTip xs -> AnnTip (HardForkBlock xs)
+undistribAnnTip :: forall xs. SListI xs => NS AnnTip xs -> AnnTip (HardForkBlock xs)
 undistribAnnTip = hcollapse . himap undistrib
   where
     undistrib :: Index xs blk
diff --git a/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Embed/Nary.hs b/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Embed/Nary.hs
@@ -25,11 +25,13 @@ module Ouroboros.Consensus.HardFork.Combinator.Embed.Nary (
 import           Data.Bifunctor (first)
 import           Data.Coerce (Coercible, coerce)
 import           Data.SOP.BasicFunctors
+import           Data.SOP.Constraint
 import           Data.SOP.Counting (Exactly (..))
 import           Data.SOP.Dict (Dict (..))
 import           Data.SOP.Index
 import qualified Data.SOP.InPairs as InPairs
 import           Data.SOP.Strict
+import qualified Data.SOP.Telescope as Telescope
 import           Ouroboros.Consensus.Block
 import           Ouroboros.Consensus.Config
 import           Ouroboros.Consensus.HardFork.Combinator
@@ -75,11 +77,9 @@ newtype InjectionIndex xs x =
 
 -- | Many instances of 'Inject' do not need the 'History.Bound's, eg those that
 -- do not construct 'HardForkState's
-forgetInjectionIndex :: InjectionIndex xs x -> Index xs x
-forgetInjectionIndex (InjectionIndex tele) = case tele of
-    TZ (State.Current _k Refl) -> IZ
-    TS _k tele'                ->
-        IS (forgetInjectionIndex (InjectionIndex tele'))
+forgetInjectionIndex :: SListI xs => InjectionIndex xs x -> Index xs x
+forgetInjectionIndex (InjectionIndex tele) =
+  Index $ hmap State.currentState $ Telescope.tip tele
 
 -- | Build an 'InjectionIndex' from oracular 'History.Bound's and an 'Index'
 --
@@ -89,18 +89,17 @@ forgetInjectionIndex (InjectionIndex tele) = case tele of
 -- INVARIANT: the result is completely independent of the 'history.Bound's for
 -- eras /after/ the given 'Index'.
 oracularInjectionIndex ::
-     Exactly xs History.Bound
+     SListI xs
+  => Exactly xs History.Bound
   -> Index xs x
   -> InjectionIndex xs x
-oracularInjectionIndex (Exactly np) idx = case (idx, np) of
-    (IZ     , K start :* _      ) ->
-        InjectionIndex
-      $ TZ State.Current { currentStart = start, currentState = Refl }
-    (IS idx', kstart  :* kstarts) ->
-        let InjectionIndex iidx =
-              oracularInjectionIndex (Exactly kstarts) idx'
-        in
-        InjectionIndex (TS kstart iidx)
+oracularInjectionIndex (Exactly np) (Index idx) =
+  InjectionIndex
+    $ Telescope.bihzipWith
+      (\x (K ()) -> x)
+      (\(K start) Refl -> State.Current { currentStart = start, currentState = Refl })
+      np
+    $ Telescope.fromTip idx
 
 -- | NOT EXPORTED
 --
diff --git a/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Ledger.hs b/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Ledger.hs
@@ -157,7 +157,7 @@ instance CanHardFork xs => IsLedger (LedgerState (HardForkBlock xs)) where
       extended :: HardForkState LedgerState xs
       extended = State.extendToSlot cfg slot st
 
-tickOne :: SingleEraBlock blk
+tickOne :: (SListI xs, SingleEraBlock blk)
         => EpochInfo (Except PastHorizonException)
         -> SlotNo
         -> ComputeLedgerEvents
@@ -207,7 +207,7 @@ instance CanHardFork xs
       error "reapplyBlockLedgerResult: can't be from other era"
       )
 
-apply :: SingleEraBlock blk
+apply :: (SListI xs, SingleEraBlock blk)
       => STS.ValidationPolicy
       -> ComputeLedgerEvents
       -> Index xs                                           blk
@@ -728,14 +728,14 @@ ledgerViewInfo :: forall blk f. SingleEraBlock blk
                => f blk -> LedgerEraInfo blk
 ledgerViewInfo _ = LedgerEraInfo $ singleEraInfo (Proxy @blk)
 
-injectLedgerError :: Index xs blk -> LedgerError blk -> HardForkLedgerError xs
+injectLedgerError :: SListI xs => Index xs blk -> LedgerError blk -> HardForkLedgerError xs
 injectLedgerError index =
       HardForkLedgerErrorFromEra
     . OneEraLedgerError
     . injectNS index
     . WrapLedgerErr
 
-injectLedgerEvent :: Index xs blk -> AuxLedgerEvent (LedgerState blk) -> OneEraLedgerEvent xs
+injectLedgerEvent :: SListI xs => Index xs blk -> AuxLedgerEvent (LedgerState blk) -> OneEraLedgerEvent xs
 injectLedgerEvent index =
       OneEraLedgerEvent
     . injectNS index
diff --git a/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Mempool.hs b/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Mempool.hs
@@ -354,14 +354,14 @@ ledgerInfo :: forall blk. SingleEraBlock blk
            => State.Current (Ticked :.: LedgerState) blk -> LedgerEraInfo blk
 ledgerInfo _ = LedgerEraInfo $ singleEraInfo (Proxy @blk)
 
-injectApplyTxErr :: Index xs blk -> ApplyTxErr blk -> HardForkApplyTxErr xs
+injectApplyTxErr :: SListI xs => Index xs blk -> ApplyTxErr blk -> HardForkApplyTxErr xs
 injectApplyTxErr index =
       HardForkApplyTxErrFromEra
     . OneEraApplyTxErr
     . injectNS index
     . WrapApplyTxErr
 
-injectValidatedGenTx :: Index xs blk -> Validated (GenTx blk) -> Validated (GenTx (HardForkBlock xs))
+injectValidatedGenTx :: SListI xs => Index xs blk -> Validated (GenTx blk) -> Validated (GenTx (HardForkBlock xs))
 injectValidatedGenTx index =
       HardForkValidatedGenTx
     . OneEraValidatedGenTx
diff --git a/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Protocol.hs b/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Protocol.hs
@@ -31,6 +31,7 @@ module Ouroboros.Consensus.HardFork.Combinator.Protocol (
 import           Control.Monad.Except
 import           Data.Functor.Product
 import           Data.SOP.BasicFunctors
+import           Data.SOP.Constraint
 import           Data.SOP.Index
 import           Data.SOP.InPairs (InPairs (..))
 import qualified Data.SOP.InPairs as InPairs
@@ -304,7 +305,7 @@ update HardForkConsensusConfig{..}
   where
     cfgs = getPerEraConsensusConfig hardForkConsensusConfigPerEra
 
-updateEra :: forall xs blk. SingleEraBlock blk
+updateEra :: forall xs blk. (SListI xs, SingleEraBlock blk)
           => EpochInfo (Except PastHorizonException)
           -> SlotNo
           -> Index xs blk
@@ -378,7 +379,8 @@ translateConsensus ei HardForkConsensusConfig{..} =
     pcfgs = getPerEraConsensusConfig hardForkConsensusConfigPerEra
     cfgs  = hcmap proxySingle (completeConsensusConfig'' ei) pcfgs
 
-injectValidationErr :: Index xs blk
+injectValidationErr :: SListI xs
+                    => Index xs blk
                     -> ValidationErr (BlockProtocol blk)
                     -> HardForkValidationErr xs
 injectValidationErr index =
diff --git a/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Serialisation/Common.hs b/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Serialisation/Common.hs
@@ -446,7 +446,7 @@ encodeNS es ns = mconcat [
     , hcollapse $ hzipWith apFn es ns
     ]
 
-decodeNS :: SListI xs => NP (Decoder s :.: f) xs -> Decoder s (NS f xs)
+decodeNS :: forall xs f s. SListI xs => NP (Decoder s :.: f) xs -> Decoder s (NS f xs)
 decodeNS ds = do
     enforceSize "decodeNS" 2
     i <- Dec.decodeWord8
@@ -460,7 +460,7 @@ decodeNS ds = do
         -> K (Decoder s (NS f xs)) blk
     aux index (Comp dec) (K ()) = K $ injectNS index <$> dec
 
-decodeAnnNS :: SListI xs
+decodeAnnNS :: forall xs f. SListI xs
             => NP (AnnDecoder f) xs
             -> forall s. Decoder s (Lazy.ByteString -> NS f xs)
 decodeAnnNS ds = do
diff --git a/sop-extras/changelog.d/js-index.md b/sop-extras/changelog.d/js-index.md
@@ -0,0 +1,4 @@
+### Non-Breaking
+
+- Refactor `Index` to use `NS ((:~:) x) xs`. Pattern synonyms `IZ` and `IS` are
+  provided, making this change transparent for users.
diff --git a/sop-extras/src/Data/SOP/Index.hs b/sop-extras/src/Data/SOP/Index.hs
@@ -4,16 +4,17 @@
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE StandaloneKindSignatures #-}
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE ViewPatterns #-}
 
 module Data.SOP.Index (
     -- * Indexing SOP types
-    Index (..)
+    Index (.., IS, IZ)
   , dictIndexAll
   , indices
   , injectNS
@@ -36,43 +37,52 @@ module Data.SOP.Index (
   ) where
 
 import           Data.Coerce
-import           Data.Kind (Type)
 import           Data.Proxy
 import           Data.SOP.BasicFunctors
 import           Data.SOP.Constraint
 import           Data.SOP.Dict (Dict (..))
 import           Data.SOP.Sing
 import           Data.SOP.Strict
+import           Data.Type.Equality ((:~:) (..))
 import           Data.Word
 
-type Index :: [k] -> k -> Type
-data Index xs x where
-  IZ ::               Index (x ': xs) x
-  IS :: Index xs x -> Index (y ': xs) x
+newtype Index xs x = Index { getIndex :: NS ((:~:) x) xs }
+
+pattern IZ ::
+     ()
+  => xs ~ (x ': xs1)
+  => Index xs x
+pattern IZ = Index (Z Refl)
+
+pattern IS ::
+     ()
+  => xs ~ (x' ': xs')
+  => Index xs' x
+  -> Index xs x
+pattern IS z <- Index (S (Index -> z)) where
+  IS (Index z) = Index (S z)
+
+{-# COMPLETE IZ, IS #-}
 
 indices :: forall xs. SListI xs => NP (Index xs) xs
 indices = case sList @xs of
     SNil  -> Nil
     SCons -> IZ :* hmap IS indices
 
-dictIndexAll :: All c xs => Proxy c -> Index xs x -> Dict c x
-dictIndexAll p = \case
-    IZ      -> Dict
-    IS idx' -> dictIndexAll p idx'
+dictIndexAll :: forall c xs x. All c xs => Proxy c -> Index xs x -> Dict c x
+dictIndexAll p (Index idx) = hcollapse $ hcmap p (\Refl -> K Dict) idx
 
-injectNS :: forall f x xs. Index xs x -> f x -> NS f xs
-injectNS idx x = case idx of
-    IZ      -> Z x
-    IS idx' -> S (injectNS idx' x)
+injectNS :: forall f x xs. All Top xs => Index xs x -> f x -> NS f xs
+injectNS (Index idx) x = hmap (\Refl -> x) idx
 
 injectNS' ::
-     forall f a b x xs. (Coercible a (f x), Coercible b (NS f xs))
+     forall f a b x xs. All Top xs => (Coercible a (f x), Coercible b (NS f xs))
   => Proxy f -> Index xs x -> a -> b
 injectNS' _ idx = coerce . injectNS @f idx . coerce
 
-projectNP :: Index xs x -> NP f xs -> f x
-projectNP IZ       (x :* _ ) = x
-projectNP (IS idx) (_ :* xs) = projectNP idx xs
+projectNP :: All Top xs => Index xs x -> NP f xs -> f x
+projectNP (Index idx) np =
+  hcollapse $ hliftA2 (\f Refl -> K f) np idx
 
 {-------------------------------------------------------------------------------
   Zipping with indices
@@ -198,9 +208,9 @@ nsFromIndex n = go 0 sList
     go !_ SNil    = Nothing
 
 toWord8 :: Index xs x -> Word8
-toWord8 = go 0
+toWord8 = go 0 . getIndex
   where
-    go :: Word8 -> Index ys y -> Word8
+    go :: Word8 -> NS ((:~:) y) ys -> Word8
     go !n = \case
-      IZ      -> n
-      IS idx' -> go (n + 1) idx'
+      Z Refl -> n
+      S idx  -> go (n + 1) idx
