Try a smallarray

IntersectMBO · Mar 17, 2023 · 33a8a98 · 33a8a98
1 parent bbbcac8
commit 33a8a98
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 13 deletions.
diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/Cek/Internal.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/Cek/Internal.hs
@@ -759,10 +759,8 @@ enterComputeCek = computeCek
     -- some reason.
     -- Skip index 7, that's the total counter!
     -- See Note [Structure of the step counter]
-    {-# INLINE spend #-}
     spend !i !w = unless (i == 7) $ let kind = toEnum i in spendBudgetCek (BStep kind) (stimes w (cekStepCost ?cekCosts kind))
 
-    {-# INLINE stepAndMaybeSpend #-}
     -- | Accumulate a step, and maybe spend the budget that has accumulated for a number of machine steps, but only if we've exceeded our slippage.
     stepAndMaybeSpend :: StepKind -> CekM uni fun s ()
     stepAndMaybeSpend !kind = do

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
@@ -5,22 +5,28 @@ import Control.Monad.Primitive
 import Data.Primitive qualified as P
 import Data.Word
 
-newtype StepCounter s = StepCounter (P.MutablePrimArray s Word8)
+newtype StepCounter s = StepCounter (P.SmallMutableArray s Word8)
 
 {-# INLINE newCounter #-}
 newCounter :: PrimMonad m => Int -> m (StepCounter (PrimState m))
-newCounter sz = do
-  c <- StepCounter <$> P.newPrimArray sz
-  resetCounter c
-  pure c
+newCounter sz = StepCounter <$> P.newSmallArray sz 0
 
 {-# INLINE resetCounter #-}
 resetCounter :: PrimMonad m => StepCounter (PrimState m) -> m ()
-resetCounter (StepCounter arr) = P.setPrimArray arr 0 (P.sizeofMutablePrimArray arr) 0
+resetCounter (StepCounter arr) = do
+  -- TODO: better impl?
+  let
+    sz = P.sizeofSmallMutableArray arr
+    go !i
+      | i < sz = do
+          P.writeSmallArray arr i 0
+          go (i+1)
+      | otherwise = pure ()
+  go 0
 
 {-# INLINE readCounter #-}
 readCounter :: PrimMonad m => StepCounter (PrimState m) -> Int -> m Word8
-readCounter (StepCounter arr) = P.readPrimArray arr
+readCounter (StepCounter arr) = P.readSmallArray arr
 
 {-# INLINE writeCounter #-}
 writeCounter
@@ -29,7 +35,7 @@ writeCounter
   -> Int
   -> Word8
   -> m ()
-writeCounter (StepCounter arr) ix v = P.writePrimArray arr ix v
+writeCounter (StepCounter arr) ix v = P.writeSmallArray arr ix v
 
 {-# INLINE overIndex #-}
 overIndex
@@ -47,12 +53,12 @@ overIndex i f c = do
 {-# INLINE itraverseCounter_ #-}
 itraverseCounter_ :: (PrimMonad m) => (Int -> Word8 -> m ()) -> StepCounter (PrimState m) -> m ()
 itraverseCounter_ f (StepCounter arr) = do
-  arr' <- P.unsafeFreezePrimArray arr
+  arr' <- P.unsafeFreezeSmallArray arr
   let
-    sz = P.sizeofPrimArray arr'
+    sz = P.sizeofSmallArray arr'
     go !i
       | i < sz = do
-          f i (P.indexPrimArray arr' i)
+          f i (P.indexSmallArray arr' i)
           go (i+1)
       | otherwise = pure ()
   go 0