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

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

@@ -617,6 +617,7 @@ hoistExpr
 hoistExpr var t = do
   wrapUnsafeDataAsConstrName <-
     lookupGhcName 'PlutusTx.AsData.Internal.wrapUnsafeDataAsConstr
+  bindingAnchorName <- lookupGhcName 'PlutusTx.Plugin.Utils.bindingAnchor
   let name = GHC.getName var
       lexName = LexName name
 
@@ -635,7 +636,7 @@ hoistExpr var t = do
   -- See Note [Dependency tracking]
   modifyCurDeps (Set.insert lexName)
   maybeDef <- PIR.lookupTerm lexName
-  let varSpan = getVarSourceSpan var
+  let varSpan = findBindingAnchorLoc bindingAnchorName t <|> getVarSourceSpan var
       addSpan = case varSpan of
         Nothing -> id
         Just src -> fmap . fmap . addSrcSpan $ src ^. srcSpanIso
@@ -650,27 +651,25 @@ hoistExpr var t = do
         (PIR.Def var' (PIR.mkVar var', PIR.Strict))
         mempty
 
-      t' <- maybeProfileRhs var var' =<< addSpan (compileExpr Nothing t)
+      t' <- maybeProfileRhs varSpan var 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.
 maybeProfileRhs
   :: CompilingDefault uni fun m ann
-  => GHC.Var
+  => Maybe GHC.RealSrcSpan
+  -> GHC.Var
   -> PLCVar uni
   -> PIRTerm uni fun
   -> m (PIRTerm uni fun)
-maybeProfileRhs ghcVar var t = do
+maybeProfileRhs mSpan ghcVar 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) <> ")"
 
@@ -937,6 +936,7 @@ compileExpr mloc e = do
   boolOperatorAnd <- lookupGhcName '(PlutusTx.Bool.&&)
   inlineName <- lookupGhcName 'PlutusTx.Optimize.Inline.inline
   anchorName <- lookupGhcName 'PlutusTx.Plugin.Utils.anchor
+  bindingAnchorName <- lookupGhcName 'PlutusTx.Plugin.Utils.bindingAnchor
   unsupportedName <- lookupGhcName 'PlutusTx.Plugin.Utils.unsupported
 
   caseIntegerName <- lookupGhcName 'Builtins.caseInteger
@@ -985,7 +985,8 @@ compileExpr mloc e = do
       throwPlain . UnsupportedError $ T.pack msg
     Nothing -> pure ()
 
-  case extractLoc anchorName maybeModBreaks e of
+  let (_, eWithoutBindingAnchor) = extractBindingAnchor bindingAnchorName e
+  case extractLoc anchorName maybeModBreaks eWithoutBindingAnchor of
     (Just loc, e') -> do
       res <- compileExpr (Just loc) e'
       CompileContext {ccOpts = coverageOpts} <- ask
@@ -995,8 +996,8 @@ compileExpr mloc e = do
         (coverageCompile e' (GHC.exprType e') loc)
         (addSrcSpan (loc ^. srcSpanIso) <$> res)
         anns
-    _ -> traceCompilationL 2 (traceExprMsg mloc GHC.$$ GHC.ppr e) mloc $ do
-      case e of
+    _ -> traceCompilationL 2 (traceExprMsg mloc GHC.$$ GHC.ppr eWithoutBindingAnchor) mloc $ do
+      case eWithoutBindingAnchor of
         -- caseInteger: dispatch on an integer index to select a branch.
         GHC.App (GHC.App (GHC.App (GHC.Var var) (GHC.Type resTy)) scrut) li
           -- BuiltinCasing: compile to native UPLC case on the integer.
@@ -1328,7 +1329,8 @@ 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 = findBindingAnchorLoc bindingAnchorName rhs <|> getVarSourceSpan b
+            rhs'' <- maybeProfileRhs bSpan b v rhs'
             let binds = pure $ PIR.TermBind annMayInline PIR.NonStrict v rhs''
             body' <- compileExpr Nothing body
             pure $ PIR.Let annMayInline PIR.NonRec binds body'
@@ -1337,7 +1339,8 @@ 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 = findBindingAnchorLoc bindingAnchorName rhs <|> getVarSourceSpan ghcVar
+              rhs' <- maybeProfileRhs gvSpan ghcVar v =<< compileExpr Nothing rhs
               -- See Note [Non-strict let-bindings]
               pure $ PIR.TermBind annMayInline PIR.NonStrict v rhs'
             body' <- compileExpr Nothing body
@@ -1726,6 +1729,46 @@ extractLoc anchorName modBreaks = go Nothing
             go (Just ss) e
       other -> (acc, other)
 
