[Builtins] Add an inlinable version of 'geq'

plutus-core/plutus-core/src/PlutusCore/Builtin/KnownType.hs

@@ -1,22 +1,25 @@
-{-# LANGUAGE BlockArguments         #-}
-{-# LANGUAGE ConstraintKinds        #-}
-{-# LANGUAGE DataKinds              #-}
-{-# LANGUAGE DefaultSignatures      #-}
-{-# LANGUAGE FlexibleInstances      #-}
-{-# LANGUAGE FunctionalDependencies #-}
-{-# LANGUAGE LambdaCase             #-}
-{-# LANGUAGE MultiParamTypeClasses  #-}
-{-# LANGUAGE OverloadedStrings      #-}
-{-# LANGUAGE TemplateHaskell        #-}
-{-# LANGUAGE TypeApplications       #-}
-{-# LANGUAGE TypeFamilies           #-}
-{-# LANGUAGE TypeOperators          #-}
-{-# LANGUAGE UndecidableInstances   #-}
-
-{-# LANGUAGE StrictData             #-}
+{-# LANGUAGE BlockArguments           #-}
+{-# LANGUAGE ConstraintKinds          #-}
+{-# LANGUAGE DataKinds                #-}
+{-# LANGUAGE DefaultSignatures        #-}
+{-# LANGUAGE FlexibleInstances        #-}
+{-# LANGUAGE FunctionalDependencies   #-}
+{-# LANGUAGE LambdaCase               #-}
+{-# LANGUAGE MultiParamTypeClasses    #-}
+{-# LANGUAGE OverloadedStrings        #-}
+{-# LANGUAGE StandaloneKindSignatures #-}
+{-# LANGUAGE TemplateHaskell          #-}
+{-# LANGUAGE TypeApplications         #-}
+{-# LANGUAGE TypeFamilies             #-}
+{-# LANGUAGE TypeOperators            #-}
+{-# LANGUAGE UndecidableInstances     #-}
+
+{-# LANGUAGE StrictData               #-}
 
 module PlutusCore.Builtin.KnownType
     ( BuiltinError
+    , GEqL (..)
+    , LoopBreaker (..)
     , KnownBuiltinTypeIn
     , KnownBuiltinType
     , BuiltinResult (..)
@@ -48,17 +51,52 @@ import Control.Monad.Except
 import Data.Bifunctor
 import Data.Either.Extras
 import Data.Functor.Identity
+import Data.Kind qualified as GHC
 import Data.String
 import GHC.Exts (inline, oneShot)
 import GHC.TypeLits
 import Prettyprinter
 import Text.PrettyBy.Internal
 import Universe
 
+-- | A version of 'GEq' that fixes @a@ in place, which allows us to create an inlinable recursive
+-- implementation of 'geqL'.
+--
+-- The way it works is that whenever there's recursion, we look up the recursive case in the current
+-- context (i.e. the dictionary) instead of actually calling 'geqL' recursively (even though it's
+-- gonna look like we do exactly that, because there's no way to distinguish between a recursive
+-- call and a dictionary lookup as the two share the same name, although to help GHC choose a lookup
+-- we sprinkle the perhaps unreliable 'LoopBreaker' in the 'DefaultUni' instance of this class).
+--
+-- Alligning things this way allows us to inline arbitrarily deep recursion for as long as types
+-- keep being monomorphic.
+--
+-- For example, the 'MapData' builtin accepts a @[(Data, Data)]@ and with 'geqL' matching on all of
+-- 'DefaultUniProtoList', 'DefaultUniProtoPair' and 'DefaultUniData' gets inlined in the denotation
+-- of the builtin. For the 'Constr' builtin that resulted in a 4.3% speedup at the time this comment
+-- was written.
+type GEqL :: (GHC.Type -> GHC.Type) -> GHC.Type -> GHC.Constraint
+class GEqL f a where
+    geqL :: f (Esc a) -> f (Esc b) -> EvaluationResult (a :~: b)
+
+-- | In @f = ... f ...@ where @f@ is a class method, how do you know if @f@ is going to be a
+-- recursive call or a type class method call? If both type check, then you don't really know how
+-- GHC is going to play it. So we add this data type to make sure that the RHS @f@ will have to
+-- become a type class method call.
+--
+-- Can GHC turn that method call into a recursive one once type classes are resolved? Dunno, but at
+-- least we've introduced an obstacle preventing GHC from immediately creating a non-inlinable
+-- recursive definition.
+newtype LoopBreaker uni a = LoopBreaker (uni a)
+
+instance GEqL uni a => GEqL (LoopBreaker uni) a where
+    geqL = coerce $ geqL @uni
+    {-# INLINE geqL #-}
+
 -- | A constraint for \"@a@ is a 'ReadKnownIn' and 'MakeKnownIn' by means of being included
 -- in @uni@\".
 type KnownBuiltinTypeIn uni val a =
-    (HasConstantIn uni val, PrettyParens (SomeTypeIn uni), GEq uni, uni `HasTermLevel` a)
+    (HasConstantIn uni val, PrettyParens (SomeTypeIn uni), GEqL uni a, uni `HasTermLevel` a)
 
 -- | A constraint for \"@a@ is a 'ReadKnownIn' and 'MakeKnownIn' by means of being included
 -- in @UniOf term@\".
@@ -277,9 +315,10 @@ readKnownConstant val = asConstant val >>= oneShot \case
         -- 'geq' matches on its first argument first, so we make the type tag that will be known
         -- statically (because this function will be inlined) go first in order for GHC to
         -- optimize some of the matching away.
-        case uniExp `geq` uniAct of
-            Just Refl -> pure x
-            Nothing   -> throwError $ BuiltinUnliftingEvaluationError $ typeMismatchError uniExp uniAct
+        case uniExp `geqL` uniAct of
+            EvaluationSuccess Refl -> pure x
+            EvaluationFailure ->
+                throwError . BuiltinUnliftingEvaluationError $ typeMismatchError uniExp uniAct
 {-# INLINE readKnownConstant #-}
 
 -- | A non-empty spine. Isomorphic to 'NonEmpty', except is strict and is defined as a single

plutus-core/plutus-core/src/PlutusCore/Default/Builtins.hs

@@ -26,6 +26,7 @@ import PlutusCore.Evaluation.Machine.ExBudgetStream (ExBudgetStream)
 import PlutusCore.Evaluation.Machine.ExMemoryUsage (ExMemoryUsage, IntegerCostedLiterally (..),
                                                     NumBytesCostedAsNumWords (..), memoryUsage,
                                                     singletonRose)
+import PlutusCore.Evaluation.Result (EvaluationResult (..))
 import PlutusCore.Pretty (PrettyConfigPlc)
 import PlutusCore.Value (Value)
 import PlutusCore.Value qualified as Value
@@ -844,12 +845,12 @@ Our final example is this:
                 :: SomeConstant uni a -> SomeConstant uni [a] -> BuiltinResult (Opaque val [a])
             mkConsDenotation
               (SomeConstant (Some (ValueOf uniA x)))
-              (SomeConstant (Some (ValueOf uniListA xs))) = do
+              (SomeConstant (Some (ValueOf uniListA xs))) =
                 case uniListA of
                     DefaultUniList uniA' -> case uniA `geq` uniA' of       -- [1]
                         Just Refl ->                                       -- [2]
                             pure . fromValueOf uniListA $ x : xs           -- [3]
-                        _ -> throwError $ structuralUnliftingError
+                        Nothing -> throwError $ structuralUnliftingError
                             "The type of the value does not match the type of elements in the list"
                     _ -> throwError $ structuralUnliftingError "Expected a list but got something else"
             {-# INLINE mkConsDenotation #-}
@@ -1425,7 +1426,7 @@ instance uni ~ DefaultUni => ToBuiltinMeaning uni DefaultFun where
                 case uniListA of
                     DefaultUniList uniA' -> case uniA `geq` uniA' of
                         Just Refl -> pure . fromValueOf uniListA $ x : xs
-                        _         -> throwError $ structuralUnliftingError
+                        Nothing   -> throwError $ structuralUnliftingError
                             "The type of the value does not match the type of elements in the list"
                     _ -> throwError $ structuralUnliftingError "Expected a list but got something else"
             {-# INLINE mkConsDenotation #-}

plutus-core/plutus-core/src/PlutusCore/Default/Universe.hs

@@ -131,6 +131,57 @@ pattern DefaultUniArray uniA =
 pattern DefaultUniPair uniA uniB =
     DefaultUniProtoPair `DefaultUniApply` uniA `DefaultUniApply` uniB
 
+-- Removing 'LoopBreaker' didn't change anything at the time this comment was written, but we kept
+-- it, because it hopefully provides some additional assurance that 'geqL' will not get elaborated
+-- as a recursive definition.
+instance AllBuiltinArgs DefaultUni (GEqL DefaultUni) a => GEqL DefaultUni a where
+    geqL DefaultUniInteger a2 = do
+        DefaultUniInteger <- pure a2
+        pure Refl
+    geqL DefaultUniByteString a2 = do
+        DefaultUniByteString <- pure a2
+        pure Refl
+    geqL DefaultUniString a2 = do
+        DefaultUniString <- pure a2
+        pure Refl
+    geqL DefaultUniUnit a2 = do
+        DefaultUniUnit <- pure a2
+        pure Refl
+    geqL DefaultUniBool a2 = do
+        DefaultUniBool <- pure a2
+        pure Refl
+    geqL (DefaultUniProtoList `DefaultUniApply` a1) listA2 = do
+        DefaultUniProtoList `DefaultUniApply` a2 <- pure listA2
+        Refl <- geqL (LoopBreaker a1) (LoopBreaker a2)
+        pure Refl
+    geqL (DefaultUniProtoArray `DefaultUniApply` a1) arrayA2 = do
+        DefaultUniProtoArray `DefaultUniApply` a2 <- pure arrayA2
+        Refl <- geqL (LoopBreaker a1) (LoopBreaker a2)
+        pure Refl
+    geqL (DefaultUniProtoPair `DefaultUniApply` a1 `DefaultUniApply` b1) pairA2 = do
+        DefaultUniProtoPair `DefaultUniApply` a2 `DefaultUniApply` b2 <- pure pairA2
+        Refl <- geqL (LoopBreaker a1) (LoopBreaker a2)
+        Refl <- geqL (LoopBreaker b1) (LoopBreaker b2)
+        pure Refl
+    geqL (f `DefaultUniApply` _ `DefaultUniApply` _ `DefaultUniApply` _) _ =
+        noMoreTypeFunctions f
+    geqL DefaultUniData a2 = do
+        DefaultUniData <- pure a2
+        pure Refl
+    geqL DefaultUniBLS12_381_G1_Element a2 = do
+        DefaultUniBLS12_381_G1_Element <- pure a2
+        pure Refl
+    geqL DefaultUniBLS12_381_G2_Element a2 = do
+        DefaultUniBLS12_381_G2_Element <- pure a2
+        pure Refl
+    geqL DefaultUniBLS12_381_MlResult a2 = do
+        DefaultUniBLS12_381_MlResult <- pure a2
+        pure Refl
+    geqL DefaultUniValue a2 = do
+        DefaultUniValue <- pure a2
+        pure Refl
+    {-# INLINE geqL #-}
+
 instance GEq DefaultUni where
     -- We define 'geq' manually instead of using 'deriveGEq', because the latter creates a single
     -- recursive definition and we want two instead. The reason why we want two is because this
@@ -140,59 +191,63 @@ instance GEq DefaultUni where
     -- (we're not really sure if this is a reliable solution, but if it stops working, we won't miss
     -- very much and we've failed to settle on any other approach).
     --
-    -- This trick gives us a 1% speedup across validation benchmarks (some are up to 4% faster) and
-    -- a more sensible generated Core where things like @geq DefaulUniBool@ are reduced away.
-    geq = geqStep where
-        geqStep :: DefaultUni a1 -> DefaultUni a2 -> Maybe (a1 :~: a2)
-        geqStep DefaultUniInteger a2 = do
-            DefaultUniInteger <- Just a2
-            Just Refl
-        geqStep DefaultUniByteString a2 = do
-            DefaultUniByteString <- Just a2
-            Just Refl
-        geqStep DefaultUniString a2 = do
-            DefaultUniString <- Just a2
-            Just Refl
-        geqStep DefaultUniUnit a2 = do
-            DefaultUniUnit <- Just a2
-            Just Refl
-        geqStep DefaultUniBool a2 = do
-            DefaultUniBool <- Just a2
-            Just Refl
-        geqStep DefaultUniProtoList a2 = do
-            DefaultUniProtoList <- Just a2
-            Just Refl
-        geqStep DefaultUniProtoArray a2 = do
-            DefaultUniProtoArray <- Just a2
-            Just Refl
-        geqStep DefaultUniProtoPair a2 = do
-            DefaultUniProtoPair <- Just a2
-            Just Refl
-        geqStep (DefaultUniApply f1 x1) a2 = do
-            DefaultUniApply f2 x2 <- Just a2
-            Refl <- geqRec f1 f2
-            Refl <- geqRec x1 x2
-            Just Refl
-        geqStep DefaultUniData a2 = do
-            DefaultUniData <- Just a2
-            Just Refl
-        geqStep DefaultUniBLS12_381_G1_Element a2 = do
-            DefaultUniBLS12_381_G1_Element <- Just a2
-            Just Refl
-        geqStep DefaultUniBLS12_381_G2_Element a2 = do
-            DefaultUniBLS12_381_G2_Element <- Just a2
-            Just Refl
-        geqStep DefaultUniBLS12_381_MlResult a2 = do
-            DefaultUniBLS12_381_MlResult <- Just a2
-            Just Refl
-        geqStep DefaultUniValue a2 = do
-            DefaultUniValue <- Just a2
-            Just Refl
-        {-# INLINE geqStep #-}
-
-        geqRec :: DefaultUni a1 -> DefaultUni a2 -> Maybe (a1 :~: a2)
-        geqRec = geqStep
-        {-# OPAQUE geqRec #-}
+    -- On the critical path this definition should only be used for builtins that perform equality
+    -- checking of statically unknown runtime type tags ('MkCons' is one such builtin for
+    -- example). All other builtins should use 'geqL' (the latter is internal to 'readKnownConstant'
+    -- and is therefore hidden from the person adding a new builtin).
+    --
+    -- We use @NOINLINE@ instead of @OPAQUE@, because we don't actually care about the recursive
+    -- definition not being inlined, we just want it to be chosen as the loop breaker.
+    geq = goStep where
+        goStep, goRec :: DefaultUni a1 -> DefaultUni a2 -> Maybe (a1 :~: a2)
+        goStep DefaultUniInteger a2 = do
+            DefaultUniInteger <- pure a2
+            pure Refl
+        goStep DefaultUniByteString a2 = do
+            DefaultUniByteString <- pure a2
+            pure Refl
+        goStep DefaultUniString a2 = do
+            DefaultUniString <- pure a2
+            pure Refl
+        goStep DefaultUniUnit a2 = do
+            DefaultUniUnit <- pure a2
+            pure Refl
+        goStep DefaultUniBool a2 = do
+            DefaultUniBool <- pure a2
+            pure Refl
+        goStep DefaultUniProtoList a2 = do
+            DefaultUniProtoList <- pure a2
+            pure Refl
+        goStep DefaultUniProtoArray a2 = do
+            DefaultUniProtoArray <- pure a2
+            pure Refl
+        goStep DefaultUniProtoPair a2 = do
+            DefaultUniProtoPair <- pure a2
+            pure Refl
+        goStep (DefaultUniApply f1 x1) a2 = do
+            DefaultUniApply f2 x2 <- pure a2
+            Refl <- oneShot goRec f1 f2
+            Refl <- oneShot goRec x1 x2
+            pure Refl
+        goStep DefaultUniData a2 = do
+            DefaultUniData <- pure a2
+            pure Refl
+        goStep DefaultUniBLS12_381_G1_Element a2 = do
+            DefaultUniBLS12_381_G1_Element <- pure a2
+            pure Refl
+        goStep DefaultUniBLS12_381_G2_Element a2 = do
+            DefaultUniBLS12_381_G2_Element <- pure a2
+            pure Refl
+        goStep DefaultUniBLS12_381_MlResult a2 = do
+            DefaultUniBLS12_381_MlResult <- pure a2
+            pure Refl
+        goStep DefaultUniValue a2 = do
+            DefaultUniValue <- pure a2
+            pure Refl
+        {-# INLINE goStep #-}
+
+        goRec = goStep
+        {-# NOINLINE goRec #-}
 
 -- | For pleasing the coverage checker.
 noMoreTypeFunctions :: DefaultUni (Esc (f :: a -> b -> c -> d)) -> any