Add per-node source ranges to SpecExpr (#810)

Add a `loc : SourceRange` field (required, no default) to every SpecExpr
constructor so that each node in the expression tree carries its own
source location from the Python AST.

Changes:
- Decls.lean: Add `loc : SourceRange` to all 17 constructors
- DDM.lean: toDDM uses `loc` as the DDM annotation; fromDDM extracts it
- Specs.lean: Construction sites pass `e.ann` (Python expr source range)
where available. Structural wrappers pass the enclosing statement's
source range: `pred.ann` for implies/not (from the if-condition),
`s.ann` for forallList/forallDict (from the for-statement).
`assumeCondition` now takes a `loc` parameter.
- ToLaurel.lean: specExprToLaurel uses per-node `loc` to create metadata
via a `nodeMd` helper, falling back to the function-level `md` when
`loc` is default. This resolves the md-reuse concern from PR #804.

By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

---------

Co-authored-by: Kiro <kiro-agent@users.noreply.github.com>

Author: tautschnig

Strata/Languages/Python/Specs.lean

@@ -770,13 +770,13 @@ def extractKwargsField (e : expr SourceRange)
 partial def extractSubject (e : expr SourceRange)
     : SpecAssertionM (Option SpecExpr) := do
   match ← extractKwargsField e with
-  | some (kn, fn) => return some (.getIndex (.var kn) fn)
+  | some (kn, fn) => return some (.getIndex (.var kn (loc := e.ann)) fn (loc := e.ann))
   | none => pure ()
   match e with
-  | .Name _ ⟨_, name⟩ (.Load _) => return some (.var name)
+  | .Name _ ⟨_, name⟩ (.Load _) => return some (.var name (loc := e.ann))
   | .Subscript _ inner (.Constant _ (.ConString _ fieldName) _) (.Load _) =>
     match ← extractSubject inner with
-    | some subj => return some (.getIndex subj fieldName.val)
+    | some subj => return some (.getIndex subj fieldName.val (loc := e.ann))
     | none => return none
   | _ => return none
 
@@ -790,7 +790,7 @@ def transCondition (e : expr SourceRange) : SpecAssertionM (Option SpecExpr) :=
         match ops.val[0] with
         | .In _ =>
           match ← extractSubject comparators.val[0] with
-          | some subj => return some (.containsKey subj key.val)
+          | some subj => return some (.containsKey subj key.val (loc := e.ann))
           | none => pure ()
         | _ => pure ()
     pure ()
@@ -800,7 +800,7 @@ def transCondition (e : expr SourceRange) : SpecAssertionM (Option SpecExpr) :=
 /-- Run an action that may produce assertions, then wrap each new assertion's
     formula with `implies cond ...` (or `implies (not cond) ...` for else branches).
     If `cond` is `none`, assertions pass through unchanged. -/
-def assumeCondition (cond : Option SpecExpr) (act : SpecAssertionM Unit)
+def assumeCondition (cond : Option SpecExpr) (loc : SourceRange) (act : SpecAssertionM Unit)
     : SpecAssertionM Unit := do
   let prevAssertions := (←get).assertions
   modify fun s => { s with assertions := #[] }
@@ -809,7 +809,7 @@ def assumeCondition (cond : Option SpecExpr) (act : SpecAssertionM Unit)
   match cond with
   | some c =>
     let wrapped := newAssertions.map fun a =>
-      { a with formula := .implies c a.formula }
+      { a with formula := .implies c a.formula loc }
     modify fun s => { s with assertions := prevAssertions ++ wrapped }
   | none =>
     modify fun s => { s with assertions := prevAssertions ++ newAssertions }
@@ -826,10 +826,10 @@ def transMessageExpr (e : expr SourceRange)
   | .Call _ (.Name _ funcName (.Load _)) args _ =>
     if funcName.val == "len" && args.val.size == 1 then
       match ← extractSubject args.val[0]! with
-      | some subj => return .len subj
-      | none => return .placeholder
-    else return .placeholder
-  | _ => return .placeholder
+      | some subj => return .len subj (loc := e.ann)
+      | none => return .placeholder (loc := e.ann)
+    else return .placeholder (loc := e.ann)
+  | _ => return .placeholder (loc := e.ann)
 
 /-- Look up a field in a TypedDict SpecType, returning its type if found. -/
 def lookupTypedDictField (tp : SpecType) (field : String) : Option SpecType := do
@@ -951,24 +951,26 @@ private def makeComparison
     (isFloat isInt : Bool)
     (subj : SpecExpr) (bound : expr SourceRange)
     : SpecAssertionM (Option SpecExpr) := do
+  let loc := bound.ann
   if isFloat then
     match extractFloatBound bound with
-    | some s => return some (floatCtor subj (.floatLit s))
+    | some s => return some (floatCtor subj (.floatLit s (loc := loc)))
     | none =>
       match extractIntBound bound with
-      | some n => return some (floatCtor subj (.floatLit (toString n)))
+      | some n => return some (floatCtor subj (.floatLit (toString n) (loc := loc)))
       | none => return none
   else if isInt then
     match extractIntBound bound with
-    | some n => return some (intCtor subj (.intLit n))
+    | some n => return some (intCtor subj (.intLit n (loc := loc)))
     | none => return none
   else
     match extractIntBound bound with
-    | some n => return some (intCtor subj (.intLit n))
+    | some n => return some (intCtor subj (.intLit n (loc := loc)))
     | none => return none
 
 def transAssertExpr (e : expr SourceRange)
     : SpecAssertionM SpecExpr := do
+  let loc := e.ann
   -- isinstance(subject, T)
   match e with
   | .Call _ (.Name _ funcName (.Load _)) args _ =>
@@ -978,10 +980,10 @@ def transAssertExpr (e : expr SourceRange)
         | some subj =>
           match args.val[1] with
           | .Name _ typeName (.Load _) =>
-            return .isInstanceOf subj typeName.val
+            return .isInstanceOf subj typeName.val (loc := loc)
           | _ =>
             specWarning e.ann s!"isinstance: unsupported type argument"
-            return .placeholder
+            return .placeholder (loc := loc)
         | none => pure () -- fall through
     if funcName.val == "len" && args.val.size == 1 then
       -- This is just len(x), not a comparison; fall through
@@ -998,8 +1000,8 @@ def transAssertExpr (e : expr SourceRange)
             match ← extractSubject callArgs.val[0] with
             | some subj =>
               match ops.val[0], extractIntBound comparators.val[0] with
-              | .GtE _, some n => return .intGe (.len subj) (.intLit n)
-              | .LtE _, some n => return .intLe (.len subj) (.intLit n)
+              | .GtE _, some n => return .intGe (.len subj (loc := loc)) (.intLit n (loc := comparators.val[0].ann)) (loc := loc)
+              | .LtE _, some n => return .intLe (.len subj (loc := loc)) (.intLit n (loc := comparators.val[0].ann)) (loc := loc)
               | _, _ => pure ()
             | none => pure ()
   | _ => pure ()
@@ -1014,8 +1016,8 @@ def transAssertExpr (e : expr SourceRange)
           let isFloat := subjType.any isFloatType
           let isInt := subjType.any isIntType
           let cmp ← match ops.val[0] with
-          | .GtE _ => makeComparison .floatGe .intGe isFloat isInt subj comparators.val[0]
-          | .LtE _ => makeComparison .floatLe .intLe isFloat isInt subj comparators.val[0]
+          | .GtE _ => makeComparison (.floatGe · · (loc := loc)) (.intGe · · (loc := loc)) isFloat isInt subj comparators.val[0]
+          | .LtE _ => makeComparison (.floatLe · · (loc := loc)) (.intLe · · (loc := loc)) isFloat isInt subj comparators.val[0]
           | _ => pure none
           match cmp with
           | some expr => return expr
@@ -1025,7 +1027,7 @@ def transAssertExpr (e : expr SourceRange)
   -- subject == "A" or subject == "B" or ...
   match ← collectEnumValues e with
   | some (subj, vals) =>
-    return .enumMember subj vals
+    return .enumMember subj vals (loc := loc)
   | none => pure ()
   -- compile("pattern").search(subject) is not None
   match e with
@@ -1044,14 +1046,14 @@ def transAssertExpr (e : expr SourceRange)
                 | some subj =>
                   match ops.val[0], comparators.val[0] with
                   | .IsNot _, .Constant _ (.ConNone _) _ =>
-                    return .regexMatch subj pattern.val
+                    return .regexMatch subj pattern.val (loc := loc)
                   | _, _ => pure ()
                 | none => pure ()
               | _ => pure ()
   | _ => pure ()
   -- Fallback: unrecognized pattern
   specWarning e.ann s!"unrecognized assert pattern: {eformat e.toAst}"
-  return .placeholder
+  return .placeholder (loc := loc)
 
 mutual
 
@@ -1113,7 +1115,7 @@ def blockStmt (s : stmt SourceRange) : SpecAssertionM Unit := do
           blockStmts body.val
         let bodyAssertions := (←get).assertions
         let wrapped := bodyAssertions.map fun a =>
-          { a with formula := .forallList listExpr varName a.formula }
+          { a with formula := .forallList listExpr varName a.formula s.ann }
         modify fun s => { s with assertions := prevAssertions ++ wrapped }
       | none =>
         specWarning s.ann s!"For: cannot extract iterable expression"
@@ -1154,7 +1156,7 @@ def blockStmt (s : stmt SourceRange) : SpecAssertionM Unit := do
               blockStmts body.val
             let bodyAssertions := (←get).assertions
             let wrapped := bodyAssertions.map fun a =>
-              { a with formula := .forallDict dictSubj keyVar valVar a.formula }
+              { a with formula := .forallDict dictSubj keyVar valVar a.formula s.ann }
             modify fun st => { st with assertions := prevAssertions ++ wrapped }
           | none =>
             specWarning s.ann s!"For: cannot extract dict expression"
@@ -1166,9 +1168,9 @@ def blockStmt (s : stmt SourceRange) : SpecAssertionM Unit := do
     let cond ← transCondition pred
     if cond.isNone then
       specWarning pred.ann s!"if: unrecognized condition pattern: {eformat pred.toAst}"
-    assumeCondition cond <| blockStmts t.val
+    assumeCondition cond pred.ann <| blockStmts t.val
     if f.val.size > 0 then
-      assumeCondition (cond.map .not) <| blockStmts f.val
+      assumeCondition (cond.map (.not · pred.ann)) pred.ann <| blockStmts f.val
   | .Pass _ =>
     pure ()
   | _ => specError s.ann s!"Unsupported statement: {eformat s.toAst}"

Strata/Languages/Python/Specs/DDM.lean

@@ -243,31 +243,31 @@ private def Arg.toDDM (d : Arg) : DDM.ArgDecl SourceRange :=
 
 protected def SpecExpr.toDDM (e : SpecExpr) : DDM.SpecExprDecl SourceRange :=
   match e with
-  | .placeholder => .placeholderExpr .none
-  | .var name => .varExpr .none ⟨.none, name⟩
-  | .getIndex subj field => .getIndexExpr .none subj.toDDM ⟨.none, field⟩
-  | .isInstanceOf subj tn => .isInstanceOfExpr .none subj.toDDM ⟨.none, tn⟩
-  | .len subj => .lenExpr .none subj.toDDM
-  | .intLit v => .intExpr .none (toDDMInt .none v)
-  | .intGe subj bound => .intGeExpr .none subj.toDDM bound.toDDM
-  | .intLe subj bound => .intLeExpr .none subj.toDDM bound.toDDM
-  | .floatLit v => .floatExpr .none ⟨.none, v⟩
-  | .floatGe subj bound => .floatGeExpr .none subj.toDDM bound.toDDM
-  | .floatLe subj bound => .floatLeExpr .none subj.toDDM bound.toDDM
-  | .enumMember subj values =>
-    .enumMemberExpr .none subj.toDDM
-      ⟨.none, values.map (⟨.none, ·⟩)⟩
-  | .regexMatch subj pattern =>
-    .regexMatchExpr .none subj.toDDM ⟨.none, pattern⟩
-  | .containsKey container key =>
-    .containsKeyExpr .none container.toDDM ⟨.none, key⟩
-  | .implies cond body =>
-    .impliesExpr .none cond.toDDM body.toDDM
-  | .not e => .notExpr .none e.toDDM
-  | .forallList list varName body =>
-    .forallListExpr .none list.toDDM ⟨.none, varName⟩ body.toDDM
-  | .forallDict dict keyVar valVar body =>
-    .forallDictExpr .none dict.toDDM ⟨.none, keyVar⟩ ⟨.none, valVar⟩ body.toDDM
+  | .placeholder loc => .placeholderExpr loc
+  | .var name loc => .varExpr loc ⟨loc, name⟩
+  | .getIndex subj field loc => .getIndexExpr loc subj.toDDM ⟨loc, field⟩
+  | .isInstanceOf subj tn loc => .isInstanceOfExpr loc subj.toDDM ⟨loc, tn⟩
+  | .len subj loc => .lenExpr loc subj.toDDM
+  | .intLit v loc => .intExpr loc (toDDMInt loc v)
+  | .intGe subj bound loc => .intGeExpr loc subj.toDDM bound.toDDM
+  | .intLe subj bound loc => .intLeExpr loc subj.toDDM bound.toDDM
+  | .floatLit v loc => .floatExpr loc ⟨loc, v⟩
+  | .floatGe subj bound loc => .floatGeExpr loc subj.toDDM bound.toDDM
+  | .floatLe subj bound loc => .floatLeExpr loc subj.toDDM bound.toDDM
+  | .enumMember subj values loc =>
+    .enumMemberExpr loc subj.toDDM
+      ⟨loc, values.map (⟨loc, ·⟩)⟩
+  | .regexMatch subj pattern loc =>
+    .regexMatchExpr loc subj.toDDM ⟨loc, pattern⟩
+  | .containsKey container key loc =>
+    .containsKeyExpr loc container.toDDM ⟨loc, key⟩
+  | .implies cond body loc =>
+    .impliesExpr loc cond.toDDM body.toDDM
+  | .not e loc => .notExpr loc e.toDDM
+  | .forallList list varName body loc =>
+    .forallListExpr loc list.toDDM ⟨loc, varName⟩ body.toDDM
+  | .forallDict dict keyVar valVar body loc =>
+    .forallDictExpr loc dict.toDDM ⟨loc, keyVar⟩ ⟨loc, valVar⟩ body.toDDM
 
 def specExprFormatContext : FormatContext :=
   .ofDialects DDM.PythonSpecs_map
@@ -383,26 +383,26 @@ private def DDM.ArgDecl.fromDDM (d : DDM.ArgDecl SourceRange) : Specs.Arg :=
 
 private def DDM.SpecExprDecl.fromDDM (d : DDM.SpecExprDecl SourceRange) : Specs.SpecExpr :=
   match d with
-  | .placeholderExpr _ => .placeholder
-  | .varExpr _ ⟨_, name⟩ => .var name
-  | .getIndexExpr _ subj ⟨_, field⟩ => .getIndex subj.fromDDM field
-  | .isInstanceOfExpr _ subj ⟨_, tn⟩ => .isInstanceOf subj.fromDDM tn
-  | .lenExpr _ subj => .len subj.fromDDM
-  | .intExpr _ i => .intLit i.ofDDM
-  | .intGeExpr _ subj bound => .intGe subj.fromDDM bound.fromDDM
-  | .intLeExpr _ subj bound => .intLe subj.fromDDM bound.fromDDM
-  | .floatExpr _ ⟨_, v⟩ => .floatLit v
-  | .floatGeExpr _ subj bound => .floatGe subj.fromDDM bound.fromDDM
-  | .floatLeExpr _ subj bound => .floatLe subj.fromDDM bound.fromDDM
-  | .enumMemberExpr _ subj ⟨_, values⟩ => .enumMember subj.fromDDM (values.map (·.2))
-  | .regexMatchExpr _ subj ⟨_, pattern⟩ => .regexMatch subj.fromDDM pattern
-  | .containsKeyExpr _ container ⟨_, key⟩ => .containsKey container.fromDDM key
-  | .impliesExpr _ cond body => .implies cond.fromDDM body.fromDDM
-  | .notExpr _ e => .not e.fromDDM
-  | .forallListExpr _ list ⟨_, varName⟩ body =>
-    .forallList list.fromDDM varName body.fromDDM
-  | .forallDictExpr _ dict ⟨_, keyVar⟩ ⟨_, valVar⟩ body =>
-    .forallDict dict.fromDDM keyVar valVar body.fromDDM
+  | .placeholderExpr loc => .placeholder loc
+  | .varExpr loc ⟨_, name⟩ => .var name loc
+  | .getIndexExpr loc subj ⟨_, field⟩ => .getIndex subj.fromDDM field loc
+  | .isInstanceOfExpr loc subj ⟨_, tn⟩ => .isInstanceOf subj.fromDDM tn loc
+  | .lenExpr loc subj => .len subj.fromDDM loc
+  | .intExpr loc i => .intLit i.ofDDM loc
+  | .intGeExpr loc subj bound => .intGe subj.fromDDM bound.fromDDM loc
+  | .intLeExpr loc subj bound => .intLe subj.fromDDM bound.fromDDM loc
+  | .floatExpr loc ⟨_, v⟩ => .floatLit v loc
+  | .floatGeExpr loc subj bound => .floatGe subj.fromDDM bound.fromDDM loc
+  | .floatLeExpr loc subj bound => .floatLe subj.fromDDM bound.fromDDM loc
+  | .enumMemberExpr loc subj ⟨_, values⟩ => .enumMember subj.fromDDM (values.map (·.2)) loc
+  | .regexMatchExpr loc subj ⟨_, pattern⟩ => .regexMatch subj.fromDDM pattern loc
+  | .containsKeyExpr loc container ⟨_, key⟩ => .containsKey container.fromDDM key loc
+  | .impliesExpr loc cond body => .implies cond.fromDDM body.fromDDM loc
+  | .notExpr loc e => .not e.fromDDM loc
+  | .forallListExpr loc list ⟨_, varName⟩ body =>
+    .forallList list.fromDDM varName body.fromDDM loc
+  | .forallDictExpr loc dict ⟨_, keyVar⟩ ⟨_, valVar⟩ body =>
+    .forallDict dict.fromDDM keyVar valVar body.fromDDM loc
 
 private def DDM.MessagePart.fromDDM (d : DDM.MessagePart SourceRange) : Specs.MessagePart :=
   match d with

Strata/Languages/Python/Specs/Decls.lean

@@ -316,39 +316,39 @@ and `placeholder`; interior nodes represent operations like `len`, `getIndex`,
 inductive SpecExpr where
 /-- Stands in for an assert pattern not yet supported by the translator.
     The original Python expression is preserved in `Assertion.message`. -/
-| placeholder
-| var (name : String)
-| getIndex (subject : SpecExpr) (field : String)
-| isInstanceOf (subject : SpecExpr) (typeName : String)
-| len (subject : SpecExpr)
-| intLit (value : Int)
-| intGe (subject : SpecExpr) (bound : SpecExpr)
-| intLe (subject : SpecExpr) (bound : SpecExpr)
+| placeholder (loc : SourceRange)
+| var (name : String) (loc : SourceRange)
+| getIndex (subject : SpecExpr) (field : String) (loc : SourceRange)
+| isInstanceOf (subject : SpecExpr) (typeName : String) (loc : SourceRange)
+| len (subject : SpecExpr) (loc : SourceRange)
+| intLit (value : Int) (loc : SourceRange)
+| intGe (subject : SpecExpr) (bound : SpecExpr) (loc : SourceRange)
+| intLe (subject : SpecExpr) (bound : SpecExpr) (loc : SourceRange)
 /-- A floating-point literal, stored as a string to preserve precision. -/
-| floatLit (value : String)
-| floatGe (subject : SpecExpr) (bound : SpecExpr)
-| floatLe (subject : SpecExpr) (bound : SpecExpr)
-| enumMember (subject : SpecExpr) (values : Array String)
+| floatLit (value : String) (loc : SourceRange)
+| floatGe (subject : SpecExpr) (bound : SpecExpr) (loc : SourceRange)
+| floatLe (subject : SpecExpr) (bound : SpecExpr) (loc : SourceRange)
+| enumMember (subject : SpecExpr) (values : Array String) (loc : SourceRange)
 /-- `regexMatch subject pattern` asserts that `subject` matches the regular
     expression `pattern`. Corresponds to `compile(pattern).search(subject) is not None`
     in the Python source. -/
-| regexMatch (subject : SpecExpr) (pattern : String)
+| regexMatch (subject : SpecExpr) (pattern : String) (loc : SourceRange)
 /-- `containsKey container key` asserts that `key` is present in `container`.
     Corresponds to `"key" in container` in the Python source. -/
-| containsKey (container : SpecExpr) (key : String)
+| containsKey (container : SpecExpr) (key : String) (loc : SourceRange)
 /-- `implies condition body` asserts that if `condition` holds then `body` holds.
     Used to represent conditional assertions like `if "field" in kwargs: assert ...`. -/
-| implies (condition : SpecExpr) (body : SpecExpr)
+| implies (condition : SpecExpr) (body : SpecExpr) (loc : SourceRange)
 /-- Logical negation. Used for else-branch conditions. -/
-| not (e : SpecExpr)
+| not (e : SpecExpr) (loc : SourceRange)
 /-- `forallList list varName body` asserts that `body` holds for every element
     of `list`, with `varName` bound to each element in turn. Only `body` may
     refer to `varName`. Corresponds to `for varName in list: assert body`. -/
-| forallList (list : SpecExpr) (varName : String) (body : SpecExpr)
+| forallList (list : SpecExpr) (varName : String) (body : SpecExpr) (loc : SourceRange)
 /-- `forallDict dict keyVar valVar body` asserts that `body` holds for every
     key-value pair in `dict`. Both `keyVar` and `valVar` are bound in `body`.
     Corresponds to `for keyVar, valVar in dict.items(): assert body`. -/
-| forallDict (dict : SpecExpr) (keyVar : String) (valVar : String) (body : SpecExpr)
+| forallDict (dict : SpecExpr) (keyVar : String) (valVar : String) (body : SpecExpr) (loc : SourceRange)
 deriving Inhabited
 
 inductive MessagePart where