+extractBindingAnchor
+  :: GHC.Name
+  -> GHC.CoreExpr
+  -> (Maybe GHC.RealSrcSpan, GHC.CoreExpr)
+extractBindingAnchor bindingAnchorName expr =
+  case go Nothing expr of
+    (Nothing, _) -> (Nothing, expr)
+    found -> found
+  where
+    go acc = \case
+      GHC.App
+        ( GHC.App
+            (GHC.App (GHC.Var f) (GHC.Type (GHC.LitTy (GHC.StrTyLit loc))))
+            (GHC.Type _eTy)
+          )
+        e
+          | GHC.getName f == bindingAnchorName ->
+              go (acc <|> decodeSrcSpan (GHC.unpackFS loc)) e
+      other -> (acc, other)
+
+findBindingAnchorLoc :: GHC.Name -> GHC.CoreExpr -> Maybe GHC.RealSrcSpan
+findBindingAnchorLoc bindingAnchorName = 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 == bindingAnchorName -> decodeSrcSpan (GHC.unpackFS loc)
+      GHC.Lam _ body -> go body
+      GHC.Let (GHC.NonRec _ rhs) body -> go body <|> go rhs
+      GHC.Let (GHC.Rec binds) body -> go body <|> firstJust (map snd binds)
+      GHC.Case _ _ _ alts -> firstJust [rhs | GHC.Alt _ _ rhs <- alts]
+      GHC.Cast body _ -> go body
+      GHC.Tick _ body -> go body
+      _ -> Nothing
+
 extractUnsupported
   :: GHC.Name
   -> GHC.CoreExpr

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

@@ -107,6 +107,7 @@ data PluginCtx = PluginCtx
   , pcFamEnvs :: GHC.FamInstEnvs
   , pcMarkerName :: GHC.Name
   , pcAnchorName :: GHC.Name
+  , pcBindingAnchorName :: GHC.Name
   , pcModuleName :: GHC.ModuleName
   , pcModuleModBreaks :: Maybe GHC.ModBreaks
   , pcPackageName :: String
@@ -158,9 +159,10 @@ installCorePlugin markerTHName args rest = do
       : pluginPass
       : rest
 
-plinthcModName, anchorName :: String
+plinthcModName, anchorName, bindingAnchorName :: String
 plinthcModName = fromJust $ TH.nameModule 'PlutusTx.Plugin.Utils.anchor
 anchorName = TH.nameBase 'PlutusTx.Plugin.Utils.anchor
+bindingAnchorName = TH.nameBase 'PlutusTx.Plugin.Utils.bindingAnchor
 
 -- | Wrap certain @HsExpr@s in the typed checked module with @anchor@.
 injectAnchors
@@ -174,17 +176,22 @@ injectAnchors env = do
         hscEnv
         (GHC.mkModuleName plinthcModName)
         GHC.NoPkgQual
-  anchorId <- case findResult of
+  (anchorId, bindingAnchorId) <- case findResult of
     GHC.Found _ m -> do
-      GHC.tcLookupId =<< GHC.lookupOrig m (GHC.mkVarOcc anchorName)
+      anchorId <- GHC.tcLookupId =<< GHC.lookupOrig m (GHC.mkVarOcc anchorName)
+      bindingAnchorId <- GHC.tcLookupId =<< GHC.lookupOrig m (GHC.mkVarOcc bindingAnchorName)
+      pure (anchorId, bindingAnchorId)
     _ ->
       GHC.pprPanic
         "Plinth.Plugin"
         (GHC.text $ "Could not find module " <> plinthcModName)
   let binds = GHC.tcg_binds env
       bindsAnchored =
         Compat.modifyBinds
-          (transformBi (stripGuardAnchors anchorId) . transformBi (anchorExpr anchorId))
+          ( transformBi (stripGuardAnchors anchorId)
+              . transformBi (anchorBinding bindingAnchorId)
+              . transformBi (anchorExpr anchorId)
+          )
           binds
   pure env {GHC.tcg_binds = bindsAnchored}
 
@@ -193,13 +200,53 @@ 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
+
+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
+
+anchorBinding
+  :: GHC.Id
+  -> GHC.HsBindLR GHC.GhcTc GHC.GhcTc
+  -> GHC.HsBindLR GHC.GhcTc GHC.GhcTc
+anchorBinding bindingAnchorId = \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
+    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 bindingAnchorId sp body)
+      other -> other
 
 isAnchorWorthy :: GHC.Id -> GHC.HsExpr GHC.GhcTc -> Bool
 isAnchorWorthy marker expr
