pyspec: SpecExpr arithmetic/comparison/boolean vocabulary

StrataPython/StrataPython/PythonLaurelTypedExpr.lean

@@ -95,6 +95,10 @@ def or (x y : TypedStmtExpr .TBool)
     (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TBool :=
   .ofStmt (.PrimitiveOp .Or [x.stmt, y.stmt]) source
 
+def and (x y : TypedStmtExpr .TBool)
+    (source : Option FileRange := x.stmt.source) : TypedStmtExpr .TBool :=
+  .ofStmt (.PrimitiveOp .And [x.stmt, y.stmt]) source
+
 abbrev tyDictStrAny : HighType := .UserDefined "DictStrAny"
 
 def anyIsfromNone (v : TypedStmtExpr tyAny)
@@ -105,6 +109,13 @@ def fromInt (v : TypedStmtExpr .TInt)
     (source : Option FileRange := v.stmt.source) : TypedStmtExpr tyAny :=
   .ofStmt (.StaticCall (mkId "from_int") [v.stmt]) source
 
+def fromBool (v : TypedStmtExpr .TBool)
+    (source : Option FileRange := v.stmt.source) : TypedStmtExpr tyAny :=
+  .ofStmt (.StaticCall (mkId "from_bool") [v.stmt]) source
+
+def fromNone (source : Option FileRange := none) : TypedStmtExpr tyAny :=
+  .ofStmt (.StaticCall (mkId "from_None") []) source
+
 def anyAsInt (a : TypedStmtExpr tyAny)
     (source : Option FileRange := a.stmt.source) : TypedStmtExpr .TInt :=
   .ofStmt (.StaticCall (mkId "Any..as_int!") [a.stmt]) source

StrataPython/StrataPython/Specs.lean

@@ -728,25 +728,13 @@ private def makeComparison
     (floatCtor intCtor : SpecExpr → SpecExpr → SpecExpr)
     (subj : SpecExpr) (subjType : SpecType) (bound : SpecExpr) (boundType : SpecType)
     : Option SpecExpr :=
-  if subjType.isFloatType then
+  if subjType.isFloatType || boundType.isFloatType then
+    -- float comparison; coerce an int-literal operand to float
     match bound with
-    | .floatLit .. => some (floatCtor subj bound)
     | .intLit v loc => some (floatCtor subj (.floatLit (toString v) (loc := loc)))
-    | _ => none
-  else if subjType.isIntType then
-    match bound with
-    | .intLit .. => some (intCtor subj bound)
-    | _ => 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
+    | _ => some (floatCtor subj bound)
   else
-    none
+    some (intCtor subj bound)
 
 private def transCompare (loc : SourceRange)
     (lhsExpr : SpecExpr) (lhsType : SpecType)
@@ -777,15 +765,26 @@ private def transCompare (loc : SourceRange)
       return some lhsExpr
     | _, _ => pure ()
   | _ => pure ()
-  -- subject == "A" (single enum value)
+  -- subject == "A" (single enum value), otherwise a general equality `a == b`.
   match ops[0] with
   | .Eq _ =>
     match comparators[0] with
     | .Constant _ (.ConString _ ⟨_, strVal⟩) _ =>
       return some (.enumMember lhsExpr #[strVal] (loc := loc))
-    | _ => pure ()
+    | _ =>
+      let (clean, (rhsExpr, _)) ← runNoWarn (transExpr comparators[0])
+      if clean then
+        return some (.eq lhsExpr rhsExpr (loc := loc))
   | _ => pure ()
 
+  -- `x is None` / `x is not None` -> PEq/PNEq against None.
+  match ops[0], comparators[0] with
+  | .Is _, .Constant _ (.ConNone _) _ =>
+    return some (.pcmp "PEq" lhsExpr (.noneLit loc) (loc := loc))
+  | .IsNot _, .Constant _ (.ConNone _) _ =>
+    return some (.pcmp "PNEq" lhsExpr (.noneLit loc) (loc := loc))
+  | _, _ => pure ()
+
   -- subject >= N / subject <= N (type-checked: int or float)
   let (clean, (boundExpr, boundType)) ← runNoWarn (transExpr comparators[0])
   if not clean then
@@ -796,6 +795,12 @@ 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 some (.pcmp "PGt" lhsExpr boundExpr (loc := loc))
+  | .Lt _ => return some (.pcmp "PLt" lhsExpr boundExpr (loc := loc))
+  | .NotEq _ => return some (.pcmp "PNEq" lhsExpr boundExpr (loc := loc))
+  -- `x in c` / `x not in c` -> PIn/PNotIn membership.
+  | .In _ => return some (.pcmp "PIn" lhsExpr boundExpr (loc := loc))
+  | .NotIn _ => return some (.pcmp "PNotIn" lhsExpr boundExpr (loc := loc))
   | _ =>
     return none
 
@@ -837,11 +842,43 @@ partial def transExpr (e : expr SourceRange)
       else
         specWarning loc s!"subscript subject is not a TypedDict"
     return (.getIndex innerExpr fieldName (loc := loc), fieldTp.getD anyType)
+  -- Attribute access `obj.field` (e.g. `self.x`) -> getIndex(obj, "field").
+  | .Attribute _ value ⟨_, attrName⟩ (.Load _) =>
+    let (valueExpr, _valueTp) ← transExpr value
+    return (.getIndex valueExpr attrName (loc := loc), anyType)
+  -- Binary arithmetic `lhs <op> rhs` -> the matching arithmetic ctor.
+  | .BinOp _ lhs op rhs => do
+    let binArith (ctor : SpecExpr → SpecExpr → SourceRange → SpecExpr)
+        : SpecAssertionM (SpecExpr × SpecType) := do
+      let (lclean, (lhsExpr, _)) ← runNoWarn (transExpr lhs)
+      let (rclean, (rhsExpr, _)) ← runNoWarn (transExpr rhs)
+      if lclean && rclean then
+        return (ctor lhsExpr rhsExpr loc, intType)
+      specWarning loc s!"unsupported binary operation: {eformat e.toAst}"
+      return placeholder
+    match op with
+    | .Add _ => binArith (.add · · ·)
+    | .Sub _ => binArith (.sub · · ·)
+    | .Mult _ => binArith (.mul · · ·)
+    | .FloorDiv _ => binArith (.floorDiv · · ·)
+    | .Mod _ => binArith (.mod · · ·)
+    | .Pow _ => binArith (.pow · · ·)
+    -- TODO: LShift/RShift and bitwise ops (BitAnd/BitOr/BitXor, Invert) unsupported.
+    | _ =>
+      specWarning loc s!"unsupported expression: {eformat e.toAst}"
+      return placeholder
   -- Integer literal
   | .Constant _ (.ConPos _ ⟨_, n⟩) _ =>
     return (.intLit (Int.ofNat n) (loc := loc), intType)
   | .Constant _ (.ConNeg _ ⟨_, n⟩) _ =>
     return (.intLit (Int.negOfNat n) (loc := loc), intType)
+  -- Boolean / None literals
+  | .Constant _ (.ConTrue _) _ =>
+    return (.boolLit true (loc := loc), boolType)
+  | .Constant _ (.ConFalse _) _ =>
+    return (.boolLit false (loc := loc), boolType)
+  | .Constant _ (.ConNone _) _ =>
+    return (.noneLit (loc := loc), SpecType.noneType loc)
   -- Float literal
   | .Constant _ (.ConFloat _ ⟨_, s⟩) _ =>
     return (.floatLit s (loc := loc), floatType)
@@ -850,6 +887,13 @@ 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)
+  -- UnaryOp(USub) on a non-literal -> `.neg` (PNeg); placed after the literal arms above.
+  | .UnaryOp _ (.USub _) operand =>
+    let (clean, (operandExpr, operandTp)) ← runNoWarn (transExpr operand)
+    if clean then
+      return (.neg operandExpr (loc := loc), operandTp)
+    specWarning loc s!"unsupported unary minus: {eformat e.toAst}"
+    return placeholder
   -- String literal (extract value for use in messages)
   | .Constant _ (.ConString _ ⟨_, _s⟩) _ =>
     return (.placeholder (loc := loc), SpecType.ident loc .builtinsStr)
@@ -906,7 +950,19 @@ partial def transExpr (e : expr SourceRange)
       | none => pure ()
     specWarning loc s!"unsupported comparison: {eformat e.toAst}"
     return placeholder
-  -- BoolOp(Or): try to merge enum values
+  -- BoolOp(And): general conjunction
+  | .BoolOp _ (.And _) ⟨_, values⟩ =>
+    let (clean, branches) ← runNoWarn <|
+      values.attach.mapM fun ⟨v, _⟩ => transExpr v
+    if clean then
+      match branches.toList with
+      | (b0, _) :: rest =>
+        let conj := rest.foldl (init := b0) fun acc (b, _) => SpecExpr.and acc b (loc := loc)
+        return (conj, boolType)
+      | [] => pure ()
+    specWarning loc s!"unsupported and-expression: {eformat e.toAst}"
+    return placeholder
+  -- BoolOp(Or): try to merge enum values, else general disjunction
   | .BoolOp _ (.Or _) ⟨_, values⟩ =>
     let (clean, branches) ← runNoWarn <|
       values.attach.mapM fun ⟨v, _⟩ => transExpr v
@@ -924,8 +980,20 @@ partial def transExpr (e : expr SourceRange)
         | _ => none
       if let some (subj, vals) := merged then
         return (.enumMember subj vals (loc := loc), boolType)
+      match branches.toList with
+      | (b0, _) :: rest =>
+        let disj := rest.foldl (init := b0) fun acc (b, _) => SpecExpr.or acc b (loc := loc)
+        return (disj, boolType)
+      | [] => pure ()
     specWarning loc s!"unsupported or-expression: {eformat e.toAst}"
     return placeholder
+  -- UnaryOp(Not): boolean negation
+  | .UnaryOp _ (.Not _) operand =>
+    let (clean, (inner, _)) ← runNoWarn (transExpr operand)
+    if clean then
+      return (.not inner (loc := loc), boolType)
+    specWarning loc s!"unsupported not-expression: {eformat e.toAst}"
+    return placeholder
   | _ =>
     specWarning loc s!"unsupported expression: {eformat e.toAst}"
     return placeholder

StrataPython/StrataPython/Specs/DDM.lean

@@ -68,7 +68,26 @@ op mkKwargsDecl(name : Ident, kwargsType : SpecType) : KwargsDecl =>
 
 category SpecExprDecl;
 op placeholderExpr() : SpecExprDecl => "placeholder";
+op trueLitExpr() : SpecExprDecl => "trueLit";
+op falseLitExpr() : SpecExprDecl => "falseLit";
+op noneLitExpr() : SpecExprDecl => "noneLit";
 op varExpr(name : Ident) : SpecExprDecl => name;
+op eqExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  "eq" "(" lhs ", " rhs ")";
+op addExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  "add" "(" lhs ", " rhs ")";
+op subExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  "sub" "(" lhs ", " rhs ")";
+op mulExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  "mul" "(" lhs ", " rhs ")";
+op floorDivExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  "floorDiv" "(" lhs ", " rhs ")";
+op modExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  "mod" "(" lhs ", " rhs ")";
+op powExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  "pow" "(" lhs ", " rhs ")";
+op negExpr(operand : SpecExprDecl) : SpecExprDecl =>
+  "neg" "(" operand ")";
 op getIndexExpr(subject : SpecExprDecl, field : Ident) : SpecExprDecl =>
   @[prec(50)] subject "[" field "]";
 op isInstanceOfExpr(subject : SpecExprDecl, typeName : Str) : SpecExprDecl =>
@@ -97,6 +116,12 @@ op impliesExpr(condition : SpecExprDecl, body : SpecExprDecl) : SpecExprDecl =>
   @[prec(10), rightassoc] condition " => " body;
 op notExpr(e : SpecExprDecl) : SpecExprDecl =>
   "not" "(" e ")";
+op andExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  "and" "(" lhs ", " rhs ")";
+op orExpr(lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  "or" "(" lhs ", " rhs ")";
+op pcmpExpr(opName : Str, lhs : SpecExprDecl, rhs : SpecExprDecl) : SpecExprDecl =>
+  "pcmp" "(" opName ", " lhs ", " rhs ")";
 op forallListExpr(list : SpecExprDecl, varName : Ident, body : SpecExprDecl) : SpecExprDecl =>
   "forall" "(" list ", " varName ", " body ")";
 op forallDictExpr(dict : SpecExprDecl, keyVar : Ident,
@@ -263,7 +288,17 @@ def Arg.toDDM (d : Arg) : DDM.ArgDecl SourceRange :=
 protected def SpecExpr.toDDM (e : SpecExpr) : DDM.SpecExprDecl SourceRange :=
   match e with
   | .placeholder loc => .placeholderExpr loc
+  | .boolLit b loc => if b then .trueLitExpr loc else .falseLitExpr loc
+  | .noneLit loc => .noneLitExpr loc
   | .var name loc => .varExpr loc ⟨loc, name⟩
+  | .eq lhs rhs loc => .eqExpr loc lhs.toDDM rhs.toDDM
+  | .add lhs rhs loc => .addExpr loc lhs.toDDM rhs.toDDM
+  | .sub lhs rhs loc => .subExpr loc lhs.toDDM rhs.toDDM
+  | .mul lhs rhs loc => .mulExpr loc lhs.toDDM rhs.toDDM
+  | .floorDiv lhs rhs loc => .floorDivExpr loc lhs.toDDM rhs.toDDM
+  | .mod lhs rhs loc => .modExpr loc lhs.toDDM rhs.toDDM
+  | .pow lhs rhs loc => .powExpr loc lhs.toDDM rhs.toDDM
+  | .neg operand loc => .negExpr loc operand.toDDM
   | .getIndex subj field loc => .getIndexExpr loc subj.toDDM ⟨loc, field⟩
   | .isInstanceOf subj tn loc => .isInstanceOfExpr loc subj.toDDM ⟨loc, tn⟩
   | .stringLen subj loc => .stringLenExpr loc subj.toDDM
@@ -283,6 +318,9 @@ protected def SpecExpr.toDDM (e : SpecExpr) : DDM.SpecExprDecl SourceRange :=
   | .implies cond body loc =>
     .impliesExpr loc cond.toDDM body.toDDM
   | .not e loc => .notExpr loc e.toDDM
+  | .and lhs rhs loc => .andExpr loc lhs.toDDM rhs.toDDM
+  | .or lhs rhs loc => .orExpr loc lhs.toDDM rhs.toDDM
+  | .pcmp op lhs rhs loc => .pcmpExpr loc ⟨loc, op⟩ lhs.toDDM rhs.toDDM
   | .forallList list varName body loc =>
     .forallListExpr loc list.toDDM ⟨loc, varName⟩ body.toDDM
   | .forallDict dict keyVar valVar body loc =>
@@ -419,7 +457,18 @@ def DDM.ArgDecl.fromDDM (d : DDM.ArgDecl SourceRange) : FromDDM Specs.Arg := do
 def DDM.SpecExprDecl.fromDDM (d : DDM.SpecExprDecl SourceRange) : Specs.SpecExpr :=
   match d with
   | .placeholderExpr loc => .placeholder loc
+  | .trueLitExpr loc => .boolLit true loc
+  | .falseLitExpr loc => .boolLit false loc
+  | .noneLitExpr loc => .noneLit loc
   | .varExpr loc ⟨_, name⟩ => .var name loc
+  | .eqExpr loc lhs rhs => .eq lhs.fromDDM rhs.fromDDM loc
+  | .addExpr loc lhs rhs => .add lhs.fromDDM rhs.fromDDM loc
+  | .subExpr loc lhs rhs => .sub lhs.fromDDM rhs.fromDDM loc
+  | .mulExpr loc lhs rhs => .mul lhs.fromDDM rhs.fromDDM loc
+  | .floorDivExpr loc lhs rhs => .floorDiv lhs.fromDDM rhs.fromDDM loc
+  | .modExpr loc lhs rhs => .mod lhs.fromDDM rhs.fromDDM loc
+  | .powExpr loc lhs rhs => .pow lhs.fromDDM rhs.fromDDM loc
+  | .negExpr loc operand => .neg operand.fromDDM loc
   | .getIndexExpr loc subj ⟨_, field⟩ => .getIndex subj.fromDDM field loc
   | .isInstanceOfExpr loc subj ⟨_, tn⟩ => .isInstanceOf subj.fromDDM tn loc
   | .lenExpr loc subj => .stringLen subj.fromDDM loc
@@ -435,6 +484,9 @@ def DDM.SpecExprDecl.fromDDM (d : DDM.SpecExprDecl SourceRange) : Specs.SpecExpr
   | .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
+  | .andExpr loc lhs rhs => .and lhs.fromDDM rhs.fromDDM loc
+  | .orExpr loc lhs rhs => .or lhs.fromDDM rhs.fromDDM loc
+  | .pcmpExpr loc ⟨_, op⟩ lhs rhs => .pcmp op lhs.fromDDM rhs.fromDDM loc
   | .forallListExpr loc list ⟨_, varName⟩ body =>
     .forallList list.fromDDM varName body.fromDDM loc
   | .forallDictExpr loc dict ⟨_, keyVar⟩ ⟨_, valVar⟩ body =>

StrataPython/StrataPython/Specs/Decls.lean

@@ -486,8 +486,28 @@ inductive SpecExpr where
     Used in preconditions like `assert len(name) >= 1`. -/
 | stringLen (subject : SpecExpr) (loc : SourceRange)
 | intLit (value : Int) (loc : SourceRange)
+| boolLit (value : Bool) (loc : SourceRange)
+| noneLit (loc : SourceRange)
 | intGe (subject : SpecExpr) (bound : SpecExpr) (loc : SourceRange)
 | intLe (subject : SpecExpr) (bound : SpecExpr) (loc : SourceRange)
+/-- General equality `lhs == rhs` (non-enum). -/
+| eq (lhs : SpecExpr) (rhs : SpecExpr) (loc : SourceRange)
+/-- Addition `lhs + rhs`. -/
+| add (lhs : SpecExpr) (rhs : SpecExpr) (loc : SourceRange)
+/-- Subtraction / multiplication. -/
+| sub (lhs : SpecExpr) (rhs : SpecExpr) (loc : SourceRange)
+| mul (lhs : SpecExpr) (rhs : SpecExpr) (loc : SourceRange)
+| floorDiv (lhs : SpecExpr) (rhs : SpecExpr) (loc : SourceRange)
+| mod (lhs : SpecExpr) (rhs : SpecExpr) (loc : SourceRange)
+/-- Exponentiation `lhs ** rhs`. -/
+| pow (lhs : SpecExpr) (rhs : SpecExpr) (loc : SourceRange)
+/-- Unary minus `-operand` (non-literal; negative literals fold to int/float). -/
+| neg (operand : SpecExpr) (loc : SourceRange)
+/-- Boolean conjunction / disjunction. -/
+| and (lhs : SpecExpr) (rhs : SpecExpr) (loc : SourceRange)
+| or (lhs : SpecExpr) (rhs : SpecExpr) (loc : SourceRange)
+/-- Runtime comparison `op(lhs, rhs)`; `op` is a prelude name (`PLt`/`PGt`/`PNEq`). -/
+| pcmp (op : String) (lhs : SpecExpr) (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)
@@ -525,6 +545,19 @@ 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₂
+  | .eq l₁ r₁ _, .eq l₂ r₂ _ => l₁.softBEq l₂ && r₁.softBEq r₂
+  | .add l₁ r₁ _, .add l₂ r₂ _ => l₁.softBEq l₂ && r₁.softBEq r₂
+  | .sub l₁ r₁ _, .sub l₂ r₂ _ => l₁.softBEq l₂ && r₁.softBEq r₂
+  | .mul l₁ r₁ _, .mul l₂ r₂ _ => l₁.softBEq l₂ && r₁.softBEq r₂
+  | .floorDiv l₁ r₁ _, .floorDiv l₂ r₂ _ => l₁.softBEq l₂ && r₁.softBEq r₂
+  | .mod l₁ r₁ _, .mod l₂ r₂ _ => l₁.softBEq l₂ && r₁.softBEq r₂
+  | .pow l₁ r₁ _, .pow l₂ r₂ _ => l₁.softBEq l₂ && r₁.softBEq r₂
+  | .neg o₁ _, .neg o₂ _ => o₁.softBEq o₂
+  | .and l₁ r₁ _, .and l₂ r₂ _ => l₁.softBEq l₂ && r₁.softBEq r₂
+  | .or l₁ r₁ _, .or l₂ r₂ _ => l₁.softBEq l₂ && r₁.softBEq r₂
+  | .pcmp op₁ l₁ r₁ _, .pcmp op₂ l₂ r₂ _ => op₁ == op₂ && l₁.softBEq l₂ && r₁.softBEq r₂
+  | .boolLit b₁ _, .boolLit b₂ _ => b₁ == b₂
+  | .noneLit _, .noneLit _ => true
   | .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₂