Introduce HasAnn typeclass; fix Apply span regression

plutus-core/changelog.d/20260421_120000_zeme_has_ann_typeclass.md

@@ -0,0 +1,5 @@
+### Fixed
+
+- Fixed `SrcSpan` annotations on inner `Apply`/`TyInst` nodes from the PLC, PIR, and
+  UPLC parsers. Parser error messages and tooling that reads annotations off parsed
+  terms will now point to the specific argument involved.

plutus-core/plutus-core/src/PlutusCore.hs

@@ -65,8 +65,7 @@ module PlutusCore
   , UniqueSet (..)
   , Normalized (..)
   , latestVersion
-  , termAnn
-  , typeAnn
+  , HasAnn (..)
   , tyVarDeclAnn
   , tyVarDeclName
   , tyVarDeclKind
@@ -161,6 +160,6 @@ applyProgram
   -> m (Program tyname name uni fun a)
 applyProgram (Program a1 v1 t1) (Program a2 v2 t2)
   | v1 == v2 =
-      pure $ Program (a1 <> a2) v1 (Apply (termAnn t1 <> termAnn t2) t1 t2)
+      pure $ Program (a1 <> a2) v1 (Apply (getAnn t1 <> getAnn t2) t1 t2)
 applyProgram (Program _a1 v1 _t1) (Program _a2 v2 _t2) =
   throwError $ MkApplyProgramError v1 v2

plutus-core/plutus-core/src/PlutusCore/Check/Normal.hs

@@ -67,7 +67,7 @@ neutralType
 neutralType TyVar {} = pure ()
 neutralType (TyBuiltin ann someUni) = neutralUni ann someUni
 neutralType (TyApp _ ty1 ty2) = neutralType ty1 >> normalType ty2
-neutralType ty = Left (BadType (typeAnn ty) ty "neutral type")
+neutralType ty = Left (BadType (getAnn ty) ty "neutral type")
 
 -- See Note [Normalization of built-in types].
 neutralUni

plutus-core/plutus-core/src/PlutusCore/Check/Value.hs

@@ -17,4 +17,4 @@ termValue (IWrap _ _ _ term) = termValue term
 termValue LamAbs {} = pure ()
 termValue TyAbs {} = pure ()
 termValue Constant {} = pure ()
-termValue t = Left $ BadTerm (termAnn t) t "term value"
+termValue t = Left $ BadTerm (getAnn t) t "term value"

plutus-core/plutus-core/src/PlutusCore/Core/Type.hs

@@ -38,9 +38,8 @@ module PlutusCore.Core.Type
   , Binder (..)
   , module Export
 
-    -- * Helper functions
-  , termAnn
-  , typeAnn
+    -- * Annotations
+  , HasAnn (..)
   , mapFun
   , tyVarDeclAnn
   , tyVarDeclName
@@ -263,29 +262,49 @@ type instance
   HasUniques (Program tyname name uni fun ann) =
     HasUniques (Term tyname name uni fun ann)
 