@@ -363,9 +410,10 @@ mkPluginPass markerTHName opts = GHC.CoreDoPluginPass "Core to PLC" $ \guts -> d
   -- See Note [Marker resolution]
   maybeMarkerName <- GHC.thNameToGhcName markerTHName
   maybeanchorGhcName <- GHC.thNameToGhcName 'PlutusTx.Plugin.Utils.anchor
-  case (maybeMarkerName, maybeanchorGhcName) of
+  maybeBindingAnchorGhcName <- GHC.thNameToGhcName 'PlutusTx.Plugin.Utils.bindingAnchor
+  case (maybeMarkerName, maybeanchorGhcName, maybeBindingAnchorGhcName) of
     -- See Note [Marker resolution]
-    (Just markerName, Just anchorGhcName) -> do
+    (Just markerName, Just anchorGhcName, Just bindingAnchorGhcName) -> do
       hscEnv <- GHC.getHscEnv
       let thisModule = GHC.mg_module guts
           pctx =
@@ -374,6 +422,7 @@ mkPluginPass markerTHName opts = GHC.CoreDoPluginPass "Core to PLC" $ \guts -> d
               , pcFamEnvs = (p_fam_env, GHC.mg_fam_inst_env guts)
               , pcMarkerName = markerName
               , pcAnchorName = anchorGhcName
+              , pcBindingAnchorName = bindingAnchorGhcName
               , pcModuleName = GHC.moduleName thisModule
               , pcModuleModBreaks = GHC.mg_modBreaks guts
               , pcPackageName = getPackageName hscEnv thisModule
@@ -510,7 +559,8 @@ compileMarkedExprs :: GHC.CoreExpr -> PluginM PLC.DefaultUni PLC.DefaultFun GHC.
 compileMarkedExprs expr = do
   markerName <- asks pcMarkerName
   anchorGhcName <- asks pcAnchorName
-  case expr of
+  bindingAnchorGhcName <- asks pcBindingAnchorName
+  case stripCoreAnchors bindingAnchorGhcName expr of
     -- This clause is for the `plc` marker. It can be removed when we remove `plc`.
     GHC.App
       ( GHC.App
@@ -619,6 +669,7 @@ compileMarkedExpr _locStr codeTy origE = do
            , 'mkNil
            , 'PlutusTx.Builtins.equalsInteger
            , 'PlutusTx.Plugin.Utils.anchor
+           , 'PlutusTx.Plugin.Utils.bindingAnchor
            , 'PlutusTx.Plugin.Utils.unsupported
            ]
   modBreaks <- asks pcModuleModBreaks

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

@@ -1,7 +1,7 @@
-CPU:                 1_621_088
-Memory:                  8_192
-AST Size:                   88
-Flat Size:                 452
+CPU:                 1_669_088
+Memory:                  8_492
+AST Size:                   91
+Flat Size:                 454
 
 No Trace Produced
 

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

@@ -1,7 +1,7 @@
-CPU:                 1_621_088
-Memory:                  8_192
-AST Size:                   88
-Flat Size:                 452
+CPU:                 1_669_088
+Memory:                  8_492
+AST Size:                   91
+Flat Size:                 454
 
 No Trace Produced
 

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

@@ -1,6 +1,6 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE NoImplicitPrelude #-}
-{-# OPTIONS_GHC -fexpose-all-unfoldings #-}
+{-# OPTIONS_GHC -fexpose-all-unfoldings -fplugin Plinth.Plugin -fplugin-opt Plinth.Plugin:preserve-source-locations #-}
 
 module CallTrace.OtherModule where
 

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

@@ -30,11 +30,19 @@ plinthc _ = SerializedCode (mustBeReplaced "plc") (mustBeReplaced "pir") (mustBe
 {-# OPAQUE plinthc #-}
 
 {-| This function is used in `typeCheckResultAction` to mark the given expression
-with its source location. -}
+with its source location. The added source location information is used for generating
+improved line number reporting on error messages. -}
 anchor :: forall (loc :: Symbol) a. a -> a
 anchor a = a
 {-# OPAQUE anchor #-}
 
+{-| This function is used in `typeCheckResultAction` to mark a binding with
+its binder source location. The added binding information is used for generating
+function enterance and exit information when `profile-all` option is enabled. -}
+bindingAnchor :: forall (loc :: Symbol) a. a -> a
+bindingAnchor a = a
+{-# OPAQUE bindingAnchor #-}
+
 {-| This function is used in `typeCheckResultAction` to mark the given expression
 as unsupported by Plinth. -}
 unsupported :: forall (err :: Symbol) (loc :: Symbol) a. a -> a