Add binding anchors; make profile-all use srcspan from anchors

plutus-tx-plugin/src/PlutusTx/Compiler/Expr.hs

@@ -620,6 +620,7 @@ hoistExpr
 hoistExpr var t = do
   wrapUnsafeDataAsConstrName <-
     lookupGhcName 'PlutusTx.AsData.Internal.wrapUnsafeDataAsConstr
+  anchorName <- lookupGhcName 'PlutusTx.Plugin.Utils.anchor
   let name = GHC.getName var
       lexName = LexName name
 
@@ -638,10 +639,15 @@ hoistExpr var t = do
   -- See Note [Dependency tracking]
   modifyCurDeps (Set.insert lexName)
   maybeDef <- PIR.lookupTerm lexName
-  let varSpan = getVarSourceSpan var
+  -- Source spans only come from @anchor@ wrappers injected by Plinth.Plugin.
+  -- @GHC.nameSrcSpan@ on imported Vars is unreliable because it depends on the
+  -- defining module being recompiled from source (issue #7203/#7722), so it is
+  -- deliberately not consulted — we prefer a missing span over a flaky one.
+  let varSpan = findAnchorLoc anchorName t
       addSpan = case varSpan of
         Nothing -> id
         Just src -> fmap . fmap . addSrcSpan $ src ^. srcSpanIso
+      nameStr = GHC.occNameString $ GHC.occName $ GHC.varName var
   case maybeDef of
     Just term -> pure term
     -- See Note [Dependency tracking]
@@ -653,27 +659,27 @@ hoistExpr var t = do
         (PIR.Def var' (PIR.mkVar var', PIR.Strict))
         mempty
 
-      t' <- maybeProfileRhs var var' =<< addSpan (compileExpr Nothing t)
+      t' <- maybeProfileRhs varSpan nameStr var' =<< addSpan (compileExpr Nothing t)
       -- See Note [Non-strict let-bindings]
       PIR.modifyTermDef lexName (const $ PIR.Def var' (t', PIR.NonStrict))
       pure $ PIR.mkVar var'
 
