pyspec: extend predicate bodies with strict comparisons, equality, and integer arithmetic

Strata/Languages/Python/PythonLaurelTypedExpr.lean

@@ -74,6 +74,45 @@ def intLeq (x y : TypedStmtExpr .TInt)
     (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TBool :=
   .ofStmt (.PrimitiveOp .Leq [x.stmt, y.stmt]) source
 
+def intGt (x y : TypedStmtExpr .TInt)
+    (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TBool :=
+  .ofStmt (.PrimitiveOp .Gt [x.stmt, y.stmt]) source
+
+def intLt (x y : TypedStmtExpr .TInt)
+    (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TBool :=
+  .ofStmt (.PrimitiveOp .Lt [x.stmt, y.stmt]) source
+
+def intEq (x y : TypedStmtExpr .TInt)
+    (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TBool :=
+  .ofStmt (.PrimitiveOp .Eq [x.stmt, y.stmt]) source
+
+def intNe (x y : TypedStmtExpr .TInt)
+    (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TBool :=
+  .ofStmt (.PrimitiveOp .Neq [x.stmt, y.stmt]) source
+
+def intAdd (x y : TypedStmtExpr .TInt)
+    (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TInt :=
+  .ofStmt (.PrimitiveOp .Add [x.stmt, y.stmt]) source
+
+def intSub (x y : TypedStmtExpr .TInt)
+    (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TInt :=
+  .ofStmt (.PrimitiveOp .Sub [x.stmt, y.stmt]) source
+
+def intMul (x y : TypedStmtExpr .TInt)
+    (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TInt :=
+  .ofStmt (.PrimitiveOp .Mul [x.stmt, y.stmt]) source
+
+/-- Python `//` is floor (Euclidean) division → Laurel `Div`. -/
+def intFloorDiv (x y : TypedStmtExpr .TInt)
+    (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TInt :=
+  .ofStmt (.PrimitiveOp .Div [x.stmt, y.stmt]) source
+
+/-- Python `%` is Euclidean modulus → Laurel `Mod`. -/
+def intMod (x y : TypedStmtExpr .TInt)
+    (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TInt :=
+  .ofStmt (.PrimitiveOp .Mod [x.stmt, y.stmt]) source
+
+
 def realGeq (x y : TypedStmtExpr .TReal)
     (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TBool :=
   .ofStmt (.PrimitiveOp .Geq [x.stmt, y.stmt]) source
@@ -82,6 +121,22 @@ def realLeq (x y : TypedStmtExpr .TReal)
     (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TBool :=
   .ofStmt (.PrimitiveOp .Leq [x.stmt, y.stmt]) source
 
+def realGt (x y : TypedStmtExpr .TReal)
+    (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TBool :=
+  .ofStmt (.PrimitiveOp .Gt [x.stmt, y.stmt]) source
+
+def realLt (x y : TypedStmtExpr .TReal)
+    (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TBool :=
+  .ofStmt (.PrimitiveOp .Lt [x.stmt, y.stmt]) source
+
+def realEq (x y : TypedStmtExpr .TReal)
+    (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TBool :=
+  .ofStmt (.PrimitiveOp .Eq [x.stmt, y.stmt]) source
+
+def realNe (x y : TypedStmtExpr .TReal)
+    (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TBool :=
+  .ofStmt (.PrimitiveOp .Neq [x.stmt, y.stmt]) source
+
 def not (x : TypedStmtExpr .TBool)
     (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TBool :=
   .ofStmt (.PrimitiveOp .Not [x.stmt]) source

Strata/Languages/Python/Specs.lean

@@ -802,15 +802,23 @@ private def makeComparison
   else if subjType.isIntType then
     match bound with
     | .intLit .. => some (intCtor subj bound)
-    | _ => none
+    | _ =>
+      -- subj is int, bound is int- or Any-typed (a parameter etc.) →
+      -- use the int comparison constructor.
+      if boundType.isIntType || boundType.isAnyType then some (intCtor subj bound) else none
   else if boundType.isFloatType then
     match bound with
     | .floatLit .. => some (floatCtor subj bound)
     | _ => none
   else if boundType.isIntType then
     match bound with
     | .intLit .. => some (intCtor subj bound)
-    | _ => none
+    | _ =>
+      if subjType.isIntType || subjType.isAnyType then some (intCtor subj bound) else none
+  else if subjType.isAnyType && boundType.isAnyType then
+    -- Both sides Any-typed (e.g. two function parameters whose types
+    -- aren't seeded). Trust the user and use the int comparison.
+    some (intCtor subj bound)
   else
     none
 
@@ -825,7 +833,7 @@ private def transCompare (loc : SourceRange)
   let .isTrue h₂ := decideProp (comparators.size = 1)
     | return none
   match lhsExpr with
-  -- len(subject) >= N / len(subject) <= N
+  -- len(subject) >= N / len(subject) <= N / len(subject) > N / len(subject) < N
   | .stringLen _ _ =>
     let (clean, (boundExpr, _)) ← runNoWarn (transExpr comparators[0])
     if clean then
@@ -834,6 +842,8 @@ private def transCompare (loc : SourceRange)
         match ops[0] with
         | .GtE _ => return some (.intGe lhsExpr boundExpr (loc := loc))
         | .LtE _ => return some (.intLe lhsExpr boundExpr (loc := loc))
+        | .Gt _ => return some (.intGt lhsExpr boundExpr (loc := loc))
+        | .Lt _ => return some (.intLt lhsExpr boundExpr (loc := loc))
         | _ => pure ()
       | _ => pure ()
     -- compile("pattern").search(subject) is not None
@@ -852,7 +862,7 @@ private def transCompare (loc : SourceRange)
     | _ => pure ()
   | _ => pure ()
 
-  -- subject >= N / subject <= N (type-checked: int or float)
+  -- subject {>=, <=, >, <, ==, !=} N (type-checked: int or float)
   let (clean, (boundExpr, boundType)) ← runNoWarn (transExpr comparators[0])
   if not clean then
     return none
@@ -862,6 +872,14 @@ private def transCompare (loc : SourceRange)
     return makeComparison (.floatGe · · (loc := loc)) (.intGe · · (loc := loc)) lhsExpr lhsType boundExpr boundType
   | .LtE _ =>
     return makeComparison (.floatLe · · (loc := loc)) (.intLe · · (loc := loc)) lhsExpr lhsType boundExpr boundType
+  | .Gt _ =>
+    return makeComparison (.floatGt · · (loc := loc)) (.intGt · · (loc := loc)) lhsExpr lhsType boundExpr boundType
+  | .Lt _ =>
+    return makeComparison (.floatLt · · (loc := loc)) (.intLt · · (loc := loc)) lhsExpr lhsType boundExpr boundType
+  | .Eq _ =>
+    return makeComparison (.floatEq · · (loc := loc)) (.intEq · · (loc := loc)) lhsExpr lhsType boundExpr boundType
+  | .NotEq _ =>
+    return makeComparison (.floatNe · · (loc := loc)) (.intNe · · (loc := loc)) lhsExpr lhsType boundExpr boundType
   | _ =>
     return none
 
@@ -916,6 +934,29 @@ partial def transExpr (e : expr SourceRange)
     return (.intLit (Int.negOfNat n) (loc := loc), intType)
   | .UnaryOp _ (.USub _) (.Constant _ (.ConFloat _ ⟨_, s⟩) _) =>
     return (.floatLit s!"-{s}" (loc := loc), floatType)
+  -- Integer arithmetic in predicate bodies. Both sides must translate
+  -- cleanly to int-typed (or Any-typed, e.g. function parameters whose
+  -- types weren't seeded into `localTypes`) SpecExprs; the result is
+  -- int-typed. Anything other than int / Any (float, str, …) falls
+  -- through to placeholder so transExpr doesn't claim a type it can't
+  -- justify.
+  | .BinOp _ left op right =>
+    let (lhs, lhsTp) ← transExpr left
+    let (rhs, rhsTp) ← transExpr right
+    let okType (tp : SpecType) : Bool := tp.isIntType || tp.isAnyType
+    if okType lhsTp && okType rhsTp then
+      match op with
+      | .Add _ => return (.intAdd lhs rhs (loc := loc), intType)
+      | .Sub _ => return (.intSub lhs rhs (loc := loc), intType)
+      | .Mult _ => return (.intMul lhs rhs (loc := loc), intType)
+      | .FloorDiv _ => return (.intDiv lhs rhs (loc := loc), intType)
+      | .Mod _ => return (.intMod lhs rhs (loc := loc), intType)
+      | _ =>
+        specWarning loc s!"unsupported BinOp in predicate: {repr op}"
+        return placeholder
+    else
+      specWarning loc s!"BinOp with non-int operand: lhs={repr lhsTp}, rhs={repr rhsTp}"
+      return placeholder
   -- String literal (extract value for use in messages)
   | .Constant _ (.ConString _ ⟨_, _s⟩) _ =>
     return (.placeholder (loc := loc), SpecType.ident loc .builtinsStr)

Strata/Languages/Python/Specs/DDM.lean

@@ -87,11 +87,37 @@ op intGeExpr(subject : SpecExprDecl, bound : SpecExprDecl) : SpecExprDecl =>
   @[prec(15)] subject " >=_int " bound;
 op intLeExpr(subject : SpecExprDecl, bound : SpecExprDecl) : SpecExprDecl =>
   @[prec(15)] subject " <=_int " bound;
+op intGtExpr(subject : SpecExprDecl, bound : SpecExprDecl) : SpecExprDecl =>
+  @[prec(15)] subject " >_int " bound;
+op intLtExpr(subject : SpecExprDecl, bound : SpecExprDecl) : SpecExprDecl =>
+  @[prec(15)] subject " <_int " bound;
+op intEqExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  @[prec(15)] lhs " ==_int " rhs;
+op intNeExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  @[prec(15)] lhs " !=_int " rhs;
+op intAddExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  @[prec(20), leftassoc] lhs " +_int " rhs;
+op intSubExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  @[prec(20), leftassoc] lhs " -_int " rhs;
+op intMulExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  @[prec(25), leftassoc] lhs " *_int " rhs;
+op intDivExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  @[prec(25), leftassoc] lhs " //_int " rhs;
+op intModExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  @[prec(25), leftassoc] lhs " %_int " rhs;
 op floatExpr(value : Str) : SpecExprDecl => value;
 op floatGeExpr(subject : SpecExprDecl, bound : SpecExprDecl) : SpecExprDecl =>
   @[prec(15)] subject " >=_float " bound;
 op floatLeExpr(subject : SpecExprDecl, bound : SpecExprDecl) : SpecExprDecl =>
   @[prec(15)] subject " <=_float " bound;
+op floatGtExpr(subject : SpecExprDecl, bound : SpecExprDecl) : SpecExprDecl =>
+  @[prec(15)] subject " >_float " bound;
+op floatLtExpr(subject : SpecExprDecl, bound : SpecExprDecl) : SpecExprDecl =>
+  @[prec(15)] subject " <_float " bound;
+op floatEqExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  @[prec(15)] lhs " ==_float " rhs;
+op floatNeExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  @[prec(15)] lhs " !=_float " rhs;
 op enumMemberExpr(subject : SpecExprDecl, values : Seq Str) : SpecExprDecl =>
   "enum" "(" subject ", [" values "]" ")";
 op regexMatchExpr(subject : SpecExprDecl, pattern : Str) : SpecExprDecl =>
@@ -275,9 +301,22 @@ protected def SpecExpr.toDDM (e : SpecExpr) : DDM.SpecExprDecl SourceRange :=
   | .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
+  | .intGt subj bound loc => .intGtExpr loc subj.toDDM bound.toDDM
+  | .intLt subj bound loc => .intLtExpr loc subj.toDDM bound.toDDM
+  | .intEq l r loc => .intEqExpr loc l.toDDM r.toDDM
+  | .intNe l r loc => .intNeExpr loc l.toDDM r.toDDM
+  | .intAdd l r loc => .intAddExpr loc l.toDDM r.toDDM
+  | .intSub l r loc => .intSubExpr loc l.toDDM r.toDDM
+  | .intMul l r loc => .intMulExpr loc l.toDDM r.toDDM
+  | .intDiv l r loc => .intDivExpr loc l.toDDM r.toDDM
+  | .intMod l r loc => .intModExpr loc l.toDDM r.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
+  | .floatGt subj bound loc => .floatGtExpr loc subj.toDDM bound.toDDM
+  | .floatLt subj bound loc => .floatLtExpr loc subj.toDDM bound.toDDM
+  | .floatEq l r loc => .floatEqExpr loc l.toDDM r.toDDM
+  | .floatNe l r loc => .floatNeExpr loc l.toDDM r.toDDM
   | .enumMember subj values loc =>
     .enumMemberExpr loc subj.toDDM
       ⟨loc, values.map (⟨loc, ·⟩)⟩
@@ -416,9 +455,22 @@ private def DDM.SpecExprDecl.fromDDM (d : DDM.SpecExprDecl SourceRange) : Specs.
   | .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
+  | .intGtExpr loc subj bound => .intGt subj.fromDDM bound.fromDDM loc
+  | .intLtExpr loc subj bound => .intLt subj.fromDDM bound.fromDDM loc
+  | .intEqExpr loc l r => .intEq l.fromDDM r.fromDDM loc
+  | .intNeExpr loc l r => .intNe l.fromDDM r.fromDDM loc
+  | .intAddExpr loc l r => .intAdd l.fromDDM r.fromDDM loc
+  | .intSubExpr loc l r => .intSub l.fromDDM r.fromDDM loc
+  | .intMulExpr loc l r => .intMul l.fromDDM r.fromDDM loc
+  | .intDivExpr loc l r => .intDiv l.fromDDM r.fromDDM loc
+  | .intModExpr loc l r => .intMod l.fromDDM r.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
+  | .floatGtExpr loc subj bound => .floatGt subj.fromDDM bound.fromDDM loc
+  | .floatLtExpr loc subj bound => .floatLt subj.fromDDM bound.fromDDM loc
+  | .floatEqExpr loc l r => .floatEq l.fromDDM r.fromDDM loc
+  | .floatNeExpr loc l r => .floatNe l.fromDDM r.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

Strata/Languages/Python/Specs/Decls.lean

@@ -387,6 +387,8 @@ def asIdent (tp : SpecType) : Option PythonIdent := do
 
 def isIntType (tp : SpecType) : Bool := tp.asIdent == some .builtinsInt
 
+def isAnyType (tp : SpecType) : Bool := tp.asIdent == some .typingAny
+
 def isFloatType (tp : SpecType) : Bool := tp.asIdent == some .builtinsFloat
 
 def isStringType (tp : SpecType) : Bool := tp.asIdent == some .builtinsStr
@@ -488,10 +490,25 @@ inductive SpecExpr where
 | intLit (value : Int) (loc : SourceRange)
 | intGe (subject : SpecExpr) (bound : SpecExpr) (loc : SourceRange)
 | intLe (subject : SpecExpr) (bound : SpecExpr) (loc : SourceRange)
+| intGt (subject : SpecExpr) (bound : SpecExpr) (loc : SourceRange)
+| intLt (subject : SpecExpr) (bound : SpecExpr) (loc : SourceRange)
+| intEq (lhs : SpecExpr) (rhs : SpecExpr) (loc : SourceRange)
+| intNe (lhs : SpecExpr) (rhs : SpecExpr) (loc : SourceRange)
+| intAdd (lhs rhs : SpecExpr) (loc : SourceRange)
+| intSub (lhs rhs : SpecExpr) (loc : SourceRange)
+| intMul (lhs rhs : SpecExpr) (loc : SourceRange)
+/-- Python `//` (floor division). -/
+| intDiv (lhs rhs : SpecExpr) (loc : SourceRange)
+/-- Python `%` (modulus). -/
+| intMod (lhs rhs : SpecExpr) (loc : SourceRange)
 /-- A floating-point literal, stored as a string to preserve precision. -/
 | floatLit (value : String) (loc : SourceRange)
 | floatGe (subject : SpecExpr) (bound : SpecExpr) (loc : SourceRange)
 | floatLe (subject : SpecExpr) (bound : SpecExpr) (loc : SourceRange)
+| floatGt (subject : SpecExpr) (bound : SpecExpr) (loc : SourceRange)
+| floatLt (subject : SpecExpr) (bound : SpecExpr) (loc : SourceRange)
+| floatEq (lhs : SpecExpr) (rhs : SpecExpr) (loc : SourceRange)
+| floatNe (lhs : SpecExpr) (rhs : 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`
@@ -525,9 +542,22 @@ def SpecExpr.softBEq : SpecExpr → SpecExpr → Bool
   | .intLit v₁ _, .intLit v₂ _ => v₁ == v₂
   | .intGe s₁ b₁ _, .intGe s₂ b₂ _ => s₁.softBEq s₂ && b₁.softBEq b₂
   | .intLe s₁ b₁ _, .intLe s₂ b₂ _ => s₁.softBEq s₂ && b₁.softBEq b₂
+  | .intGt s₁ b₁ _, .intGt s₂ b₂ _ => s₁.softBEq s₂ && b₁.softBEq b₂
+  | .intLt s₁ b₁ _, .intLt s₂ b₂ _ => s₁.softBEq s₂ && b₁.softBEq b₂
+  | .intEq a₁ b₁ _, .intEq a₂ b₂ _ => a₁.softBEq a₂ && b₁.softBEq b₂
+  | .intNe a₁ b₁ _, .intNe a₂ b₂ _ => a₁.softBEq a₂ && b₁.softBEq b₂
+  | .intAdd a₁ b₁ _, .intAdd a₂ b₂ _ => a₁.softBEq a₂ && b₁.softBEq b₂
+  | .intSub a₁ b₁ _, .intSub a₂ b₂ _ => a₁.softBEq a₂ && b₁.softBEq b₂
+  | .intMul a₁ b₁ _, .intMul a₂ b₂ _ => a₁.softBEq a₂ && b₁.softBEq b₂
+  | .intDiv a₁ b₁ _, .intDiv a₂ b₂ _ => a₁.softBEq a₂ && b₁.softBEq b₂
+  | .intMod a₁ b₁ _, .intMod a₂ b₂ _ => a₁.softBEq a₂ && b₁.softBEq b₂
   | .floatLit v₁ _, .floatLit v₂ _ => v₁ == v₂
   | .floatGe s₁ b₁ _, .floatGe s₂ b₂ _ => s₁.softBEq s₂ && b₁.softBEq b₂
   | .floatLe s₁ b₁ _, .floatLe s₂ b₂ _ => s₁.softBEq s₂ && b₁.softBEq b₂
+  | .floatGt s₁ b₁ _, .floatGt s₂ b₂ _ => s₁.softBEq s₂ && b₁.softBEq b₂
+  | .floatLt s₁ b₁ _, .floatLt s₂ b₂ _ => s₁.softBEq s₂ && b₁.softBEq b₂
+  | .floatEq a₁ b₁ _, .floatEq a₂ b₂ _ => a₁.softBEq a₂ && b₁.softBEq b₂
+  | .floatNe a₁ b₁ _, .floatNe a₂ b₂ _ => a₁.softBEq a₂ && b₁.softBEq b₂
   | .enumMember s₁ v₁ _, .enumMember s₂ v₂ _ => s₁.softBEq s₂ && v₁ == v₂
   | .regexMatch s₁ p₁ _, .regexMatch s₂ p₂ _ => s₁.softBEq s₂ && p₁ == p₂
   | .containsKey c₁ k₁ _, .containsKey c₂ k₂ _ => c₁.softBEq c₂ && k₁ == k₂