Kwxm/expose optimisers 3 (PLT-1145) (#5189)

* WIP

* WIP

* WIP

* Remove old compilation function

* Remove old compilation function

* Simplify lens code

* Simplify lens code

* Fix Reader problem

* Tidy up some imports

* Remove unused type

* Add changelog entry

* Update code to use SrcSpan

* Fix some prettyprinting exports

* Tidy up types

* Update executables

* Incorporate some changes from abandoned branch

* Merge code from various branches

* Change error name

* Tidying up

* Tidying up

* Tidy up; fix renaming problem

* Tidy up

* Add a Note about why we don't have de Bruijn for PIR

* Forgot file

* Be precise about

* Remove unexpectedly superfluous import

* Fix for broken debugger

* Fix for broken Agda

* Fix for broken Agda

* Address PR comments

* Tidy up; fix problem with fakeTyNameDeBruijn

* Adjust indentation

* Adjust formatting

* Remove superfluous changelog entry

* Update comment and remove empty 'where' clause

plutus-core/changelog.d/20230220_085551_kenneth.mackenzie_expose_optimisers.md

@@ -0,0 +1,9 @@
+### Added
+- `optimise` options for the `pir`, `plc`, and `uplc` commands.
+- A `compile` option for the `pir` command which allows a PIR file to be
+  compiled to PLC or UPLC.
+- Functions for mapping over names and typenames in the PLC AST.
+
+### Changed
+- Some of the `pir` commands have been extended to allow both `flat` and textual
+  input.

plutus-core/executables/plc/Main.hs

@@ -34,6 +34,7 @@ data EraseOptions     = EraseOptions Input Format Output Format PrintMode
 -- Main commands
 data Command = Apply     ApplyOptions
              | Typecheck TypecheckOptions
+             | Optimise  OptimiseOptions
              | Convert   ConvertOptions
              | Print     PrintOptions
              | Example   ExampleOptions
@@ -64,15 +65,15 @@ plutus langHelpText =
       (fullDesc <> header "Typed Plutus Core Tool" <> progDesc langHelpText)
 
 plutusOpts :: Parser Command
-plutusOpts = hsubparser (
+plutusOpts = hsubparser $
        command "apply"
            (info (Apply <$> applyOpts)
             (progDesc $
             "Given a list of input scripts f g1 g2 ... gn, output a script consisting of "
               ++ "(... ((f g1) g2) ... gn); "
               ++ "for example, "
               ++ "'plc apply --if flat Validator.flat Datum.flat Redeemer.flat Context.flat"
-              ++" --of flat -o Script.flat'"))
+              ++" --of flat -o Script.flat'."))
     <> command "print"
            (info (Print <$> printOpts)
             (progDesc "Parse a program then prettyprint it."))
@@ -88,6 +89,9 @@ plutusOpts = hsubparser (
     <> command "typecheck"
            (info (Typecheck <$> typecheckOpts)
             (progDesc "Typecheck a typed Plutus Core program."))
+    <> command "optimise" (optimise $ "Run the PLC optimisation pipeline on the input.  "
+                                        ++ "At present there are no PLC optimisations.")
+    <> command "optimize" (optimise "Same as 'optimise'.")
     <> command "erase"
            (info (Erase <$> eraseOpts)
             (progDesc "Convert a typed Plutus Core program to an untyped one."))
@@ -96,18 +100,19 @@ plutusOpts = hsubparser (
             (progDesc "Evaluate a typed Plutus Core program using the CK machine."))
     <> command "dump-model"
            (info (pure DumpModel)
-            (progDesc "Dump the cost model parameters"))
+            (progDesc "Dump the cost model parameters."))
     <> command "print-builtin-signatures"
            (info (pure PrintBuiltinSignatures)
-            (progDesc "Print the signatures of the built-in functions"))
-  )
+            (progDesc "Print the signatures of the built-in functions."))
+    where optimise desc = info (Optimise <$> optimiseOpts) $ progDesc desc
+
 
 ---------------- Script application ----------------
 
 -- | Apply one script to a list of others.
 runApply :: ApplyOptions -> IO ()
 runApply (ApplyOptions inputfiles ifmt outp ofmt mode) = do
