[Evaluation] Statically unroll 'itraverseCounter_'

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,41 +80,52 @@ 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
   -- are done with the frozen version. That ordering is enforced by the fact that
   -- 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 ()