Implement missing bits of PlutusIR API

plutus-core/plutus-core.cabal

@@ -110,6 +110,7 @@ library
         PlutusIR.Parser
         PlutusIR.MkPir
         PlutusIR.Purity
+        PlutusIR.Subst
         PlutusIR.Transform.DeadCode
         PlutusIR.Transform.Substitute
         PlutusIR.Transform.ThunkRecursions

plutus-core/plutus-ir/src/PlutusIR/Core/Plated.hs

@@ -1,23 +1,37 @@
 {-# LANGUAGE LambdaCase #-}
-module PlutusIR.Core.Plated (
-    termSubterms,
-    termSubtypes,
-    termBindings,
-    typeSubtypes,
-    datatypeSubtypes,
-    bindingSubterms,
-    bindingSubtypes,
-    bindingIds
+{-# LANGUAGE RankNTypes #-}
+module PlutusIR.Core.Plated
+    ( termSubterms
+    , termSubtermsDeep
+    , termSubtypes
+    , termSubtypesDeep
+    , termBindings
+    , typeSubtypes
+    , typeSubtypesDeep
+    , typeUniques
+    , typeUniquesDeep
+    , datatypeSubtypes
+    , bindingSubterms
+    , bindingSubtypes
+    , bindingIds
+    , termUniques
+    , termUniquesDeep
     ) where
 
-import           PlutusCore         (typeSubtypes)
-import qualified PlutusCore         as PLC
-import           PlutusCore.Flat    ()
-import qualified PlutusCore.Name    as PLC
+import qualified PlutusCore             as PLC
+import           PlutusCore.Core.Plated (typeSubtypes, typeSubtypesDeep, typeUniques, typeUniquesDeep)
+import           PlutusCore.Flat        ()
+import qualified PlutusCore.Name        as PLC
 
 import           PlutusIR.Core.Type
 
-import           Control.Lens       hiding (Strict)
+import           Control.Lens           hiding (Strict)
+
+infixr 6 <^>
+
+-- | Compose two folds to make them run in parallel. The results are concatenated.
+(<^>) :: Fold s a -> Fold s a -> Fold s a
+(f1 <^> f2) g s = f1 g s *> f2 g s
 
 {-# INLINE bindingSubterms #-}
 -- | Get all the direct child 'Term's of the given 'Binding'.
@@ -82,6 +96,10 @@ termSubterms f = \case
     c@Constant {}     -> pure c
     b@Builtin {}      -> pure b
 
+-- | Get all the transitive child 'Term's of the given 'Term'.
+termSubtermsDeep :: Fold (Term tyname name uni fun ann) (Term tyname name uni fun ann)
+termSubtermsDeep = cosmosOf termSubterms
+
 {-# INLINE termSubtypes #-}
 -- | Get all the direct child 'Type's of the given 'Term', including those within 'Binding's.
 termSubtypes :: Traversal' (Term tyname name uni fun a) (Type tyname uni a)
@@ -98,9 +116,37 @@ termSubtypes f = \case
     c@Constant {}     -> pure c
     b@Builtin {}      -> pure b
 
+-- | Get all the transitive child 'Type's of the given 'Term'.
+termSubtypesDeep :: Fold (Term tyname name uni fun ann) (Type tyname uni ann)
+termSubtypesDeep = termSubtermsDeep . termSubtypes . typeSubtypesDeep
+
 {-# INLINE termBindings #-}
 -- | Get all the direct child 'Binding's of the given 'Term'.
 termBindings :: Traversal' (Term tyname name uni fun a) (Binding tyname name uni fun a)
 termBindings f = \case
     Let x r bs t -> Let x r <$> traverse f bs <*> pure t
     t            -> pure t
+
+
+-- | Get all the direct child 'Unique's of the given 'Term' (including the type-level ones).
+termUniques
+    :: PLC.HasUniques (Term tyname name uni fun ann)
+    => Traversal' (Term tyname name uni fun ann) PLC.Unique
+termUniques f = \case
+    Let ann r bs t    -> Let ann r <$> (traverse . bindingIds) f bs <*> pure t
+    Var ann n         -> PLC.theUnique f n  <&> Var ann
+    TyAbs ann tn k t  -> PLC.theUnique f tn <&> \tn' -> TyAbs ann tn' k t
+    LamAbs ann n ty t -> PLC.theUnique f n  <&> \n'  -> LamAbs ann n' ty t
+    a@Apply{}         -> pure a
+    c@Constant{}      -> pure c
+    b@Builtin{}       -> pure b
+    t@TyInst{}        -> pure t
+    e@Error{}         -> pure e
+    i@IWrap{}         -> pure i
+    u@Unwrap{}        -> pure u
+
+-- | Get all the transitive child 'Unique's of the given 'Term' (including the type-level ones).
+termUniquesDeep
+    :: PLC.HasUniques (Term tyname name uni fun ann)
+    => Fold (Term tyname name uni fun ann) PLC.Unique
+termUniquesDeep = termSubtermsDeep . (termSubtypes . typeUniquesDeep <^> termUniques)

plutus-core/plutus-ir/src/PlutusIR/Core/Type.hs

@@ -29,6 +29,7 @@ import           PlutusCore.Constant (AsConstant (..), FromConstant (..))
 import           PlutusCore.Core     (UniOf)
 import           PlutusCore.Flat     ()
 import           PlutusCore.MkPlc    (Def (..), TermLike (..), TyVarDecl (..), VarDecl (..))
+import qualified PlutusCore.Name     as PLC
 
 
 import qualified Data.Text           as T
@@ -185,3 +186,6 @@ instance ( PLC.Closed uni
          , Flat tyname
          , Flat name
          ) => Flat (Program tyname name uni fun a)
+
+type instance PLC.HasUniques (Term tyname name uni fun ann) = (PLC.HasUnique tyname PLC.TypeUnique, PLC.HasUnique name PLC.TermUnique)
+type instance PLC.HasUniques (Program tyname name uni fun ann) = PLC.HasUniques (Term tyname name uni fun ann)

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

@@ -4,20 +4,32 @@ module PlutusIR.Mark
     , markNonFreshProgram
     ) where
 
-import qualified PlutusCore.Core           as PLC
+import qualified PlutusCore.Core  as PLC
+import qualified PlutusCore.Name  as PLC
 
-import           PlutusCore.Name
 import           PlutusCore.Quote
 
-import           PlutusIR.Transform.Rename
-
 import           PlutusIR.Core
+import           PlutusIR.Subst
 
+-- | Marks all the 'Unique's in a term as used, so they will not be generated in future. Useful if you
+-- have a term which was not generated in 'Quote'.
 markNonFreshTerm
     :: (PLC.HasUniques (Term tyname name uni fun ann), MonadQuote m)
     => Term tyname name uni fun ann -> m ()
-markNonFreshTerm = markNonFreshMax . undefined
+markNonFreshTerm = markNonFreshMax . uniquesTerm
 
-markNonFreshType = undefined
+-- | Marks all the 'Unique's in a type as used, so they will not be generated in future. Useful if you
+-- have a type which was not generated in 'Quote'.
+markNonFreshType
+    :: (PLC.HasUniques (Type tyname uni ann), MonadQuote m)
+    => Type tyname uni ann -> m ()
+markNonFreshType = markNonFreshMax . uniquesType
 
-markNonFreshProgram = undefined
+-- | Marks all the 'Unique's in a program as used, so they will not be generated in future. Useful if you
+-- have a program which was not generated in 'Quote'.
+markNonFreshProgram
+    :: (PLC.HasUniques (Program tyname name uni fun ann), MonadQuote m)
+    => Program tyname name uni fun ann
+    -> m ()
+markNonFreshProgram (Program _ body) = markNonFreshTerm body

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

@@ -0,0 +1,21 @@
+module PlutusIR.Subst
+    ( uniquesTerm
+    , uniquesType
+    ) where
+
+import qualified PlutusCore.Core.Type as PLC
+import qualified PlutusCore.Name      as PLC
+import           PlutusCore.Subst     (uniquesType)
+
+import           PlutusIR.Core
+
+import           Control.Lens
+import           Data.Set             as Set
+
+setOf :: Getting (Set a) s a -> s -> Set a
+setOf g = foldMapOf g singleton
+
+uniquesTerm
+    :: PLC.HasUniques (Term tyname name uni fun ann)
+    => Term tyname name uni fun ann -> Set PLC.Unique
+uniquesTerm = setOf termUniquesDeep

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

@@ -57,7 +57,7 @@ Two problems arise:
    'PLC.ScopedRenameM' is for performing the renaming (the second stage).
 -}
 
-type instance PLC.HasUniques (Term tyname name uni fun ann) = PLC.HasUniques (PLC.Term tyname name uni fun ann)
+-- type instance PLC.HasUniques (Term tyname name uni fun ann) = PLC.HasUniques (PLC.Term tyname name uni fun ann)
 
 instance PLC.HasUniques (Term tyname name uni fun ann) => PLC.Rename (Term tyname name uni fun ann) where
     -- TODO: the Plutus Core codebase uses marking in order to prevent clashing with existing