-  scripts <- mapM ((getProgram ifmt ::  Input -> IO (PlcProg PLC.SrcSpan)) . FileInput) inputfiles
+  scripts <- mapM ((readProgram ifmt ::  Input -> IO (PlcProg PLC.SrcSpan)) . FileInput) inputfiles
   let appliedScript =
         case map (\case p -> () <$ p) scripts of
           []          -> errorWithoutStackTrace "No input files"
@@ -118,7 +123,7 @@ runApply (ApplyOptions inputfiles ifmt outp ofmt mode) = do
 
 runTypecheck :: TypecheckOptions -> IO ()
 runTypecheck (TypecheckOptions inp fmt) = do
-  prog <- getProgram fmt inp
+  prog <- readProgram fmt inp
   case PLC.runQuoteT $ do
     tcConfig <- PLC.getDefTypeCheckConfig ()
     PLC.inferTypeOfProgram tcConfig (void prog)
@@ -128,11 +133,20 @@ runTypecheck (TypecheckOptions inp fmt) = do
       Right ty                                               ->
         T.putStrLn (PP.displayPlcDef ty) >> exitSuccess
 
+---------------- Optimisation ----------------
+
+runOptimisations:: OptimiseOptions -> IO ()
+runOptimisations (OptimiseOptions inp ifmt outp ofmt mode) = do
+  prog <- readProgram ifmt inp  :: IO (PlcProg PLC.SrcSpan)
+  let optimised = prog  -- No PLC optimisations at present!
+  writeProgram outp ofmt mode optimised
+
+
 ---------------- Evaluation ----------------
 
 runEval :: EvalOptions -> IO ()
 runEval (EvalOptions inp ifmt printMode timingMode) = do
-  prog <- getProgram ifmt inp
+  prog <- readProgram ifmt inp
   let evaluate = Ck.evaluateCkNoEmit PLC.defaultBuiltinsRuntime
       term = void $ prog ^. PLC.progTerm
       !_ = rnf term
@@ -153,7 +167,7 @@ runPlcPrintExample = runPrintExample getPlcExamples
 -- | Input a program, erase the types, then output it
 runErase :: EraseOptions -> IO ()
 runErase (EraseOptions inp ifmt outp ofmt mode) = do
-  typedProg <- (getProgram ifmt inp :: IO (PlcProg PLC.SrcSpan))
+  typedProg <- (readProgram ifmt inp :: IO (PlcProg PLC.SrcSpan))
   let untypedProg = () <$ PLC.eraseProgram typedProg
   case ofmt of
     Textual       -> writePrettyToFileOrStd outp mode untypedProg
@@ -165,12 +179,13 @@ main :: IO ()
 main = do
     options <- customExecParser (prefs showHelpOnEmpty) plcInfoCommand
     case options of
-        Apply     opts         -> runApply        opts
-        Typecheck opts         -> runTypecheck    opts
-        Eval      opts         -> runEval         opts
-        Example   opts         -> runPlcPrintExample opts
-        Erase     opts         -> runErase        opts
-        Print     opts         -> runPrint        opts
+        Apply     opts         -> runApply            opts
+        Typecheck opts         -> runTypecheck        opts
+        Optimise  opts         -> runOptimisations    opts
+        Eval      opts         -> runEval             opts
+        Example   opts         -> runPlcPrintExample  opts
+        Erase     opts         -> runErase            opts
+        Print     opts         -> runPrint   @PlcProg opts
         Convert   opts         -> runConvert @PlcProg opts
         DumpModel              -> runDumpModel
         PrintBuiltinSignatures -> runPrintBuiltinSignatures

plutus-core/executables/src/PlutusCore/Executable/AstIO.hs

