Add custom SCCs to PlutusTx/Plugin.hs to profiled compilation stages

.gitignore

@@ -49,6 +49,8 @@ stack.yaml.lock
 *.chi
 *.chs.h
 *.prof
+*.hp
+*.ps
 .liquid/
 
 # Agda

nix/outputs.nix

@@ -32,8 +32,8 @@ let
   project = import ./project.nix
     { inherit inputs pkgs lib metatheory r-with-packages utils; };
 
-  mkShell = ghc: import ./shell.nix
-    { inherit inputs pkgs lib project agda-tools metatheory r-with-packages ghc; };
+  mkShell = project-variant: import ./shell.nix
+    { inherit inputs pkgs lib project agda-tools metatheory r-with-packages project-variant; };
 
   exposed-haskell-packages = {
     plutus-core-test = project.flake'.packages."plutus-core:test:plutus-core-test";
@@ -113,7 +113,7 @@ let
 
   devShells =
     (non-profiled-shells) //
-    { profiled = mkShell project.projectVariants.ghc96-profiled; } //
+    { profiled = mkShell "ghc96-profiled"; } //
     { metatheory-jailbreak = metatheory-jailbreak-shell; };
 
   full-nested-ci-jobs = {

nix/shell.nix

@@ -1,27 +1,30 @@
 # editorconfig-checker-disable-file
 
-{ inputs, pkgs, lib, project, agda-tools, metatheory, r-with-packages, ghc }:
+{ inputs, pkgs, lib, project, agda-tools, metatheory, r-with-packages, project-variant }:
 
 let
 
   # Toolchain versions used in dev shells. Consumed by `project.shellFor`.
-  all-tools = {
+  tools = rec {
     "ghc96".cabal = project.projectVariants.ghc96.tool "cabal" "3.12.1.0";
     "ghc96".cabal-fmt = project.projectVariants.ghc96.tool "cabal-fmt" "latest";
     "ghc96".haskell-language-server = project.projectVariants.ghc96.tool "haskell-language-server" "latest";
     "ghc96".stylish-haskell = project.projectVariants.ghc96.tool "stylish-haskell" "latest";
     "ghc96".fourmolu = project.projectVariants.ghc96.tool "fourmolu" "0.17.0.0"; # fourmolu 0.18.0.0 and hlint 3.10 require GHC >=9.8
     "ghc96".hlint = project.projectVariants.ghc96.tool "hlint" "3.8";
+    "ghc96".hp2ps = project.projectVariants.ghc96.tool "hp2ps" "latest";
 
     "ghc912".cabal = project.projectVariants.ghc912.tool "cabal" "latest";
     "ghc912".cabal-fmt = project.projectVariants.ghc96.tool "cabal-fmt" "latest"; # cabal-fmt not buildable with ghc9122
     "ghc912".haskell-language-server = project.projectVariants.ghc912.tool "haskell-language-server" "latest";
     "ghc912".stylish-haskell = project.projectVariants.ghc912.tool "stylish-haskell" "latest";
     "ghc912".fourmolu = project.projectVariants.ghc912.tool "fourmolu" "latest";
     "ghc912".hlint = project.projectVariants.ghc912.tool "hlint" "latest";
-  };
+    "ghc912".hp2ps = project.projectVariants.ghc912.tool "hp2ps" "latest";
+
+    "ghc96-profiled" = tools."ghc96";
 
-  tools = all-tools.${ghc};
+  }.${project-variant};
 
   # Pre-commit hooks for the repo. Injects into shell via shellHook.
   pre-commit-check = inputs.pre-commit-hooks.lib.${pkgs.system}.run {
@@ -129,8 +132,8 @@ let
       "export LOCALE_ARCHIVE=${pkgs.glibcLocales}/lib/locale/locale-archive";
 
   # Full developer shell with many tools.
-  full-shell = project.projectVariants.${ghc}.shellFor {
-    name = "plutus-shell-${ghc}";
+  full-shell = project.projectVariants.${project-variant}.shellFor {
+    name = "plutus-shell-${project-variant}";
 
     buildInputs = lib.concatLists [
       common-pkgs
@@ -151,8 +154,8 @@ let
 
 
   # Lightweight shell with minimal tools.
-  quick-shell = project.projectVariants.${ghc}.shellFor {
-    name = "plutus-shell-${ghc}";
+  quick-shell = project.projectVariants.${project-variant}.shellFor {
+    name = "plutus-shell-${project-variant}";
     tools = { cabal = "latest"; };
     shellHook = ''
       ${locale-archive-hook}
@@ -165,10 +168,11 @@ let
   # Select shell by compiler used in the project variant.
   shell = {
     ghc96 = full-shell;
+    ghc96-profiled = full-shell;
     ghc98 = quick-shell;
     ghc910 = quick-shell;
     ghc912 = full-shell;
-  }.${ghc};
+  }.${project-variant};
 
 in
 

plutus-tx-plugin/plutus-tx-plugin.cabal

@@ -57,6 +57,7 @@ library
     PlutusTx.Compiler.Error
     PlutusTx.Options
     PlutusTx.Plugin
+    PlutusTx.Compiler.Types
 
   other-modules:
     PlutusTx.Compiler.Binders
@@ -67,7 +68,6 @@ library
     PlutusTx.Compiler.Names
     PlutusTx.Compiler.Trace
     PlutusTx.Compiler.Type
-    PlutusTx.Compiler.Types
     PlutusTx.Compiler.Utils
     PlutusTx.PIRTypes
     PlutusTx.PLCTypes
@@ -241,3 +241,28 @@ test-suite size
   ghc-options:
     -fno-strictness -fno-unbox-strict-fields
     -fno-unbox-small-strict-fields -fno-full-laziness
+
+
+test-suite plutus-tx-plugin-profile-test
+  import:             lang, ghc-version-support, os-support
+  type:               exitcode-stdio-1.0
+  main-is:            ProfileTest.hs
+  hs-source-dirs:     test/Plugin/Profiling
+  build-depends:
+    , base                   >=4.9   && <5
+    , containers
+    , data-default
+    , ghc
+    , ghc-paths
+    , mtl
+    , plutus-core            ^>=1.55
+    , plutus-core:plutus-ir
+    , plutus-tx-plugin       ^>=1.55
+
+  default-extensions: Strict
+  ghc-options:        -threaded -rtsopts -with-rtsopts=-N
+
+  -- See Note [-fno-full-laziness in Plutus Tx]
+  ghc-options:
+    -fno-strictness -fno-unbox-strict-fields
+    -fno-unbox-small-strict-fields -fno-full-laziness

plutus-tx-plugin/src/PlutusTx/Plugin.hs

@@ -10,7 +10,7 @@
 -- For some reason this module is very slow to compile otherwise
 {-# OPTIONS_GHC -O0 #-}
 
-module PlutusTx.Plugin (plugin, plc) where
+module PlutusTx.Plugin (plugin, plc, runCompiler) where
 
 import PlutusPrelude
 import PlutusTx.AsData.Internal qualified
@@ -610,23 +610,31 @@ runCompiler moduleName opts expr = do
             (opts ^. posPreserveLogging)
 
   -- GHC.Core -> Pir translation.
-  pirT <- original <$> (PIR.runDefT annMayInline $ compileExprWithDefs expr)
+  pirT <-
+    {-# SCC "plinth-plugin-core-to-pir-step" #-}
+    original <$> (PIR.runDefT annMayInline $ compileExprWithDefs expr)
+
   let pirP = PIR.Program noProvenance plcVersion pirT
   when (opts ^. posDumpPir) . liftIO $
     dumpFlat (void pirP) "initial PIR program" (moduleName ++ "_initial.pir-flat")
 
   -- Pir -> (Simplified) Pir pass. We can then dump/store a more legible PIR program.
   spirP <-
+    {-# SCC "plinth-plugin-pir-to-simp-step" #-}
     flip runReaderT pirCtx $
       modifyError (NoContext . PIRError) $
         PIR.compileToReadable pirP
+
   when (opts ^. posDumpPir) . liftIO $
       dumpFlat (void spirP) "simplified PIR program" (moduleName ++ "_simplified.pir-flat")
 
   -- (Simplified) Pir -> Plc translation.
-  plcP <- flip runReaderT pirCtx $
-    modifyError (NoContext . PIRError) $
-      PIR.compileReadableToPlc spirP
+  plcP <-
+    {-# SCC "plinth-plugin-simp-to-plc-step" #-}
+    flip runReaderT pirCtx $
+      modifyError (NoContext . PIRError) $
+        PIR.compileReadableToPlc spirP
+
   when (opts ^. posDumpPlc) . liftIO $
       dumpFlat (void plcP) "typed PLC program" (moduleName ++ ".tplc-flat")
 
@@ -636,7 +644,11 @@ runCompiler moduleName opts expr = do
         modifyError PLC.TypeErrorE $
           PLC.inferTypeOfProgram plcTcConfig (plcP $> annMayInline)
 
-  (uplcP, _) <- flip runReaderT plcOpts $ PLC.compileProgramWithTrace plcP
+  (uplcP, _) <-
+    {-# SCC "plinth-plugin-plc-to-uplc-step" #-}
+    flip runReaderT plcOpts $ PLC.compileProgramWithTrace plcP
+
+
   dbP <- liftExcept $ modifyError PLC.FreeVariableErrorE $ traverseOf UPLC.progTerm UPLC.deBruijnTerm uplcP
   when (opts ^. posDumpUPlc) . liftIO $
       dumpFlat

plutus-tx-plugin/test/Plugin/Profiling/ProfileTest.hs

@@ -0,0 +1,127 @@
+{-# LANGUAGE OverloadedStrings   #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TemplateHaskell     #-}
+{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}
+
+-- | Standalone executable for profiling the plugin compilation functions.
+--
+-- This test uses the exposed 'runCompiler' from Plugin.hs to compile a simple
+-- Core expression. This allows profiling the plugin code at runtime, including
+-- the SCC annotations in Plugin.hs.
+--
+-- To run with profiling:
+--   cabal build plutus-tx-plugin-profile-test --enable-profiling
+--   cabal run plutus-tx-plugin-profile-test --enable-profiling -- +RTS -p -hc
+module Main where
+
+import Data.Default
+import Data.Foldable (fold)
+import PlutusCore qualified as PLC
+import PlutusCore.Quote
+import PlutusCore.Version qualified as PLC
+import PlutusIR.Compiler qualified as PIR
+import PlutusIR.Compiler.Types qualified as PIR
+import PlutusIR.Transform.RewriteRules
+import PlutusIR.Transform.RewriteRules.RemoveTrace (rewriteRuleRemoveTrace)
+import PlutusTx.Compiler.Types
+import PlutusTx.Options (PluginOptions (..), defaultPluginOptions)
+import PlutusTx.Plugin (runCompiler)
+
+import GHC qualified as GHC
+import GHC.Core.FamInstEnv qualified as GHC
+import GHC.Core.Opt.OccurAnal qualified as GHC
+import GHC.Driver.Session qualified as GHC
+import GHC.Paths as GHC
+import GHC.Plugins qualified as GHC
+
+import Control.Monad.Except
+import Control.Monad.Reader
+import Control.Monad.State
+import Control.Monad.Writer
+
+import Data.Map qualified as Map
+
+-- | Create a simple Core expression for testing (a literal integer)
+createSimpleCoreExpr :: GHC.DynFlags -> GHC.CoreExpr
+createSimpleCoreExpr _flags =
+  let lit = GHC.Lit (GHC.LitNumber GHC.LitNumInt 42)
+  in lit
+
+-- | Set up a minimal CompileContext for testing
+setupCompileContext
+  :: GHC.DynFlags
+  -> GHC.FamInstEnvs
+  -> NameInfo
+  -> CompileContext PLC.DefaultUni PLC.DefaultFun
+setupCompileContext flags famEnvs nameInfo =
+  let opts = defaultPluginOptions
+      coverage = CoverageOpts mempty
+   in CompileContext
+        { ccOpts =
+            CompileOptions
+              { coProfile = _posProfile opts
+              , coCoverage = coverage
+              , coDatatypeStyle =
+                  if _posPlcTargetVersion opts < PLC.plcVersion110
+                    then PIR.ScottEncoding
+                    else PIR._dcoStyle $ _posDatatypes opts
+              , coRemoveTrace = _posRemoveTrace opts
+              , coInlineFix = _posInlineFix opts
+              }
+        , ccFlags = flags
+        , ccFamInstEnvs = famEnvs
+        , ccNameInfo = nameInfo
+        , ccScope = initialScope
+        , ccBlackholed = mempty
+        , ccCurDef = Nothing
+        , ccModBreaks = Nothing
+        , ccBuiltinsInfo = def
+        , ccBuiltinCostModel = def
+        , ccDebugTraceOn = _posDumpCompilationTrace opts
+        , ccRewriteRules = makeRewriteRules opts
+        , ccSafeToInline = False
+        }
+ where
+  makeRewriteRules :: PluginOptions -> RewriteRules PLC.DefaultUni PLC.DefaultFun
+  makeRewriteRules options =
+    fold
+      [ mwhen (_posRemoveTrace options) rewriteRuleRemoveTrace
+      , defaultUniRewriteRules
+      ]
+  mwhen :: Monoid m => Bool -> m -> m
+  mwhen b m = if b then m else mempty
+
+-- | Create empty NameInfo (simplified - in real usage would need proper lookups)
+createEmptyNameInfo :: NameInfo
+createEmptyNameInfo = Map.empty
+
+main :: IO ()
+main = do
+  putStrLn "Setting up for plugin profiling test..."
+
+  -- Use GHC's API to get DynFlags
+  GHC.defaultErrorHandler GHC.defaultFatalMessager GHC.defaultFlushOut $ do
+    -- Initialize GHC session to get DynFlags
+    GHC.runGhc (Just GHC.libdir) $ do
+      -- Get DynFlags
+      flags <- GHC.getSessionDynFlags
+
+      -- Create a simple Core expression (literal integer)
+      let expr = createSimpleCoreExpr flags
+
+      -- Set up minimal context
+      let famEnvs = (GHC.emptyFamInstEnv, GHC.emptyFamInstEnv)
+          nameInfo = createEmptyNameInfo
+          ctx = setupCompileContext flags famEnvs nameInfo
+          opts = defaultPluginOptions
+          st = CompileState 0 mempty
+          moduleNameStr = "ProfileTest"
+          -- Apply occurrence analysis like the plugin does
+          expr' = GHC.occurAnalyseExpr expr
+
+      -- Call runCompiler - this is where the SCC annotations are!
+      _ <-
+        runExceptT . runWriterT . runQuoteT . flip runReaderT ctx . flip evalStateT st $
+          runCompiler moduleNameStr opts expr'
+
+      pure ()