--- 'GHC.Var' in argument is only for extracting srcspan and accurate name.
+{-| Wrap the RHS of a definition with entry/exit profile traces when profiling
+is enabled and the definition is a function/abstraction. The span is
+resolved by the caller (see 'hoistExpr'), preferring an @anchor@-carried
+span so it still works with GHC's interface cache. -}
 maybeProfileRhs
   :: CompilingDefault uni fun m ann
-  => GHC.Var
+  => Maybe GHC.RealSrcSpan
+  -> String
   -> PLCVar uni
   -> PIRTerm uni fun
   -> m (PIRTerm uni fun)
-maybeProfileRhs ghcVar var t = do
+maybeProfileRhs mSpan nameStr var t = do
   CompileContext {ccOpts = compileOpts} <- ask
   let
-    nameStr = GHC.occNameString $ GHC.occName $ GHC.varName $ ghcVar
     displayName = T.pack $
-      case getVarSourceSpan ghcVar of
-        -- When module is not compiled and GHC is using cached build from previous build, it will
-        -- lack source span. There's nothing much we can do about this here since this is GHC
-        -- behavior. Issue #7203
+      case mSpan of
         Nothing -> nameStr
         Just src -> nameStr <> " (" <> show (src ^. srcSpanIso) <> ")"
 
@@ -1331,7 +1337,9 @@ compileExpr mloc e = do
             _ -> compileTypeNorm $ GHC.varType b
           -- See Note [Non-strict let-bindings]
           withVarTyScoped b ty $ \v -> do
-            rhs'' <- maybeProfileRhs b v rhs'
+            let bSpan = findAnchorLoc anchorName rhs
+                bName = GHC.occNameString (GHC.occName (GHC.varName b))
+            rhs'' <- maybeProfileRhs bSpan bName v rhs'
             let binds = pure $ PIR.TermBind annMayInline PIR.NonStrict v rhs''
             body' <- compileExpr Nothing body
             pure $ PIR.Let annMayInline PIR.NonRec binds body'
@@ -1340,7 +1348,9 @@ compileExpr mloc e = do
             -- the bindings are scope in both the body and the args
             -- TODO: this is a bit inelegant matching the vars back up
             binds <- for (zip vars bs) $ \(v, (ghcVar, rhs)) -> do
-              rhs' <- maybeProfileRhs ghcVar v =<< compileExpr Nothing rhs
+              let gvSpan = findAnchorLoc anchorName rhs
+                  gvName = GHC.occNameString (GHC.occName (GHC.varName ghcVar))
+              rhs' <- maybeProfileRhs gvSpan gvName v =<< compileExpr Nothing rhs
               -- See Note [Non-strict let-bindings]
               pure $ PIR.TermBind annMayInline PIR.NonStrict v rhs'
             body' <- compileExpr Nothing body
@@ -1535,9 +1545,6 @@ getSourceSpan mmb GHC.HpcTick {GHC.tickId = tid} = do
   return sp
 getSourceSpan _ _ = Nothing
 
-getVarSourceSpan :: GHC.Var -> Maybe GHC.RealSrcSpan
-getVarSourceSpan = GHC.srcSpanToRealSrcSpan . GHC.nameSrcSpan . GHC.varName
-
 srcSpanIso :: Iso' GHC.RealSrcSpan SrcSpan
 srcSpanIso = iso fromGHC toGHC
   where
@@ -1729,6 +1736,32 @@ extractLoc anchorName modBreaks = go Nothing
             go (Just ss) e
       other -> (acc, other)
 
+{-| Find the first @anchor@-carried 'RealSrcSpan' anywhere inside a 'CoreExpr',
+descending through applications, lambdas, lets, cases, casts, and ticks.
+Useful when the immediate head of the expression is not an anchor call
+(e.g. a hoisted top-level RHS that starts with a lambda). -}
+findAnchorLoc :: GHC.Name -> GHC.CoreExpr -> Maybe GHC.RealSrcSpan
+findAnchorLoc anchorName = go
+  where
+    firstJust = foldr ((<|>) . go) Nothing
+    go = \case
+      GHC.App
+        ( GHC.App
+            (GHC.App (GHC.Var f) (GHC.Type (GHC.LitTy (GHC.StrTyLit loc))))
+            (GHC.Type _eTy)
+          )
+        _
+          | GHC.getName f == anchorName -> decodeSrcSpan (GHC.unpackFS loc)
+      GHC.App f x -> go f <|> go x
+      GHC.Lam _ b -> go b
+      GHC.Let (GHC.NonRec _ rhs) e -> go rhs <|> go e
+      GHC.Let (GHC.Rec bs) e -> firstJust (map snd bs) <|> go e
+      GHC.Case scrut _ _ alts ->
+        go scrut <|> firstJust [rhs | GHC.Alt _ _ rhs <- alts]
+      GHC.Cast e _ -> go e
+      GHC.Tick _ e -> go e
+      _ -> Nothing
+
 extractUnsupported
   :: GHC.Name
   -> GHC.CoreExpr

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

@@ -176,7 +176,10 @@ injectAnchors env = do
   let binds = GHC.tcg_binds env
       bindsAnchored =
         Compat.modifyBinds
-          (transformBi (stripGuardAnchors anchorId) . transformBi (anchorExpr anchorId))
+          ( transformBi (stripGuardAnchors anchorId)
+              . transformBi (anchorBinding anchorId)
+              . transformBi (anchorExpr anchorId)
+          )
           binds
   pure env {GHC.tcg_binds = bindsAnchored}
 
@@ -185,13 +188,64 @@ anchorExpr :: GHC.Id -> GHC.LHsExpr GHC.GhcTc -> GHC.LHsExpr GHC.GhcTc
 anchorExpr anchorId le@(GHC.L ann e)
   | isAnchorWorthy anchorId e
   , Just !sp <- GHC.srcSpanToRealSrcSpan (GHC.locA ann) =
+      wrapWithAnchor anchorId sp le
+  | otherwise = le
+
+{-| Wrap an 'LHsExpr' with an @anchor@ carrying the given source span. Skips
+expressions whose type might be unlifted (since @anchor@'s type variable has
+kind @Type@). -}
+wrapWithAnchor
+  :: GHC.Id
+  -> GHC.RealSrcSpan
+  -> GHC.LHsExpr GHC.GhcTc
+  -> GHC.LHsExpr GHC.GhcTc
+wrapWithAnchor anchorId sp le@(GHC.L _ e)
+  | GHC.mightBeUnliftedType (GHC.hsExprType e) = le
+  | otherwise =
       let locStr = encodeSrcSpan sp
           locTy = GHC.LitTy (GHC.StrTyLit (GHC.mkFastString locStr))
           exprTy = GHC.hsExprType e
           wrapper = GHC.WpTyApp exprTy `GHC.WpCompose` GHC.WpTyApp locTy
           anchor = GHC.mkHsWrap wrapper (GHC.HsVar GHC.noExtField $ GHC.noLocA anchorId)
        in GHC.noLocA (Compat.hsAppTc (GHC.noLocA anchor) le)
-  | otherwise = le
+
+{-| Wrap the body of every 'GRHS' of a value binding with an @anchor@ carrying
+the binder's source span. This gives 'findAnchorLoc' a reliable outer-most
+anchor for every hoisted binding, so profile-trace output identifies which
+function was entered (the span points to the binder, not an arbitrary
+sub-expression). See issue #7722. -}
+anchorBinding
+  :: GHC.Id
+  -> GHC.HsBindLR GHC.GhcTc GHC.GhcTc
+  -> GHC.HsBindLR GHC.GhcTc GHC.GhcTc
+anchorBinding anchorId = \case
+  fb@GHC.FunBind {GHC.fun_id = GHC.L ann _, GHC.fun_matches = mg}
+    | Just sp <- GHC.srcSpanToRealSrcSpan (GHC.locA ann) ->
+        fb {GHC.fun_matches = wrapMatchGroup sp mg}
+  pb@GHC.PatBind {GHC.pat_lhs = GHC.L ann _, GHC.pat_rhs = grhss}
+    | Just sp <- GHC.srcSpanToRealSrcSpan (GHC.locA ann) ->
+        pb {GHC.pat_rhs = wrapGRHSs sp grhss}
+  other -> other
+  where
+    -- Extension constructors (XMG/XMatch/XGRHSs/XGRHS) are uninhabited for
+    -- 'GhcTc', but GHC's pattern coverage checker does not know that, so each
+    -- helper has a fall-through to silence the warning.
+    wrapMatchGroup sp = \case
+      mg@GHC.MG {GHC.mg_alts = lalts} ->
+        mg {GHC.mg_alts = (fmap . fmap . fmap) (wrapMatch sp) lalts}
+      other -> other
+    wrapMatch sp = \case
+      m@GHC.Match {GHC.m_grhss = grhss} ->
+        m {GHC.m_grhss = wrapGRHSs sp grhss}
+      other -> other
+    wrapGRHSs sp = \case
+      grhss@GHC.GRHSs {GHC.grhssGRHSs = lgrhss} ->
+        grhss {GHC.grhssGRHSs = (fmap . fmap) (wrapGRHS sp) lgrhss}
+      other -> other
+    wrapGRHS sp = \case
+      GHC.GRHS x guards body ->
+        GHC.GRHS x guards (wrapWithAnchor anchorId sp body)
+      other -> other
 
 isAnchorWorthy :: GHC.Id -> GHC.HsExpr GHC.GhcTc -> Bool
 isAnchorWorthy marker expr

plutus-tx-plugin/test/CallTrace/9.6/func03.golden.eval

@@ -5,4 +5,4 @@ Caused by: error
 Trace:
 -> func (test/CallTrace/Spec.hs:84:1-84:4)
 -> functionFromOtherModule (test/CallTrace/Spec.hs:101:1-101:23)
--> errorWhenTrue (test/CallTrace/OtherModule.hs:10:1-10:13)
+-> errorWhenTrue (test/CallTrace/OtherModule.hs:15:1-15:13)

plutus-tx-plugin/test/CallTrace/9.6/func05.golden.eval

@@ -6,4 +6,4 @@ Trace:
 -> func (test/CallTrace/Spec.hs:84:1-84:4)
 -> $cmyClassFunc (test/CallTrace/Spec.hs:107:3-107:13)
 -> functionFromOtherModule (test/CallTrace/Spec.hs:101:1-101:23)
--> errorWhenTrue (test/CallTrace/OtherModule.hs:10:1-10:13)
+-> errorWhenTrue (test/CallTrace/OtherModule.hs:15:1-15:13)

plutus-tx-plugin/test/CallTrace/9.6/func07.golden.eval

@@ -4,4 +4,4 @@ Caused by: error
 
 Trace:
 -> func (test/CallTrace/Spec.hs:84:1-84:4)
--> $fMyClassInOtherModuleInteger_$cmyClassFuncInOtherModule (test/CallTrace/OtherModule.hs:24:3-24:26)
+-> $fMyClassInOtherModuleInteger_$cmyClassFuncInOtherModule (test/CallTrace/OtherModule.hs:25:3-25:26)

plutus-tx-plugin/test/CallTrace/9.6/func08.golden.eval

@@ -4,5 +4,5 @@ Caused by: error
 
 Trace:
 -> func (test/CallTrace/Spec.hs:84:1-84:4)
--> $fMyClassInOtherModuleInteger_$cmyClassFuncInOtherModule (test/CallTrace/OtherModule.hs:24:3-24:26)
--> errorWhenTrue (test/CallTrace/OtherModule.hs:10:1-10:13)
+-> $fMyClassInOtherModuleInteger_$cmyClassFuncInOtherModule (test/CallTrace/OtherModule.hs:25:3-25:26)
+-> errorWhenTrue (test/CallTrace/OtherModule.hs:15:1-15:13)

plutus-tx-plugin/test/CallTrace/9.6/func09.golden.eval

@@ -4,4 +4,4 @@ Caused by: error
 
 Trace:
 -> func (test/CallTrace/Spec.hs:84:1-84:4)
--> $fMyClassInOtherModule()_$cmyClassFuncInOtherModule (test/CallTrace/OtherModule.hs:29:3-29:26)
+-> $fMyClassInOtherModule()_$cmyClassFuncInOtherModule (test/CallTrace/OtherModule.hs:30:3-30:26)

plutus-tx-plugin/test/CallTrace/9.6/funcionFromOtherModule-error.golden.eval

@@ -4,4 +4,4 @@ Caused by: error
 
 Trace:
 -> functionFromOtherModule (test/CallTrace/Spec.hs:101:1-101:23)
--> errorWhenTrue (test/CallTrace/OtherModule.hs:10:1-10:13)
+-> errorWhenTrue (test/CallTrace/OtherModule.hs:15:1-15:13)

plutus-tx-plugin/test/CallTrace/9.6/successfullEvaluationYieldsNoTraceLog.golden.eval

@@ -1,7 +1,7 @@
-CPU:                 1_653_088
-Memory:                  8_392
-AST Size:                   90
-Flat Size:                 451
+CPU:                 2_579_080
+Memory:                 12_820
+AST Size:                  116
+Flat Size:                 489
 
 No Trace Produced
 

plutus-tx-plugin/test/CallTrace/OtherModule.hs

@@ -1,6 +1,12 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE NoImplicitPrelude #-}
 {-# OPTIONS_GHC -fexpose-all-unfoldings #-}
+{-# OPTIONS_GHC -fplugin Plinth.Plugin #-}
+
+-- It is necessary to have Plinth plugin enabled here so that anchors get correctly added.
+-- If Plinth plugin is not enabled, profile-all will not be able to correctly retrieve
+-- srcspan and in the call trace function call in this module will not have any srcspan
+-- information.
 
 module CallTrace.OtherModule where
 
@@ -9,13 +15,9 @@ import PlutusTx.Prelude
 errorWhenTrue :: Bool -> BuiltinString
 errorWhenTrue True = error ()
 errorWhenTrue False = "hi"
--- NOINLINE ensures that the function appear in the call trace.
-{-# NOINLINE errorWhenTrue #-}
 
 wraps :: Bool -> BuiltinString
 wraps = errorWhenTrue
--- NOINLINE ensures that the function appear in the call trace.
-{-# NOINLINE wraps #-}
 
 class MyClassInOtherModule a where
   myClassFuncInOtherModule :: a -> BuiltinString

plutus-tx-plugin/test/CallTrace/Spec.hs

@@ -8,7 +8,7 @@
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE NoImplicitPrelude #-}
 {-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}
-{-# OPTIONS_GHC -fplugin PlutusTx.Plugin #-}
+{-# OPTIONS_GHC -fplugin Plinth.Plugin #-}
 {-# OPTIONS_GHC -fplugin-opt PlutusTx.Plugin:context-level=0 #-}
 {-# OPTIONS_GHC -fplugin-opt PlutusTx.Plugin:datatypes=BuiltinCasing #-}
 {-# OPTIONS_GHC -fplugin-opt PlutusTx.Plugin:max-cse-iterations=0 #-}

plutus-tx-plugin/test/Plugin/Profiling/9.6/addInt.golden.pir

@@ -1,18 +1,33 @@
 let
   !addInteger : integer -> integer -> integer = addInteger
+  ~addInteger : integer -> integer -> integer
+    = \(x : integer) ->
+        let
+          !x : integer = x
+        in
+        \(y : integer) ->
+          let
+            !y : integer = y
+          in
+          trace
+            {unit -> integer}
+            "-> addInteger"
+            (\(thunk : unit) ->
+               trace {integer} "<- addInteger" (addInteger x y))
+            ()
   ~addInt : integer -> integer -> integer
     = \(x : integer) ->
         let
           !x : integer = x
         in
         trace
           {unit -> integer -> integer}
-          "-> addInt (test/Plugin/Profiling/Spec.hs:116:1-116:6)"
+          "-> addInt (test/Plugin/Profiling/Spec.hs:115:1-115:6)"
           (\(thunk : unit) ->
              trace
                {integer -> integer}
-               "<- addInt (test/Plugin/Profiling/Spec.hs:116:1-116:6)"
-               (\(y : integer) -> let !y : integer = y in addInteger x y))
+               "<- addInt (test/Plugin/Profiling/Spec.hs:115:1-115:6)"
+               (addInteger x))
           ()
 in
 addInt

plutus-tx-plugin/test/Plugin/Profiling/9.6/addInt3.golden.eval

@@ -1,2 +1,2 @@
-[ -> addInt (test/Plugin/Profiling/Spec.hs:116:1-116:6)
-, <- addInt (test/Plugin/Profiling/Spec.hs:116:1-116:6) ]
+[ -> addInt (test/Plugin/Profiling/Spec.hs:115:1-115:6)
+, <- addInt (test/Plugin/Profiling/Spec.hs:115:1-115:6) ]

plutus-tx-plugin/test/Plugin/Profiling/9.6/argMismatch1.golden.eval

@@ -1,6 +1,6 @@
 [ -> runIdentity
 , <- runIdentity
-, -> newtypeFunction (test/Plugin/Profiling/Spec.hs:176:1-176:15)
-, <- newtypeFunction (test/Plugin/Profiling/Spec.hs:176:1-176:15)
+, -> newtypeFunction (test/Plugin/Profiling/Spec.hs:175:1-175:15)
+, <- newtypeFunction (test/Plugin/Profiling/Spec.hs:175:1-175:15)
 , -> $fFoldableIdentity2
 , <- $fFoldableIdentity2 ]

plutus-tx-plugin/test/Plugin/Profiling/9.6/argMismatch2.golden.eval

@@ -1,2 +1,2 @@
-[ -> obscuredFunction (test/Plugin/Profiling/Spec.hs:184:1-184:16)
-, <- obscuredFunction (test/Plugin/Profiling/Spec.hs:184:1-184:16) ]
+[ -> obscuredFunction (test/Plugin/Profiling/Spec.hs:183:1-183:16)
+, <- obscuredFunction (test/Plugin/Profiling/Spec.hs:183:1-183:16) ]