@@ -0,0 +1,184 @@
+-- editorconfig-checker-disable-file
+{-# LANGUAGE LambdaCase       #-}
+{-# LANGUAGE TypeApplications #-}
+
+-- | Reading and writing ASTs with various name types in flat format.
+
+module PlutusCore.Executable.AstIO
+    ( serialisePirProgramFlat
+    , serialisePlcProgramFlat
+    , serialiseUplcProgramFlat
+    , loadPirASTfromFlat
+    , loadPlcASTfromFlat
+    , loadUplcASTfromFlat)
+where
+
+import PlutusCore.Executable.Types
+
+import PlutusCore qualified as PLC
+import PlutusCore.DeBruijn (fakeNameDeBruijn, fakeTyNameDeBruijn, unNameDeBruijn, unNameTyDeBruijn)
+
+import PlutusIR.Core.Instance.Pretty ()
+
+import UntypedPlutusCore qualified as UPLC
+
+import Control.Lens (traverseOf)
+import Control.Monad.Except (runExcept, runExceptT)
+import Data.ByteString.Lazy qualified as BSL
+import Data.Functor ((<&>))
+import Flat (Flat, flat, unflat)
+
+type UplcProgDB ann = UPLC.Program PLC.DeBruijn PLC.DefaultUni PLC.DefaultFun ann
+type UplcProgNDB ann = UPLC.Program PLC.NamedDeBruijn PLC.DefaultUni PLC.DefaultFun ann
+
+type PlcProgDB ann = PLC.Program PLC.TyDeBruijn PLC.DeBruijn PLC.DefaultUni PLC.DefaultFun ann
+type PlcProgNDB ann = PLC.Program PLC.NamedTyDeBruijn PLC.NamedDeBruijn PLC.DefaultUni PLC.DefaultFun ann
+
+-- | PIR does not support names involving de Bruijn indices. We do allow these
+-- formats here to facilitate code sharing, but issue the error below if they're
+-- encountered.  This should never happen in practice because the options
+-- parsers for the `pir` command only accept the Named and Textual formats.
+unsupportedNameTypeError :: AstNameType -> a
+unsupportedNameTypeError nameType = error $ "ASTs with " ++ show nameType ++ " names are not supported for PIR"
+
+---------------- Name conversions ----------------
+
+-- Untyped programs
+
+-- | Convert an untyped program to one where the 'name' type is textual names
+-- with de Bruijn indices.
+toNamedDeBruijnUPLC :: UplcProg ann -> UplcProgNDB ann
+toNamedDeBruijnUPLC prog =
+    case runExcept @PLC.FreeVariableError $ traverseOf UPLC.progTerm UPLC.deBruijnTerm prog of
+        Left e  -> error $ show e
+        Right p -> p
+
+-- | Convert an untyped program to one where the 'name' type is de Bruijn indices.
+toDeBruijnUPLC :: UplcProg ann -> UplcProgDB ann
+toDeBruijnUPLC = UPLC.programMapNames unNameDeBruijn . toNamedDeBruijnUPLC
+
+
+-- | Convert an untyped program with named de Bruijn indices to one with textual names.
+fromNamedDeBruijnUPLC :: UplcProgNDB ann -> UplcProg ann
+fromNamedDeBruijnUPLC prog =
+    case PLC.runQuote $
+         runExceptT @PLC.FreeVariableError $ traverseOf UPLC.progTerm UPLC.unDeBruijnTerm prog of
+      Left e  -> error $ show e
+      Right p -> p
+
+-- | Convert an untyped program with de Bruijn indices to one with textual names.
+fromDeBruijnUPLC :: UplcProgDB ann -> UplcProg ann
+fromDeBruijnUPLC = fromNamedDeBruijnUPLC . UPLC.programMapNames fakeNameDeBruijn
+
+-- Typed programs
+
+-- | Convert a typed program to one where the 'name' type is textual names with
+-- de Bruijn indices.
+toNamedDeBruijnPLC :: PlcProg ann -> PlcProgNDB ann
+toNamedDeBruijnPLC prog =
+    case runExcept @PLC.FreeVariableError $ traverseOf PLC.progTerm PLC.deBruijnTerm prog of
+        Left e  -> error $ show e
+        Right p -> p
+
+-- | Convert a typed program to one where the 'name' type is de Bruijn indices.
+toDeBruijnPLC :: PlcProg ann -> PlcProgDB ann
+toDeBruijnPLC = PLC.programMapNames unNameTyDeBruijn unNameDeBruijn . toNamedDeBruijnPLC
+
+
+-- | Convert a typed program with named de Bruijn indices to one with textual names.
+fromNamedDeBruijnPLC :: PlcProgNDB ann -> PlcProg ann
+fromNamedDeBruijnPLC prog = do
+  case PLC.runQuote $
+       runExceptT @PLC.FreeVariableError $ traverseOf PLC.progTerm PLC.unDeBruijnTerm prog of
+    Left e  -> error $ show e
+    Right p -> p
+
+-- | Convert a typed program with de Bruijn indices to one with textual names.
+fromDeBruijnPLC :: PlcProgDB ann -> PlcProg ann
+fromDeBruijnPLC = fromNamedDeBruijnPLC . PLC.programMapNames fakeTyNameDeBruijn fakeNameDeBruijn
+
+-- Flat serialisation in various formats.
+
+serialisePirProgramFlat
+    :: Flat ann
+    => AstNameType
+    -> PirProg ann
+    -> BSL.ByteString
+serialisePirProgramFlat =
+    \case
+      Named         -> BSL.fromStrict . flat
+      DeBruijn      -> unsupportedNameTypeError DeBruijn
+      NamedDeBruijn -> unsupportedNameTypeError NamedDeBruijn
+
+serialisePlcProgramFlat
+    :: Flat ann
+    => AstNameType
+    -> PlcProg ann
+    -> BSL.ByteString
+serialisePlcProgramFlat =
+    \case
+     Named         -> BSL.fromStrict . flat
+     DeBruijn      -> BSL.fromStrict . flat . toDeBruijnPLC
+     NamedDeBruijn -> BSL.fromStrict . flat . toNamedDeBruijnPLC
+
+serialiseUplcProgramFlat
+    :: Flat ann
+    => AstNameType
+    -> UplcProg ann
+    -> BSL.ByteString
+serialiseUplcProgramFlat =
+    \case
+     Named         -> BSL.fromStrict . flat. UPLC.UnrestrictedProgram
+     DeBruijn      -> BSL.fromStrict . flat. UPLC.UnrestrictedProgram . toDeBruijnUPLC
+     NamedDeBruijn -> BSL.fromStrict . flat .UPLC.UnrestrictedProgram . toNamedDeBruijnUPLC
+
+-- Deserialising ASTs from Flat
+
+-- Read a binary-encoded file (eg, Flat-encoded PLC)
+getBinaryInput :: Input -> IO BSL.ByteString
+getBinaryInput StdInput         = BSL.getContents
+getBinaryInput (FileInput file) = BSL.readFile file
+
+unflatOrFail :: Flat a => BSL.ByteString -> a
+unflatOrFail input =
+    case unflat input of
+     Left e  -> error $ "Flat deserialisation failure: " ++ show e
+     Right r -> r
+
+loadPirASTfromFlat
+    :: Flat a
+    => AstNameType
+    -> Input
+    -> IO (PirProg a)
+loadPirASTfromFlat flatMode inp =
+    getBinaryInput inp <&>
+    case flatMode of
+      Named -> unflatOrFail
+      _     -> unsupportedNameTypeError flatMode
+
+-- | Read and deserialise a Flat-encoded PIR/PLC AST
+loadPlcASTfromFlat
+    :: Flat a
+    => AstNameType
+    -> Input
+    -> IO (PlcProg a)
+loadPlcASTfromFlat flatMode inp =
+    getBinaryInput inp <&>
+    case flatMode of
+      Named         -> unflatOrFail
+      DeBruijn      -> unflatOrFail <&> fromDeBruijnPLC
+      NamedDeBruijn -> unflatOrFail <&> fromNamedDeBruijnPLC
+
+-- | Read and deserialise a Flat-encoded UPLC AST
+loadUplcASTfromFlat
+    :: Flat ann
+    => AstNameType
+    -> Input
+    -> IO (UplcProg ann)
+loadUplcASTfromFlat flatMode inp =
+    getBinaryInput inp <&>
+    case flatMode of
+      Named         -> unflatOrFail <&> UPLC.unUnrestrictedProgram
+      DeBruijn      -> unflatOrFail <&> UPLC.unUnrestrictedProgram <&> fromDeBruijnUPLC
+      NamedDeBruijn -> unflatOrFail <&> UPLC.unUnrestrictedProgram <&> fromNamedDeBruijnUPLC
+