-typeAnn :: Type tyname uni ann -> ann
-typeAnn (TyVar ann _) = ann
-typeAnn (TyFun ann _ _) = ann
-typeAnn (TyIFix ann _ _) = ann
-typeAnn (TyForall ann _ _ _) = ann
-typeAnn (TyBuiltin ann _) = ann
-typeAnn (TyLam ann _ _ _) = ann
-typeAnn (TyApp ann _ _) = ann
-typeAnn (TySOP ann _) = ann
-
-termAnn :: Term tyname name uni fun ann -> ann
-termAnn (Var ann _) = ann
-termAnn (TyAbs ann _ _ _) = ann
-termAnn (Apply ann _ _) = ann
-termAnn (Constant ann _) = ann
-termAnn (Builtin ann _) = ann
-termAnn (TyInst ann _ _) = ann
-termAnn (Unwrap ann _) = ann
-termAnn (IWrap ann _ _ _) = ann
-termAnn (Error ann _) = ann
-termAnn (LamAbs ann _ _ _) = ann
-termAnn (Constr ann _ _ _) = ann
-termAnn (Case ann _ _ _) = ann
+instance HasAnn (Type tyname uni) where
+  getAnn (TyVar ann _)        = ann
+  getAnn (TyFun ann _ _)      = ann
+  getAnn (TyIFix ann _ _)     = ann
+  getAnn (TyForall ann _ _ _) = ann
+  getAnn (TyBuiltin ann _)    = ann
+  getAnn (TyLam ann _ _ _)    = ann
+  getAnn (TyApp ann _ _)      = ann
+  getAnn (TySOP ann _)        = ann
+  modifyAnn f (TyVar ann x)        = TyVar (f ann) x
+  modifyAnn f (TyFun ann a b)      = TyFun (f ann) a b
+  modifyAnn f (TyIFix ann a b)     = TyIFix (f ann) a b
+  modifyAnn f (TyForall ann tn k t) = TyForall (f ann) tn k t
+  modifyAnn f (TyBuiltin ann b)    = TyBuiltin (f ann) b
+  modifyAnn f (TyLam ann tn k t)   = TyLam (f ann) tn k t
+  modifyAnn f (TyApp ann a b)      = TyApp (f ann) a b
+  modifyAnn f (TySOP ann tss)      = TySOP (f ann) tss
+
+instance HasAnn (Term tyname name uni fun) where
+  getAnn (Var ann _)           = ann
+  getAnn (LamAbs ann _ _ _)   = ann
+  getAnn (Apply ann _ _)      = ann
+  getAnn (TyAbs ann _ _ _)    = ann
+  getAnn (TyInst ann _ _)     = ann
+  getAnn (IWrap ann _ _ _)    = ann
+  getAnn (Unwrap ann _)       = ann
+  getAnn (Constr ann _ _ _)   = ann
+  getAnn (Case ann _ _ _)     = ann
+  getAnn (Constant ann _)     = ann
+  getAnn (Builtin ann _)      = ann
+  getAnn (Error ann _)        = ann
+  modifyAnn f (Var ann x)           = Var (f ann) x
+  modifyAnn f (LamAbs ann n ty t)   = LamAbs (f ann) n ty t
+  modifyAnn f (Apply ann t1 t2)     = Apply (f ann) t1 t2
+  modifyAnn f (TyAbs ann tn k t)    = TyAbs (f ann) tn k t
+  modifyAnn f (TyInst ann t ty)     = TyInst (f ann) t ty
+  modifyAnn f (IWrap ann ty1 ty2 t) = IWrap (f ann) ty1 ty2 t
+  modifyAnn f (Unwrap ann t)        = Unwrap (f ann) t
+  modifyAnn f (Constr ann ty i ts)  = Constr (f ann) ty i ts
+  modifyAnn f (Case ann ty t ts)    = Case (f ann) ty t ts
+  modifyAnn f (Constant ann c)      = Constant (f ann) c
+  modifyAnn f (Builtin ann b)       = Builtin (f ann) b
+  modifyAnn f (Error ann ty)        = Error (f ann) ty
 
 -- | Map a function over the set of built-in functions.
 mapFun :: (fun -> fun') -> Term tyname name uni fun ann -> Term tyname name uni fun' ann

plutus-core/plutus-core/src/PlutusCore/Parser.hs

@@ -1,5 +1,4 @@
 {-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE TupleSections #-}
 
 -- | Parsers for PLC terms in DefaultUni.
 module PlutusCore.Parser
@@ -13,7 +12,7 @@ module PlutusCore.Parser
   ) where
 
 import PlutusCore.Annotation
-import PlutusCore.Core (Program (..), Term (..), Type)
+import PlutusCore.Core (HasAnn (..), Program (..), Term (..), Type)
 import PlutusCore.Default
 import PlutusCore.Error (ParserError (..), ParserErrorBundle)
 import PlutusCore.MkPlc (mkIterApp, mkIterInst)
@@ -24,6 +23,7 @@ import PlutusCore.Parser.Type as Export
 import PlutusCore.Quote (MonadQuote)
 import PlutusCore.Version
 
+import Control.Arrow ((&&&))
 import Control.Monad (when)
 import Control.Monad.Except (MonadError)
 import Data.Text (Text)
@@ -49,8 +49,9 @@ lamTerm = withSpan $ \sp ->
 
 appTerm :: Parser PTerm
 appTerm = withSpan $ \sp ->
-  -- TODO: should not use the same `sp` for all arguments.
-  inBrackets $ mkIterApp <$> term <*> (fmap (sp,) <$> some term)
+    inBrackets $
+      setAnn sp <$>
+        (mkIterApp <$> term <*> (fmap (getAnn &&& id) <$> some term))
 
 conTerm :: Parser PTerm
 conTerm = withSpan $ \sp ->
@@ -62,8 +63,9 @@ builtinTerm = withSpan $ \sp ->
 
 tyInstTerm :: Parser PTerm
 tyInstTerm = withSpan $ \sp ->
-  -- TODO: should not use the same `sp` for all arguments.
-  inBraces $ mkIterInst <$> term <*> (fmap (sp,) <$> many pType)
+    inBraces $
+      setAnn sp <$>
+        (mkIterInst <$> term <*> (fmap (getAnn &&& id) <$> many pType))
 
 unwrapTerm :: Parser PTerm
 unwrapTerm = withSpan $ \sp ->

plutus-core/plutus-core/src/PlutusCore/Parser/Type.hs

@@ -1,7 +1,6 @@
 -- editorconfig-checker-disable-file
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE TupleSections #-}
 {-# LANGUAGE TypeApplications #-}
 
 module PlutusCore.Parser.Type where
@@ -65,8 +64,7 @@ appType :: Parser PType
 appType = withSpan $ \sp -> inBrackets $ do
   fn <- pType
   args <- some pType
-  -- TODO: should not use the same `sp` for all arguments.
-  pure $ mkIterTyApp fn ((sp,) <$> args)
+  pure . setAnn sp $ mkIterTyApp fn (map (getAnn &&& id) args)
 
 kind :: Parser (Kind SrcSpan)
 kind = withSpan $ \sp ->

plutus-core/plutus-core/test/Parser/Spec.hs

@@ -28,7 +28,7 @@ propTermSrcSpan = property $ do
   trailingSpaces <- forAll $ Gen.text (Range.linear 0 10) (Gen.element [' ', '\n'])
   case runQuoteT . parseTerm $ code <> trailingSpaces of
     Right parsed ->
-      let sp = termAnn parsed
+      let sp = getAnn parsed
        in (srcSpanELine sp, srcSpanECol sp) === (endingLine, endingCol + 1)
     Left err -> annotate (display err) >> failure
 
@@ -92,4 +92,32 @@ tests =
     , testCase
         "parser of Value should succeed"
         parseValueValid
+    , testCase
+        "multi-arg application has per-argument spans on inner nodes"
+        multiArgSpans
     ]
+
+-- | Test that inner Apply nodes get per-argument spans, not the bracket span.
+-- For @[ (con integer 1) (con integer 2) (con integer 3) ]@, the outer Apply
+-- should have the bracket span, but the inner Apply should have the span of
+-- its argument @(con integer 2)@, NOT the bracket span.
+multiArgSpans :: Assertion
+multiArgSpans = do
+  let code = "[ (con integer 1) (con integer 2) (con integer 3) ]"
+  case runQuoteT (parseTerm code) of
+    Left err -> assertFailure $ "parse failed: " <> show err
+    Right parsed ->
+      case parsed of
+        Apply outerAnn (Apply innerAnn _ _) _ -> do
+          -- outer should have the bracket span (col 1 to col 52)
+          assertBool "outer span should start at col 1"
+            (srcSpanSCol outerAnn == 1)
+          -- inner should NOT have the same span as outer
+          assertBool
+            ("inner Apply should have a different span than outer, but both are: " <> show outerAnn)
+            (outerAnn /= innerAnn)
+          -- inner span should start after col 1 (it's the span of the 2nd arg)
+          assertBool
+            ("inner Apply span should not start at col 1, got: " <> show innerAnn)
+            (srcSpanSCol innerAnn /= 1)
+        other -> assertFailure $ "expected nested Apply, got: " <> show other

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

@@ -1,14 +1,13 @@
 module PlutusIR
   ( -- * AST
-    Term (..)
+    HasAnn (..)
+  , Term (..)
   , progAnn
   , progVer
   , progTerm
   , termSubterms
   , termSubtypes
   , termBindings
-  , termAnn
-  , bindingAnn
   , Type (..)
   , typeSubtypes
   , Datatype (..)

plutus-core/plutus-ir/src/PlutusIR/Compiler/Recursion.hs

@@ -91,7 +91,7 @@ mkFixpoint bs = do
     case PIR.mkFunctionDef p name ty term of
       Just fun -> pure fun
       Nothing ->
-        lift $ throwError $ CompilationError (PLC.typeAnn ty) "Recursive values must be of function type"
+        lift $ throwError $ CompilationError (PLC.getAnn ty) "Recursive values must be of function type"
 
   inlineFix <- view (ccOpts . coInlineConstants)
 

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

@@ -7,7 +7,8 @@
 {-# LANGUAGE UndecidableInstances #-}
 
 module PlutusIR.Core.Type
-  ( TyName (..)
+  ( HasAnn (..)
+  , TyName (..)
   , Name (..)
   , VarDecl (..)
   , TyVarDecl (..)
@@ -24,8 +25,6 @@ module PlutusIR.Core.Type
   , Program (..)
   , Version (..)
   , applyProgram
-  , termAnn
-  , bindingAnn
   , progAnn
   , progVer
   , progTerm
@@ -227,28 +226,46 @@ applyProgram
   -> m (Program tyname name uni fun a)
 applyProgram (Program a1 v1 t1) (Program a2 v2 t2)
   | v1 == v2 =
-      pure $ Program (a1 <> a2) v1 (Apply (termAnn t1 <> termAnn t2) t1 t2)
+      pure $ Program (a1 <> a2) v1 (Apply (getAnn t1 <> getAnn t2) t1 t2)
 applyProgram (Program _a1 v1 _t1) (Program _a2 v2 _t2) =
   throwError $ MkApplyProgramError v1 v2
 
-termAnn :: Term tyname name uni fun a -> a
-termAnn = \case
-  Let a _ _ _ -> a
-  Var a _ -> a
-  TyAbs a _ _ _ -> a
-  LamAbs a _ _ _ -> a
-  Apply a _ _ -> a
-  Constant a _ -> a
-  Builtin a _ -> a
-  TyInst a _ _ -> a
-  Error a _ -> a
-  IWrap a _ _ _ -> a
-  Unwrap a _ -> a
-  Constr a _ _ _ -> a
-  Case a _ _ _ -> a
-
-bindingAnn :: Binding tyname name uni fun a -> a
-bindingAnn = \case
-  TermBind a _ _ _ -> a
-  TypeBind a _ _ -> a
-  DatatypeBind a _ -> a
+instance HasAnn (Term tyname name uni fun) where
+  getAnn = \case
+    Let a _ _ _   -> a
+    Var a _       -> a
+    TyAbs a _ _ _ -> a
+    LamAbs a _ _ _ -> a
+    Apply a _ _   -> a
+    Constant a _  -> a
+    Builtin a _   -> a
+    TyInst a _ _  -> a
+    Error a _     -> a
+    IWrap a _ _ _ -> a
+    Unwrap a _    -> a
+    Constr a _ _ _ -> a
+    Case a _ _ _  -> a
+  modifyAnn f = \case
+    Let a r bs t      -> Let (f a) r bs t
+    Var a x           -> Var (f a) x
+    TyAbs a tn k t    -> TyAbs (f a) tn k t
+    LamAbs a n ty t   -> LamAbs (f a) n ty t
+    Apply a t1 t2     -> Apply (f a) t1 t2
+    Constant a c      -> Constant (f a) c
+    Builtin a b       -> Builtin (f a) b
+    TyInst a t ty     -> TyInst (f a) t ty
+    Error a ty        -> Error (f a) ty
+    IWrap a ty1 ty2 t -> IWrap (f a) ty1 ty2 t
+    Unwrap a t        -> Unwrap (f a) t
+    Constr a ty i ts  -> Constr (f a) ty i ts
+    Case a ty t ts    -> Case (f a) ty t ts
+
+instance HasAnn (Binding tyname name uni fun) where
+  getAnn = \case
+    TermBind a _ _ _  -> a
+    TypeBind a _ _    -> a
+    DatatypeBind a _  -> a
+  modifyAnn f = \case
+    TermBind a s d t  -> TermBind (f a) s d t
+    TypeBind a d t    -> TypeBind (f a) d t
+    DatatypeBind a d  -> DatatypeBind (f a) d

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

@@ -1,5 +1,4 @@
 {-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE TupleSections #-}
 
 -- | Parsers for PIR terms in DefaultUni.
 module PlutusIR.Parser
@@ -128,13 +127,15 @@ letTerm = withSpan $ \sp ->
 
 appTerm :: Parametric
 appTerm tm = withSpan $ \sp ->
-  -- TODO: should not use the same `sp` for all arguments.
-  inBrackets $ PIR.mkIterApp <$> tm <*> (fmap (sp,) <$> some tm)
+    inBrackets $
+      setAnn sp <$>
+        (PIR.mkIterApp <$> tm <*> (fmap (getAnn &&& id) <$> some tm))
 
 tyInstTerm :: Parametric
 tyInstTerm tm = withSpan $ \sp ->
-  -- TODO: should not use the same `sp` for all arguments.
-  inBraces $ PIR.mkIterInst <$> tm <*> (fmap (sp,) <$> some pType)
+    inBraces $
+      setAnn sp <$>
+        (PIR.mkIterInst <$> tm <*> (fmap (getAnn &&& id) <$> some pType))
 
 pTerm :: Parser PTerm
 pTerm = leadingWhitespace go

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

@@ -127,11 +127,11 @@ beta = over termSubterms beta . localTransform
     localTransform = \case
       -- See Note [Multi-beta]
       -- This maybe isn't the best annotation for this term, but it will do.
-      (extractBindings -> Just (bs, t)) -> Let (termAnn t) NonRec bs t
+      (extractBindings -> Just (bs, t)) -> Let (getAnn t) NonRec bs t
       -- See Note [Multi-beta] for why we don't perform multi-beta on `TyInst`.
       TyInst _ (TyAbs a n k body) tyArg ->
         let b = TypeBind a (TyVarDecl a n k) tyArg
-         in Let (termAnn body) NonRec (pure b) body
+         in Let (getAnn body) NonRec (pure b) body
       t -> t
 
 betaPassSC