The UnconstrConstrData rewrite rule (#5605)

Co-authored-by: Nikolaos Bezirgiannis <[email protected]>

plutus-core/changelog.d/20231101_201807_bezirg_rewrite_more.md

@@ -0,0 +1,3 @@
+### Added
+
+- A PIR rewrite rule for optimizing ""(unconstr . constrdata)"

plutus-core/plutus-core.cabal

@@ -532,6 +532,8 @@ library plutus-ir
     PlutusIR.Compiler.Lower
     PlutusIR.Compiler.Recursion
     PlutusIR.Normalize
+    PlutusIR.Transform.RewriteRules.Common
+    PlutusIR.Transform.RewriteRules.UnConstrConstrData
 
   build-depends:
     , algebraic-graphs     >=0.7
@@ -607,7 +609,6 @@ test-suite plutus-ir-test
   build-depends:
     , base                                                       >=4.9   && <5
     , containers
-    , data-default-class
     , flat                                                       ^>=0.6
     , hashable
     , hedgehog

plutus-core/plutus-core/src/PlutusCore/MkPlc.hs

@@ -4,6 +4,7 @@
 {-# LANGUAGE FlexibleInstances      #-}
 {-# LANGUAGE FunctionalDependencies #-}
 {-# LANGUAGE LambdaCase             #-}
+{-# LANGUAGE OverloadedStrings      #-}
 {-# LANGUAGE PolyKinds              #-}
 {-# LANGUAGE TupleSections          #-}
 {-# LANGUAGE TypeApplications       #-}
@@ -52,13 +53,16 @@ module PlutusCore.MkPlc
     , mkIterTyApp
     , mkIterTyAppNoAnn
     , mkIterKindArrow
+    , mkFreshTermLet
     ) where
 
 import PlutusPrelude
 import Prelude hiding (error)
 
 import PlutusCore.Builtin
 import PlutusCore.Core
+import PlutusCore.Name
+import PlutusCore.Quote
 
 import Data.Word
 import Universe
@@ -317,3 +321,16 @@ mkIterKindArrow
     -> Kind ann
     -> Kind ann
 mkIterKindArrow ann kinds target = foldr (KindArrow ann) target kinds
+
+{- | A helper to create a single, fresh strict binding; It returns the fresh bound `Var`iable and
+a function `Term -> Term`, expecting an "in-Term" to form a let-expression.
+-}
+mkFreshTermLet :: (MonadQuote m, TermLike t tyname Name uni fun, Monoid a)
+               => Type tyname uni a -- ^ the type of binding
+               -> t a -- ^ the term bound to the fresh variable
+               -> m (t a, t a -> t a) -- ^ the fresh Var and a function that takes an "in" term to construct the Let
+mkFreshTermLet aT a = do
+    -- I wish this was less constrained to Name
+    genName <- freshName "generated"
+    pure (var mempty genName, termLet mempty (Def (VarDecl mempty genName aT) a))
+

plutus-core/plutus-ir/src/PlutusIR/Compiler.hs

@@ -16,7 +16,7 @@ module PlutusIR.Compiler (
     Provenance (..),
     DatatypeComponent (..),
     noProvenance,
-    CompilationOpts,
+    CompilationOpts (..),
     coOptimize,
     coPedantic,
     coVerbose,
@@ -28,6 +28,7 @@ module PlutusIR.Compiler (
     coDoSimplifierEvaluateBuiltins,
     coDoSimplifierStrictifyBindings,
     coDoSimplifierRewrite,
+    coDoSimplifierKnownCon,
     coInlineHints,
     coProfile,
     coRelaxedFloatin,

plutus-core/plutus-ir/src/PlutusIR/Compiler/Provenance.hs

@@ -37,6 +37,9 @@ instance Ord a => Semigroup (Provenance a) where
             MultipleSources ps -> ps
             other              -> S.singleton other
 
+instance Ord a => Monoid (Provenance a) where
+    mempty = noProvenance
+
 -- workaround, use a smart constructor to replace the older NoProvenance data constructor
 noProvenance :: Provenance a
 noProvenance = MultipleSources S.empty

plutus-core/plutus-ir/src/PlutusIR/Compiler/Types.hs

@@ -104,7 +104,7 @@ makeLenses ''CompilationOpts
 
 defaultCompilationOpts :: CompilationOpts a
 defaultCompilationOpts = CompilationOpts
-  { _coOptimize = True
+  { _coOptimize = True -- synonymous with max-simplifier-iterations=0
   , _coPedantic = False
   , _coVerbose = False
   , _coDebug = False

plutus-core/plutus-ir/src/PlutusIR/Subst.hs

@@ -129,6 +129,7 @@ funRes f = \case
   TyFun a dom cod -> TyFun a dom <$> funRes f cod
   t               -> f t
 
+-- TODO: these could be Traversals
 -- | Get all the term variables in a term.
 vTerm :: Fold (Term tyname name uni fun ann) name
 vTerm = termSubtermsDeep . termVars

plutus-core/plutus-ir/src/PlutusIR/Transform/RewriteRules.hs

@@ -15,14 +15,19 @@ import PlutusCore.Quote
 import PlutusIR as PIR
 import PlutusIR.Analysis.VarInfo
 import PlutusIR.Transform.RewriteRules.CommuteFnWithConst
+import PlutusIR.Transform.RewriteRules.UnConstrConstrData
 import PlutusPrelude
 
 import Control.Lens
 
 
--- | Rewrite a `Term` using the given `RewriteRules` (similar to functions of Term -> Term)
--- Normally the rewrite rules are configured at entrypoint time of the compiler.
-rewriteWith :: ( Semigroup a, t ~ Term tyname name uni fun a
+{- | Rewrite a `Term` using the given `RewriteRules` (similar to functions of Term -> Term)
+Normally the rewrite rules are configured at entrypoint time of the compiler.
+
+It goes without saying that the supplied rewrite rules must be type-preserving.
+MAYBE: enforce this with a `through typeCheckTerm`?
+-}
+rewriteWith :: ( Monoid a, t ~ Term tyname Name uni fun a
               , HasUniques t
               , MonadQuote m
               )
@@ -39,19 +44,19 @@ rewriteWith (RewriteRules rules) t =
 
 -- | A bundle of composed `RewriteRules`, to be passed at entrypoint of the compiler.
 newtype RewriteRules uni fun = RewriteRules {
-    unRewriteRules :: forall tyname name m a
-                   . (MonadQuote m, Semigroup a)
-                   => VarsInfo tyname name uni a
-                   -> PIR.Term tyname name uni fun a
-                   -> m (PIR.Term tyname name uni fun a)
+    unRewriteRules :: forall tyname m a
+                   . (MonadQuote m, Monoid a)
+                   => VarsInfo tyname Name uni a
+                   -> PIR.Term tyname Name uni fun a
+                   -> m (PIR.Term tyname Name uni fun a)
     }
 
 -- | The rules for the Default Universe/Builtin.
 defaultUniRewriteRules :: RewriteRules DefaultUni DefaultFun
-defaultUniRewriteRules = RewriteRules $ \ _vinfo ->
+defaultUniRewriteRules = RewriteRules $ \ vinfo ->
         -- The rules are composed from left to right.
         pure . commuteFnWithConst
-        -- e.g. >=> a follow-up rewrite rule
+        >=> unConstrConstrData def vinfo
 
 instance Default (RewriteRules DefaultUni DefaultFun) where
   def = defaultUniRewriteRules

plutus-core/plutus-ir/src/PlutusIR/Transform/RewriteRules/Common.hs

@@ -0,0 +1,89 @@
+{-# LANGUAGE GADTs           #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE TypeOperators   #-}
+module PlutusIR.Transform.RewriteRules.Common
+    ( seqA
+    , seqP
+    , mkFreshTermLet -- from MkPlc
+    , pattern A
+    , pattern B
+    , pattern I
+    ) where
+
+import PlutusCore.Builtin
+import PlutusCore.Quote
+import PlutusIR
+import PlutusIR.Analysis.Builtins
+import PlutusIR.Analysis.VarInfo
+import PlutusIR.MkPir
+import PlutusIR.Purity
+
+{- | A wrapper that can be more easily turned into an infix operator.
+
+e.g. `infixr 5 (***) = seqA binfo vInfo`
+-}
+seqA :: (MonadQuote m, Monoid a, ToBuiltinMeaning uni fun)
+     => BuiltinsInfo uni fun
+     -> VarsInfo tyname Name uni a
+     -> (Type tyname uni a, Term tyname Name uni fun a)
+     ->  m (Term tyname Name uni fun a)
+     -> m (Term tyname Name uni fun a)
+seqA binfo vinfo (a,aT) y = seqOpt binfo vinfo a aT <*> y
+
+
+{- | Another "infix" wrapper where second operand is a Haskell pure value.
+
+e.g. `infixr 5 (***) = seqP binfo vInfo`
+-}
+seqP :: (MonadQuote m, Monoid a, ToBuiltinMeaning uni fun)
+     => BuiltinsInfo uni fun
+     -> VarsInfo tyname Name uni a
+     -> (Type tyname uni a, Term tyname Name uni fun a)
+     -> Term tyname Name uni fun a
+     -> m (Term tyname Name uni fun a)
+seqP binfo vinfo p y = seqA binfo vinfo p (pure y)
+
+-- | An optimized version to omit call to `seq` if left operand `isPure`.
+seqOpt :: ( MonadQuote m
+       , Monoid a
+       , ToBuiltinMeaning uni fun
+       , t ~ Term tyname Name uni fun a
+       )
+     => BuiltinsInfo uni fun
+     -> VarsInfo tyname Name uni a
+     -> Type tyname uni a -- ^ the type of left operand a
+     -> t -- ^ left operand a
+     -> m (t -> t) -- ^ how to modify right operand b
+seqOpt binfo vinfo aT a | isPure binfo vinfo a = pure id
+                        | otherwise = seqUnOpt aT a
+
+{- | Takes as input a term `a` with its type `aT`,
+and constructs a function that expects another term `b`.
+When the returned function is applied to a term, the execution of the applied term
+would strictly evaluate the first term `a` only for its effects (i.e. ignoring its result)
+while returning the result of the second term `b`.
+
+The name is intentionally taken from Haskell's `GHC.Prim.seq`.
+Currently, the need for this `seq` "combinator" is in `RewriteRules`,
+to trying to retain the effects, that would otherwise be lost if that code was instead considered
+completely dead.
+
+Unfortunately, unlike Haskell's `seq`, we need the pass the correct `Type` of `a`,
+so as to apply this `seq` combinator. This is usually not a problem because we are generating
+code and we should have the type of `a` somewhere available.
+-}
+seqUnOpt :: (MonadQuote m, Monoid a, t ~ Term tyname Name uni fun a)
+    => Type tyname uni a -- ^ the type of left operand a
+    -> t  -- ^ left operand a
+    -> m (t -> t) -- ^ how to modify right operand b
+seqUnOpt aT a = snd <$> mkFreshTermLet aT a
+
+-- Some shorthands for easier pattern-matching when creating rewrite rules
+-- TODO: these patterns ignores annotations. Find a better way for easier writing rules that does
+-- not ignore annotations e.g. (pattern-PIR-quasiquoters?)
+pattern A :: Term tyname name uni fun a -> Term tyname name uni fun a -> Term tyname name uni fun a
+pattern A l r <- Apply _ l r
+pattern B :: fun -> Term tyname name uni fun a
+pattern B b <- Builtin _ b
+pattern I :: Term tyname name uni fun a -> Type tyname uni a -> Term tyname name uni fun a
+pattern I e t <- TyInst _ e t

plutus-core/plutus-ir/src/PlutusIR/Transform/RewriteRules/UnConstrConstrData.hs

@@ -0,0 +1,59 @@
+-- editorconfig-checker-disable-file
+{-# LANGUAGE GADTs         #-}
+{-# LANGUAGE TypeOperators #-}
+module PlutusIR.Transform.RewriteRules.UnConstrConstrData
+    ( unConstrConstrData
+    ) where
+
+import PlutusCore.Default
+import PlutusCore.Quote
+import PlutusIR
+import PlutusIR.Analysis.Builtins
+import PlutusIR.Analysis.VarInfo
+import PlutusIR.Transform.RewriteRules.Common
+
+{- | This rule rewrites terms of form `BUILTIN(unConstrData(constrData(x,y)))`
+, where builtin stands for `FstPair` or `SndPair`, to "x" or "y" respectively.
+
+This rewrite-rule was originally meant to rewrite `unConstrData(constrData(x,y)) => (x,y)`,
+however we do not have a (polymorphic or monomorphic) builtin constructor to create a `BuiltinPair`
+"(x,y)". See note [Representable built-in functions over polymorphic built-in types].
+
+So we adapted the original rewrite rule to try to achieve a similar goal.
+Unfortunately, the adapted rule is less applicable and will most likely never fire
+(at least for PIR code generated by plutus-tx).
+The reason is that the TH code in plutus-tx does not create such "tight" code, but uses
+way more lets that get in the way preventing the rule from firing.
+
+Possible solutions: Some more aggressive PIR inlining, rewriting the PlutusTx TH code, or
+introducing specialised pattern-matching builtins as last resort.
+Plutus Tx TH code responsible:
+<https://github.com/input-output-hk/plutus/blob/9364099b38e3aa27fb311af3299d2210e7e33e45/plutus-tx/src/PlutusTx/IsData/TH.hs#L174-L192>
+-}
+unConstrConstrData :: (MonadQuote m, t ~ Term tyname Name DefaultUni DefaultFun a, Monoid a)
+                   => BuiltinsInfo DefaultUni DefaultFun
+                   -> VarsInfo tyname Name DefaultUni a
+                   -> t
+                   -> m t
+unConstrConstrData binfo vinfo t = case t of
+    -- This rule might as well have been split into two separate rules, but kept as one
+    -- so as to reuse most of the matching pattern.
+
+    -- builtin({t1}, {t2}, unConstr(constrData(i, data)))
+    (A (I (I (B builtin) tyFst) tySnd)
+     (A (B UnConstrData) (A (A (B ConstrData) arg1) arg2))) ->
+        case builtin of
+            -- sndPair({t1}, {t2}, unConstr(constrData(i, data))) = i `seq` data
+            SndPair -> (tyFst,arg1) `seQ` arg2
+
+            -- fstPair({t1}, {t2}, unConstr(constrData(i, data))) = let !gen = i in data `seq` gen
+            FstPair ->  do
+                  (genVar, genLetIn) <- mkFreshTermLet tyFst arg1
+                  genLetIn <$>
+                      (tySnd, arg2) `seQ` genVar
+            _ -> pure t
+    _ -> pure t
+
+  where
+      infixr 5 `seQ` -- 5 so it has more precedence than <$>
+      seQ = seqP binfo vinfo

plutus-core/plutus-ir/test/PlutusIR/Transform/RewriteRules/Tests.hs

@@ -1,21 +1,37 @@
 module PlutusIR.Transform.RewriteRules.Tests where
 
 import PlutusCore.Quote
+import PlutusCore.Test hiding (ppCatch)
+import PlutusIR.Compiler qualified as PIR
 import PlutusIR.Parser
 import PlutusIR.Test
 import PlutusIR.Transform.RewriteRules as RewriteRules
+import PlutusPrelude
 
-import Data.Default.Class
 import Test.Tasty
-import Test.Tasty.Extras
 
 test_RewriteRules :: TestTree
 test_RewriteRules = runTestNestedIn ["plutus-ir/test/PlutusIR/Transform"] $
     testNested "RewriteRules" $
-    fmap
-        (goldenPir (runQuote . RewriteRules.rewriteWith def)  pTerm)
-        [ "equalsInt" -- this tests that the function works on equalInteger
-        , "divideInt" -- this tests that the function excludes not commutative functions
-        , "multiplyInt" -- this tests that the function works on multiplyInteger
-        , "let" -- this tests that it works in the subterms
+    (fmap
+        (goldenPirDoc (prettyPlcClassicDebug . runQuote . RewriteRules.rewriteWith def)  pTerm)
+        [ "equalsInt.pir" -- this tests that the function works on equalInteger
+        , "divideInt.pir" -- this tests that the function excludes not commutative functions
+        , "multiplyInt.pir" -- this tests that the function works on multiplyInteger
+        , "let.pir" -- this tests that it works in the subterms
+        , "unConstrConstrDataFst.pir"
+        , "unConstrConstrDataSnd.pir"
         ]
+    )
+    ++
+    (fmap
+        (goldenPirEvalTrace pTermAsProg)
+        [ "unConstrConstrDataFst.pir.eval"
+        ]
+    )
+
+  where
+    goldenPirEvalTrace = goldenPirM $ \ast -> ppCatch $ do
+          -- we need traces to remain for checking the evaluation-order
+          tplc <- asIfThrown $ compileWithOpts ( set (PIR.ccOpts . PIR.coPreserveLogging) True) ast
+          runUPlcLogs [void tplc]

plutus-core/plutus-ir/test/PlutusIR/Transform/RewriteRules/let.pir.golden

@@ -0,0 +1,8 @@
+(let
+  (nonrec)
+  (termbind (strict) (vardecl x_0 (con integer)) (error (con integer)))
+  [
+    [ (builtin equalsInteger) (con integer 5) ]
+    [ [ (builtin addInteger) (con integer 10) ] x_0 ]
+  ]
+)