Slightly abstract over de Bruijn environments

We're going to end up with a bunch of these, especially once we add a COLA based one. This adds a small abstraction so that we can easily swap the implementation out (although we should keep it monomorphic in the CEK machine iteslf). While I was at it, I pulled the environments stuff and their benchmarks out to an internal library and tidied the benchmarks a bit to make use of the new classs.
IntersectMBO · michaelpj · Jun 28, 2021 · Jun 25, 2021 · Jun 28, 2021 · cc98443142371c4d18fcd30df53debcdd07c490f
commit cc98443142371c4d18fcd30df53debcdd07c490f
diff --git a/nix/pkgs/haskell/materialized-darwin/.plan.nix/plutus-core.nix b/nix/pkgs/haskell/materialized-darwin/.plan.nix/plutus-core.nix
diff --git a/nix/pkgs/haskell/materialized-linux/.plan.nix/plutus-core.nix b/nix/pkgs/haskell/materialized-linux/.plan.nix/plutus-core.nix
diff --git a/plutus-core/index-envs/bench/Main.hs b/plutus-core/index-envs/bench/Main.hs
@@ -0,0 +1,97 @@
+{-# LANGUAGE BangPatterns     #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeFamilies     #-}
+module Main where
+
+import           Criterion.Main
+import           Data.Function
+import           Data.Semigroup
+import           System.Random
+import           Unsafe.Coerce
+
+import           Data.DeBruijnEnv
+import qualified Data.DeBruijnEnv                 as DBE
+import qualified Data.IntMap.Strict               as I
+import qualified Data.RAList                      as R
+import qualified Data.RandomAccessList.SkewBinary as B
+
+
+main :: IO ()
+main = defaultMain
+    -- NOTE: there is a faster/better way to create a map using fromAscList
+    -- but we want to bench cons-ing because that is what we are using in our machine.
+    [ bgroup "create" $ flip fmap [100, 250] $ \sz ->
+            bgroup (show sz) [ bench "bral" $ whnf (ext @(B.RAList ())) sz
+                             , bench "ral" $ whnf (ext @(R.RAList ())) sz
+                             , bench "rmap" $ whnf (ext @(DBE.RelativizedMap ())) sz
+                             ]
+
+    , bgroup "query/front" $ flip fmap [100, 250] $ \sz ->
+            bgroup (show sz) [ bench "bral" $ whnf (queryFront sz) (ext @(B.RAList ()) sz)
+                             , bench "ral" $ whnf (queryFront sz) (ext @(R.RAList ()) sz)
+                             , bench "rmap" $ whnf (queryFront sz) (ext @(DBE.RelativizedMap ()) sz)
+                             ]
+
+    , bgroup "query/back" $ flip fmap [100, 250] $ \sz ->
+            bgroup (show sz) [ bench "bral" $ whnf (queryBack sz) (ext @(B.RAList ()) sz)
+                             , bench "ral" $ whnf (queryBack sz) (ext @(R.RAList ()) sz)
+                             , bench "rmap" $ whnf (queryBack sz) (ext @(DBE.RelativizedMap ()) sz)
+                             ]
+
+    , bgroup "query/rand" $ flip fmap [100, 250] $ \sz ->
+            bgroup (show sz) [ bench "bral" $ whnf (uncurry queryRand) (randWord sz, ext @(B.RAList ()) sz)
+                             , bench "ral" $ whnf (uncurry queryRand) (randWord sz, ext @(R.RAList ()) sz)
+                             , bench "rmap" $ whnf (uncurry queryRand) (randWord sz, ext @(DBE.RelativizedMap ()) sz)
+                             ]
+
+    , bgroup "create/front100/cons100/back100/cons100/rand" $ flip fmap [100, 250] $ \sz ->
+            let qsize = 100
+            in bgroup (show sz) [ bench "bral" $ whnf (uncurry $ mix qsize qsize qsize qsize) (randWord sz, ext @(B.RAList ()) sz)
+                                , bench "ral" $ whnf (uncurry $ mix qsize qsize qsize qsize) (randWord sz, ext @(B.RAList ()) sz)
+                                , bench "rmap" $ whnf (uncurry $ mix qsize qsize qsize qsize) (randWord sz, ext @(DBE.RelativizedMap ()) sz)
+                                ]
+
+    ]
+  where
+        -- the Words in these lists are smaller than maxBound :: Int
+        -- so they will not overflow when unsafe coerced to Int
+        ext :: (DeBruijnEnv e, Element e ~ ()) => Word -> e
+        ext = extend empty
+        -- if the range is the same, they should produce the same numbers for word and int
+        randWord :: Word -> [Word]
+        randWord sz = take (fromIntegral sz) $ randomRs (0,sz-1) g
+        randInt :: Word -> [Int]
+        randInt sz = take (fromIntegral sz) $ randomRs (0,fromIntegral sz-1) g
+        -- note: fixed rand-seed to make benchmarks deterministic
+        g = mkStdGen 59950
+
+applyN :: Integral b => (a -> a) -> a -> b -> a
+applyN f init n = appEndo (stimes n $ Endo f) init
+
+extend :: (DeBruijnEnv e, Element e ~ ()) => e -> Word -> e
+extend = applyN $ cons ()
+
+queryFront :: (DeBruijnEnv e, Element e ~ ()) => Word -> e -> Element e
+queryFront 0 _ = ()
+queryFront !i d = index d i' `seq` queryFront i' d
+  where i' = i-1
+
+queryBack :: (DeBruijnEnv e, Element e ~ ()) => Word -> e -> Element e
+queryBack size = go 0
+ where
+   go !i d | i == size = ()
+           | otherwise = index d i `seq` go (i+1) d
+
+queryRand :: (DeBruijnEnv e, Element e ~ ()) => [Word] -> e -> Element e
+queryRand [] _     = ()
+queryRand (i:is) d = index d i `seq` queryRand is d
+
+mix :: (DeBruijnEnv e, Element e ~ ()) => Word -> Word -> Word -> Word -> [Word] -> e -> Element e
+mix front cons1 back cons2 rand d =
+    queryFront front d
+    `seq`
+    let d1 = extend d cons1
+    in queryBack back d1
+    `seq`
+    let d2 = extend d1 cons2
+    in queryRand rand d2
diff --git a/plutus-core/index-envs/src/Data/DeBruijnEnv.hs b/plutus-core/index-envs/src/Data/DeBruijnEnv.hs
@@ -0,0 +1,61 @@
+{-# LANGUAGE TypeFamilies #-}
+module Data.DeBruijnEnv (DeBruijnEnv (..), RelativizedMap (..)) where
+
+import qualified Data.IntMap.Strict               as IM
+import           Data.Kind
+import           Data.Maybe                       (fromJust)
+import qualified Data.RAList                      as RAL
+import qualified Data.RandomAccessList.SkewBinary as BRAL
+
+{-|
+A class for types that can be used to implement a de Bruijn index environment.
+-}
+class DeBruijnEnv e where
+    -- | The type of elements in the environment.
+    type Element e :: Type
+
+    -- | The empty environment.
+    empty :: e
+    -- | Prepend an element to the environment.
+    cons :: Element e -> e -> e
+    -- | Lookup an element in the environment.
+    index :: e -> Word -> Maybe (Element e)
+    {-# INLINABLE unsafeIndex #-}
+    -- | Lookup an element in the environment, partially.
+    unsafeIndex :: e -> Word -> Element e
+    unsafeIndex e i = fromJust $ index e i
+
+instance DeBruijnEnv (BRAL.RAList a) where
+    type Element (BRAL.RAList a) = a
+
+    {-# INLINABLE empty #-}
+    empty = BRAL.Nil
+    {-# INLINABLE cons #-}
+    cons = BRAL.Cons
+    {-# INLINABLE index #-}
+    index = BRAL.safeIndex
+    {-# INLINABLE unsafeIndex #-}
+    unsafeIndex = BRAL.index
+
+-- | A sequence implemented by a map from "levels" to values and a counter giving the "current" level.
+data RelativizedMap a = RelativizedMap (IM.IntMap a) {-# UNPACK #-} !Int
+
+instance DeBruijnEnv (RelativizedMap a) where
+    type Element (RelativizedMap a) = a
+
+    {-# INLINABLE empty #-}
+    empty = RelativizedMap mempty 0
+    {-# INLINABLE cons #-}
+    cons a (RelativizedMap im l) = RelativizedMap (IM.insert l a im) (l+1)
+    {-# INLINABLE index #-}
+    index (RelativizedMap im l) w = IM.lookup (l - fromIntegral w) im
+
+instance DeBruijnEnv (RAL.RAList  a) where
+    type Element (RAL.RAList a) = a
+
+    {-# INLINABLE empty #-}
+    empty = mempty
+    {-# INLINABLE cons #-}
+    cons = RAL.cons
+    {-# INLINABLE index #-}
+    index l w = l RAL.!? fromIntegral w
diff --git a/...e/src/Data/RandomAccessList/SkewBinary.hs → ...s/src/Data/RandomAccessList/SkewBinary.hs b/...e/src/Data/RandomAccessList/SkewBinary.hs → ...s/src/Data/RandomAccessList/SkewBinary.hs
@@ -5,6 +5,7 @@
 {-# LANGUAGE ViewPatterns    #-}
 module Data.RandomAccessList.SkewBinary ( RAList(Cons,Nil)
                  , index
+                 , safeIndex
                  , Data.RandomAccessList.SkewBinary.null
                  , Data.RandomAccessList.SkewBinary.head
                  , Data.RandomAccessList.SkewBinary.tail
@@ -101,4 +102,19 @@ index (BHead w t ts) !i  =
            then indexTree halfSize (offset - 1) t1
            else indexTree halfSize (offset - 1 - halfSize) t2
 
--- TODO: safeIndex
+safeIndex :: RAList a -> Word -> Maybe a
+safeIndex Nil _  = Nothing
+safeIndex (BHead w t ts) !i  =
+    if i < w
+    then indexTree w i t
+    else safeIndex ts (i-w)
+  where
+    indexTree :: Word -> Word -> 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
diff --git a/...re/untyped-plutus-core/test/TestRAList.hs → plutus-core/index-envs/test/TestRAList.hs b/...re/untyped-plutus-core/test/TestRAList.hs → plutus-core/index-envs/test/TestRAList.hs
diff --git a/plutus-core/plutus-core.cabal b/plutus-core/plutus-core.cabal
@@ -135,7 +135,6 @@ library
         Common
         Crypto
         Data.ByteString.Hash
-        Data.RandomAccessList.SkewBinary
         Data.SatInt
         Data.Text.Prettyprint.Doc.Custom
         ErrorCode
@@ -516,26 +515,40 @@ benchmark cost-model-test
     other-modules:
         CostModelCreation
 
-benchmark bral-bench
+library index-envs
+    hs-source-dirs: index-envs/src
+    default-language: Haskell2010
+    exposed-modules:
+      Data.DeBruijnEnv
+      Data.RandomAccessList.SkewBinary
+    build-depends:
+        base -any,
+        containers -any,
+        -- broken for ral-0.2 conflic with  cardano-binary:recursion-schemes
+        ral == 0.1
+
+benchmark index-envs-bench
     type: exitcode-stdio-1.0
-    hs-source-dirs:      untyped-plutus-core/bench
+    hs-source-dirs: index-envs/bench
+    default-language: Haskell2010
     main-is: Main.hs
     build-depends:
         base -any,
-        plutus-core -any,
+        index-envs -any,
         criterion >= 1.5.9.0,
         random >= 1.2.0,
         containers -any,
         -- broken for ral-0.2 conflic with  cardano-binary:recursion-schemes
         ral == 0.1
 
-test-suite bral-test
+test-suite index-envs-test
     type: exitcode-stdio-1.0
-    hs-source-dirs:      untyped-plutus-core/test
+    hs-source-dirs: index-envs/test
+    default-language: Haskell2010
     main-is: TestRAList.hs
     build-depends:
         base -any,
-        plutus-core -any,
+        index-envs -any,
         tasty -any,
         tasty-hunit -any,
         tasty-quickcheck -any