diff --git a/plutus-core/changelog.d/20241126_025636_effectfully_statically_unroll_itraverseCounter_.md b/plutus-core/changelog.d/20241126_025636_effectfully_statically_unroll_itraverseCounter_.md new file mode 100644 index 00000000000..4e84188ab82 --- /dev/null +++ b/plutus-core/changelog.d/20241126_025636_effectfully_statically_unroll_itraverseCounter_.md @@ -0,0 +1,3 @@ +### Changed + +- In #5265 made `itraverseCounter_` faster increasing overall performance of the evaluator by 2.5%. diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/Cek/StepCounter.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/Cek/StepCounter.hs index 8d4419a26ce..10fa8460fe4 100644 --- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/Cek/StepCounter.hs +++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/Cek/StepCounter.hs @@ -1,16 +1,23 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE KindSignatures #-} -{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE StandaloneKindSignatures #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} module UntypedPlutusCore.Evaluation.Machine.Cek.StepCounter where import Control.Monad.Primitive import Data.Coerce (coerce) +import Data.Kind import Data.Primitive qualified as P import Data.Proxy import Data.Word -import GHC.TypeLits (Nat) -import GHC.TypeNats (KnownNat, natVal) +import GHC.TypeNats (KnownNat, Nat, natVal, type (-)) -- See Note [Step counter data structure] -- You might think that since we can store whatever we like in here we might as well @@ -73,18 +80,39 @@ modifyCounter i f c = do pure modified {-# INLINE modifyCounter #-} +-- | The type of natural numbers in Peano form. +data Peano + = Z + | S Peano + +type NatToPeano :: Nat -> Peano +type family NatToPeano n where + NatToPeano 0 = 'Z + NatToPeano n = 'S (NatToPeano (n - 1)) + +type UpwardsM :: (Type -> Type) -> Peano -> Constraint +class Applicative f => UpwardsM f n where + -- | @upwardsM i k@ means @k i *> k (i + 1) *> ... *> k (i + n - 1)@. + -- We use this function in order to statically unroll a loop in 'itraverseCounter_' through + -- instance resolution. This makes @validation@ benchmarks a couple of percent faster. + upwardsM :: Int -> (Int -> f ()) -> f () + +instance Applicative f => UpwardsM f 'Z where + upwardsM _ _ = pure () + {-# INLINE upwardsM #-} + +instance UpwardsM f n => UpwardsM f ('S n) where + upwardsM !i k = k i *> upwardsM @f @n (i + 1) k + {-# INLINE upwardsM #-} + -- | Traverse the counters with an effectful function. itraverseCounter_ :: forall n m - . (KnownNat n, PrimMonad m) + . (UpwardsM m (NatToPeano n), PrimMonad m) => (Int -> Word8 -> m ()) -> StepCounter n (PrimState m) -> m () itraverseCounter_ f (StepCounter arr) = do - -- The implementation of this function is very performance-sensitive. I believe - -- it may be possible to do better, but it's time-consuming to experiment more - -- and unclear what improves things. - -- The safety of this operation is a little subtle. The frozen array is only -- safe to use if the underlying mutable array is not mutated 'afterwards'. -- In our case it likely _will_ be mutated afterwards... but not until we @@ -92,22 +120,12 @@ itraverseCounter_ f (StepCounter arr) = do -- the whole thing runs in 'm': future accesses to the mutable array can't -- happen until this whole function is finished. arr' <- P.unsafeFreezePrimArray arr - let - sz = fromIntegral $ natVal (Proxy @n) - go !i - | i < sz = do - f i (P.indexPrimArray arr' i) - go (i+1) - | otherwise = pure () - go 0 --- I also tried INLINABLE + SPECIALIZE, which just resulted in it getting inlined, so might --- as well just go straight for that + upwardsM @_ @(NatToPeano n) 0 $ \i -> f i $ P.indexPrimArray arr' i {-# INLINE itraverseCounter_ #-} - -- | Traverse the counters with an effectful function. iforCounter_ - :: (KnownNat n, PrimMonad m) + :: (UpwardsM m (NatToPeano n), PrimMonad m) => StepCounter n (PrimState m) -> (Int -> Word8 -> m ()) -> m ()