[Evaluation] [Names] Define all lookups in terms of 'contIndexZero' (#6702)

Instead of 4 different implementation this will make it 1 main and 3 derivative ones.

This appears to make the validation benchmarks faster by a percent and the nofib benchmarks faster by a couple of percent while simplifying the code, so a clear win.

Author: effectfully

plutus-core/changelog.d/20241126_020822_effectfully_define_1_indexing_via_0_indexing.md

@@ -0,0 +1,3 @@
+### Changed
+
+- In #6702 made variable lookup faster increasing overall performance of the evaluator by 1%.

plutus-core/index-envs/src/Data/RandomAccessList/SkewBinary.hs

@@ -7,10 +7,11 @@
 {-# LANGUAGE ViewPatterns         #-}
 module Data.RandomAccessList.SkewBinary
     ( RAList(Cons,Nil)
+    , contIndexZero
+    , contIndexOne
     , safeIndexZero
     , unsafeIndexZero
     , safeIndexOne
-    , safeIndexOneCont
     , unsafeIndexOne
     , Data.RandomAccessList.SkewBinary.null
     , uncons
@@ -48,7 +49,7 @@ data RAList a = BHead
 null :: RAList a -> Bool
 null Nil = True
 null _   = False
-{-# INLINABLE null #-}
+{-# INLINE null #-}
 
 {-# complete Cons, Nil #-}
 {-# complete BHead, Nil #-}
@@ -63,6 +64,7 @@ cons :: a -> RAList a -> RAList a
 cons x = \case
     (BHead w1 t1 (BHead w2 t2 ts')) | w1 == w2 -> BHead (2*w1+1) (Node x t1 t2) ts'
     ts                                         -> BHead 1 (Leaf x) ts
+{-# INLINE cons #-}
 
 -- /O(1)/
 uncons :: RAList a -> Maybe (a, RAList a)
@@ -74,122 +76,94 @@ uncons = \case
             -- split the node in two)
         in Just (x, BHead halfSize t1 $ BHead halfSize t2 ts)
     Nil -> Nothing
+{-# INLINE uncons #-}
 
--- 0-based
-unsafeIndexZero :: RAList a -> Word64 -> a
-unsafeIndexZero Nil _  = error "out of bounds"
-unsafeIndexZero (BHead w t ts) !i  =
-    if i < w
-    then indexTree w i t
-    else unsafeIndexZero ts (i-w)
-  where
-    indexTree :: Word64 -> Word64 -> Tree a -> a
-    indexTree 1 0 (Leaf x) = x
-    indexTree _ _ (Leaf _) = error "out of bounds"
-    indexTree _ 0 (Node x _ _) = x
-    indexTree treeSize offset (Node _ t1 t2) =
-        let halfSize = unsafeShiftR treeSize 1 -- probably faster than `div w 2`
-        in if offset <= halfSize
-           then indexTree halfSize (offset - 1) t1
-           else indexTree halfSize (offset - 1 - halfSize) t2
-
--- 0-based
-safeIndexZero :: RAList a -> Word64 -> Maybe a
-safeIndexZero Nil _  = Nothing
-safeIndexZero (BHead w t ts) !i  =
-    if i < w
-    then indexTree w i t
-    else safeIndexZero ts (i-w)
-  where
-    indexTree :: Word64 -> Word64 -> Tree a -> Maybe a
-    indexTree 1 0 (Leaf x) = Just x
-    indexTree _ _ (Leaf _) = Nothing
-    indexTree _ 0 (Node x _ _) = Just x
-    indexTree treeSize offset (Node _ t1 t2) =
-        let halfSize = unsafeShiftR treeSize 1 -- probably faster than `div w 2`
-        in if offset <= halfSize
-           then indexTree halfSize (offset - 1) t1
-           else indexTree halfSize (offset - 1 - halfSize) t2
-
--- 1-based
-unsafeIndexOne :: RAList a -> Word64 -> a
-unsafeIndexOne Nil _ = error "out of bounds"
-unsafeIndexOne (BHead w t ts) !i =
-    if i <= w
-    then indexTree w i t
-    else unsafeIndexOne ts (i-w)
-  where
-    indexTree :: Word64 -> Word64 -> Tree a -> a
-    indexTree _ 0 _ = error "index zero"
-    indexTree 1 1 (Leaf x) = x
-    indexTree _ _ (Leaf _) = error "out of bounds"
-    indexTree _ 1 (Node x _ _) = x
-    indexTree treeSize offset (Node _ t1 t2) =
-        let halfSize = unsafeShiftR treeSize 1 -- probably faster than `div w 2`
-            offset' = offset - 1
-        in if offset' <= halfSize
-           then indexTree halfSize offset' t1
-           else indexTree halfSize (offset' - halfSize) t2
-
-{- Note [Optimizations of safeIndexOneCont]
-Bangs in the local definitions of 'safeIndexOneCont' are needed to tell GHC that the functions are
+{- Note [Optimizations of contIndexZero]
+Bangs in the local definitions of 'contIndexZero' are needed to tell GHC that the functions are
 strict in the 'Word64' argument, so that GHC produces workers operating on @Word64#@.
 
 The function itself is CPS-ed, so that the arguments force the local definitions to be retained
-within 'safeIndexOneCont' instead of being pulled out via full-laziness or some other optimization
-pass. This ensures that when 'safeIndexOneCont' gets inlined, the local definitions appear directly
-in the GHC Core, allowing GHC to inline the arguments of 'safeIndexOneCont' and transform the whole
+within 'contIndexZero' instead of being pulled out via full-laziness or some other optimization
+pass. This ensures that when 'contIndexZero' gets inlined, the local definitions appear directly
+in the GHC Core, allowing GHC to inline the arguments of 'contIndexZero' and transform the whole
 thing into a beautiful recursive join point full of @Word64#@s, i.e. allocating very little if
 anything at all.
 -}
 
--- See Note [Optimizations of safeIndexOneCont].
-safeIndexOneCont :: forall a b. b -> (a -> b) -> RAList a -> Word64 -> b
-safeIndexOneCont z f = findTree where
+-- See Note [Optimizations of contIndexZero].
+contIndexZero :: forall a b. b -> (a -> b) -> RAList a -> Word64 -> b
+contIndexZero z f = findTree where
     findTree :: RAList a -> Word64 -> b
-    -- See Note [Optimizations of safeIndexOneCont].
+    -- See Note [Optimizations of contIndexZero].
     findTree Nil !_ = z
     findTree (BHead w t ts) i =
-        if i <= w
+        if i < w
         then indexTree w i t
         else findTree ts (i-w)
 
     indexTree :: Word64 -> Word64 -> Tree a -> b
-    -- See Note [Optimizations of safeIndexOneCont].
-    indexTree !w 1 t = case t of
+    -- See Note [Optimizations of contIndexZero].
+    indexTree !w 0 t = case t of
         Node x _ _ -> f x
         Leaf x     -> if w == 1 then f x else z
-    indexTree _ 0 _ = z -- "index zero"
     indexTree _ _ (Leaf _) = z
     indexTree treeSize offset (Node _ t1 t2) =
         let halfSize = unsafeShiftR treeSize 1 -- probably faster than `div w 2`
-            offset' = offset - 1
-        in if offset' <= halfSize
-           then indexTree halfSize offset' t1
-           else indexTree halfSize (offset' - halfSize) t2
-{-# INLINE safeIndexOneCont #-}
+        in if offset <= halfSize
+           then indexTree halfSize (offset - 1) t1
+           else indexTree halfSize (offset - 1 - halfSize) t2
+{-# INLINE contIndexZero #-}
+
+contIndexOne :: forall a b. b -> (a -> b) -> RAList a -> Word64 -> b
+contIndexOne z _ _ 0 = z
+contIndexOne z f t n = contIndexZero z f t (n - 1)
+{-# INLINE contIndexOne #-}
+
+-- 0-based
+unsafeIndexZero :: RAList a -> Word64 -> a
+unsafeIndexZero = contIndexZero (error "out of bounds") id
+{-# INLINE unsafeIndexZero #-}
+
+-- 0-based
+safeIndexZero :: RAList a -> Word64 -> Maybe a
+safeIndexZero = contIndexZero Nothing Just
+{-# INLINE safeIndexZero #-}
+
+-- 1-based
+unsafeIndexOne :: RAList a -> Word64 -> a
+unsafeIndexOne = contIndexOne (error "out of bounds") id
+{-# INLINE unsafeIndexOne #-}
 
 -- 1-based
 safeIndexOne :: RAList a -> Word64 -> Maybe a
-safeIndexOne = safeIndexOneCont Nothing Just
+safeIndexOne = contIndexOne Nothing Just
+{-# INLINE safeIndexOne #-}
 
 instance RAL.RandomAccessList (RAList a) where
     type Element (RAList a) = a
 
-    {-# INLINABLE empty #-}
     empty = Nil
-    {-# INLINABLE cons #-}
+    {-# INLINE empty #-}
+
     cons = Cons
-    {-# INLINABLE uncons #-}
+    {-# INLINE cons #-}
+
     uncons = uncons
-    {-# INLINABLE length #-}
-    length Nil             = 0
-    length (BHead sz _ tl) = sz + RAL.length tl
-    {-# INLINABLE indexZero #-}
+    {-# INLINE uncons #-}
+
+    length = go 0 where
+        go !acc Nil             = acc
+        go !acc (BHead sz _ tl) = go (acc + sz) tl
+    {-# INLINE length #-}
+
     indexZero = safeIndexZero
-    {-# INLINABLE indexOne #-}
+    {-# INLINE indexZero #-}
+
     indexOne = safeIndexOne
-    {-# INLINABLE unsafeIndexZero #-}
+    {-# INLINE indexOne #-}
+
     unsafeIndexZero = unsafeIndexZero
-    {-# INLINABLE unsafeIndexOne #-}
+    {-# INLINE unsafeIndexZero #-}
+
     unsafeIndexOne = unsafeIndexOne
+    {-# INLINE unsafeIndexOne #-}

plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/Cek/Internal.hs

@@ -878,7 +878,7 @@ enterComputeCek = computeCek
     -- | Look up a variable name in the environment.
     lookupVarName :: NamedDeBruijn -> CekValEnv uni fun ann -> CekM uni fun s (CekValue uni fun ann)
     lookupVarName varName@(NamedDeBruijn _ varIx) varEnv =
-        Env.safeIndexOneCont
+        Env.contIndexOne
             (throwingWithCause _MachineError OpenTermEvaluatedMachineError . Just $ Var () varName)
             pure
             varEnv