Parse and lift the SMT counterexamples to LExpr (strata-org#495)

This patch

- Parses the output of SMT solver containing counterexample using
SMTDDM,
- Lifts SMTDDM.Term to SMT.Term, then to LExpr, and
- Let Core.verify properly print the counterexample using the new
printer of Core.

To implement this,
- A new `parseCategoryFromDialect` function is added to DDM. It parses a
string using a specific syntax category.
- A bug in SpaceSepBy making only odd'th symbols to be parsed was fixed.
It seems SpaceSepBy was only used by SMTDDM.
- The old CexParser.lean is erased.

A new StrataTest/Languages/Core/CounterExampleLiftTest.lean test shows a
handful of interesting results.

---

strata-org#482 and this pull request have overlapping edits. Merging strata-org#482 does not
introduce merge conflict, but tests need updates. Would be great if strata-org#482
is reviewed first! :)

---

Thanks to Kiro :)

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

Author: aqjune-aws

Strata/DDM/Elab.lean

@@ -412,5 +412,57 @@ def parseStrataProgramFromDialect (dialects : LoadedDialects) (dialect : Dialect
         return s!"{msg}  {e.pos.line}:{e.pos.column}: {← e.data.toString}\n"
       throw (IO.userError errMsg)
 
+/--
+Parse a single syntax category from a loaded dialect set.
+
+Unlike `parseStrataProgramFromDialect`, this targets a specific syntax
+category (e.g. `q\`SMTResponse.GetValueResponse`) rather than parsing a
+full program of `Command`s. This avoids ambiguity issues that arise when
+multiple categories share the same surface syntax.
+
+Returns the elaborated `Operation` on success.
+-/
+def parseCategoryFromDialect
+    (dialects : LoadedDialects) (cat : QualifiedIdent) (input : InputContext)
+    (startPos : String.Pos.Raw := 0)
+    (stopPos : String.Pos.Raw := input.endPos)
+    : IO Strata.Operation := do
+  let leanEnv ← Lean.mkEmptyEnvironment 0
+  -- Open dialects using the same pattern as elabProgramRest: start from
+  -- initDeclState (which has Init open) and open each dialect via
+  -- ensureLoaded! to avoid panics from HashMap iteration order.
+  let s := DeclState.initDeclState
+  let s := dialects.dialects.toList.foldl (init := s) fun s d =>
+    DeclState.ensureLoaded! s dialects d.name
+  -- Run the category-level parser
+  let ps := Parser.runCatParser s.tokenTable s.parserMap leanEnv input startPos stopPos cat
+  if ps.hasError then
+    let errMsg ← ps.allErrors.foldlM (init := "Parse errors:\n") fun msg (pos, stk, err) =>
+      let m := Lean.mkErrorMessage input pos stk err
+      return s!"{msg}  {m.pos.line}:{m.pos.column}: {← m.data.toString}\n"
+    throw (IO.userError errMsg)
+  if ps.stxStack.size == 0 then
+    throw (IO.userError "Parser returned no syntax")
+  let stx := ps.stxStack.back
+  -- Elaborate the parsed syntax into an Operation
+  let ctx : ElabContext := {
+    dialects := dialects.dialects
+    syntaxElabs := dialects.syntaxElabMap
+    openDialectSet := s.openDialectSet
+    typeOrCatDeclMap := s.typeOrCatDeclMap
+    metadataDeclMap := s.metadataDeclMap
+    globalContext := s.globalContext
+    inputContext := input
+    missingImport := false
+  }
+  let es : ElabState := {}
+  let (tree, es) := elabOperation (.empty s.globalContext) stx ctx es
+  if es.errors.size > 0 then
+    let errMsg ← es.errors.foldlM (init := "Elaboration errors:\n") fun msg e =>
+      return s!"{msg}  {e.pos.line}:{e.pos.column}: {← e.data.toString}\n"
+    throw (IO.userError errMsg)
+  let op := tree.info.asOp!.op
+  return op
+
 end Strata.Elab
 end

Strata/DDM/Elab/Core.lean

@@ -1235,7 +1235,7 @@ partial def catElaborator (c : SyntaxCat) : TypingContext → Syntax → ElabM T
   | q`Init.CommaSepBy =>
     elabSeqWith c .comma "commaSepBy" (·.getSepArgs)
   | q`Init.SpaceSepBy =>
-    elabSeqWith c .space "spaceSepBy" (·.getSepArgs)
+    elabSeqWith c .space "spaceSepBy" (·.getArgs)
   | q`Init.SpacePrefixSepBy =>
     elabSeqWith c .spacePrefix "spacePrefixSepBy" (·.getArgs)
   | q`Init.NewlineSepBy =>

Strata/DDM/Format.lean

@@ -19,6 +19,8 @@ namespace Strata
 /--
 Check if a character is valid for starting a regular identifier.
 Regular identifiers must start with a letter or underscore.
+
+NOTE: When updating this function, you will want to consider updating Strata/DDM/Parser.lean as well.
 -/
 private def isIdBegin (c : Char) : Bool :=
   c.isAlpha || c == '_' || c == '$'

Strata/DDM/Parser.lean

@@ -117,11 +117,12 @@ def stringInputContext (fileName : System.FilePath) (contents : String) : InputC
   fileName := fileName.toString
   fileMap  := FileMap.ofString contents
 
+-- When updating this, you will want to consider updating Strata.isIdBegin in Strata/DDM/Format.lean as well.
 private def strataIsIdFirst (c : Char) : Bool :=
-  c.isAlpha || c == '_'
+  c.isAlpha || c == '_' || c == '$'
 
 private def strataIsIdRest (c : Char) : Bool :=
-  c.isAlphanum || c == '_' || c == '\'' || c == '.' || c == '?' || c == '!'
+  c.isAlphanum || c == '_' || c == '\'' || c == '.' || c == '?' || c == '!' || c == '$'
 
 private def isIdFirstOrBeginEscape (c : Char) : Bool :=
   strataIsIdFirst c || isIdBeginEscape c

Strata/DL/Imperative/SMTUtils.lean

@@ -5,6 +5,10 @@
 -/
 
 import Strata.DL.SMT.SMT
+import Strata.DL.SMT.DDMTransform.Parse
+import Strata.DL.SMT.DDMTransform.Translate
+import Strata.DDM.Elab
+import Strata.DDM.Format
 import Strata.DL.Imperative.PureExpr
 import Strata.DL.Imperative.EvalContext
 
@@ -16,16 +20,22 @@ namespace SMT
 
 /--
 A counterexample derived from an SMT solver is a map from an identifier
-to a string.
+to an `SMT.Term`.
 -/
-abbrev CounterEx (Ident : Type) := Map Ident String
+abbrev CounterEx (Ident : Type) := Map Ident Strata.SMT.Term
+
+/-- Render an `SMT.Term` to a string via the SMTDDM translation. -/
+private def termToString (t : Strata.SMT.Term) : String :=
+  match Strata.SMTDDM.termToString t with
+  | .ok s => s
+  | .error _ => repr t |>.pretty
 
 def CounterEx.format {Ident} [ToFormat Ident] (cex : CounterEx Ident) : Format :=
   match cex with
   | [] => ""
-  | [(id, v)] => f!"({id}, {v})"
+  | [(id, v)] => f!"({id}, {termToString v})"
   | (id, v) :: rest =>
-    (f!"({id}, {v}) ") ++ CounterEx.format rest
+    (f!"({id}, {termToString v}) ") ++ CounterEx.format rest
 
 instance {Ident} [ToFormat Ident] : ToFormat (CounterEx Ident) where
   format := CounterEx.format
@@ -83,27 +93,6 @@ def getSMTId {Ident Ty} [ToFormat Ident]
     let key : Strata.SMT.UF := { id := var', args := [], out := ty' }
     .ok (E.ufs[key]!)
 
-def getModel (m : String) : Except Format (List Strata.SMT.CExParser.KeyValue) := do
-  let cex ← Strata.SMT.CExParser.parseCEx m
-  return cex.pairs
-
-def processModel {P : PureExpr} [ToFormat P.Ident]
-    (typedVarToSMTFn : P.Ident → P.Ty → Except Format (String × Strata.SMT.TermType))
-    (vars : List P.TypedIdent) (cexs : List Strata.SMT.CExParser.KeyValue)
-    (E : Strata.SMT.EncoderState) : Except Format (CounterEx P.Ident) := do
-  match vars with
-  | [] => return []
-  | (var, ty) :: vrest =>
-    let id ← @getSMTId P.Ident P.Ty _ typedVarToSMTFn var ty E
-    let value ← findCExValue id cexs
-    let pair := (var, value)
-    let rest ← processModel typedVarToSMTFn vrest cexs E
-    .ok (pair :: rest)
-  where findCExValue id cexs : Except Format String :=
-    match cexs.find? (fun p => p.key == id) with
-    | none => .error f!"Cannot find model for id: {id}"
-    | some p => .ok p.value
-
 def runSolver (solver : String) (args : Array String) : IO IO.Process.Output := do
   let output ← IO.Process.output {
     cmd := solver
@@ -114,9 +103,90 @@ def runSolver (solver : String) (args : Array String) : IO IO.Process.Output :=
   --                         stdout: {repr output.stdout}"
   return output
 
+---------------------------------------------------------------------
+-- SMTDDM-based parsing
+---------------------------------------------------------------------
+
+/--
+Parse a verdict line ("sat", "unsat", "unknown") via the SMTResponse DDM
+dialect. Returns `some .sat`, `some .unsat`, `some .unknown`, or `none`
+on parse/conversion failure.
+-/
+private def parseVerdict (line : String) : IO (Option (Result PUnit)) := do
+  let inputCtx := Strata.Parser.stringInputContext "solver" (line ++ "\n")
+  let prg ←
+    try Strata.Elab.parseStrataProgramFromDialect
+          Strata.SMTResponseDDM.smtResponseDialects "SMTResponse" inputCtx
+    catch _ => return none
+  if prg.commands.isEmpty then return none
+  let op := prg.commands[0]!
+  match Strata.SMTResponseDDM.Command.ofAst op with
+  | .ok (.specific_success_response _ (.ssr_check_sat _ (.csr_sat _)))     => return some (.sat [])
+  | .ok (.specific_success_response _ (.ssr_check_sat _ (.csr_unsat _)))   => return some .unsat
+  | .ok (.specific_success_response _ (.ssr_check_sat _ (.csr_unknown _))) => return some .unknown
+  | _ => return none
+
+/--
+Parse a `(get-value ...)` model response using the SMTResponse DDM dialect.
+Uses `parseCategoryFromDialect` targeting `SMTResponse.GetValueResponse`
+directly, which avoids the ambiguity that arises when parsing at the
+`Command` level.
+
+Returns a list of (key-string, value-Term) pairs on success.
+-/
+private def parseModelDDM (modelStr : String) : IO (List (String × Strata.SMT.Term)) := do
+  let inputCtx := Strata.Parser.stringInputContext "solver-model" modelStr
+  let op ←
+    try Strata.Elab.parseCategoryFromDialect
+          Strata.SMTResponseDDM.smtResponseDialects q`SMTResponse.GetValueResponse inputCtx
+    catch _ => return []
+  match Strata.SMTResponseDDM.GetValueResponse.ofAst op with
+  | .ok (.get_value_response _ vps) =>
+    let pairs ← vps.val.toList.filterMapM fun vp =>
+      match vp with
+      | .valuation_pair _ t1 t2 => do
+        match Strata.SMTResponseDDM.translateFromDDMTermToUntyped t2 with
+        | .ok t2' =>
+          return .some (Strata.SMTResponseDDM.formatArg (.op (Strata.SMTResponseDDM.Term.toAst t1)),
+                  t2')
+        | .error _ =>
+          -- The model has an SMT expression (e.g., (lambda ...)) which cannot
+          -- be represented in Strata.SMT.Term. Filter out this variable from
+          -- the model.
+          return .none
+    return pairs
+  | .error _ => return []
+
+/--
+Process a parsed model (list of key-string / value-Term pairs) against the
+expected variables, matching each variable's SMT-encoded name to its
+value in the model.
+-/
+private def processModel {P : PureExpr} [ToFormat P.Ident]
+    (typedVarToSMTFn : P.Ident → P.Ty → Except Format (String × Strata.SMT.TermType))
+    (vars : List P.TypedIdent) (pairs : List (String × Strata.SMT.Term))
+    (E : Strata.SMT.EncoderState) : Except Format (CounterEx P.Ident) := do
+  match vars with
+  | [] => return []
+  | (var, ty) :: vrest =>
+    let id ← @getSMTId P.Ident P.Ty _ typedVarToSMTFn var ty E
+    let value ← findValue id pairs
+    let rest ← processModel typedVarToSMTFn vrest pairs E
+    .ok ((var, value) :: rest)
+  where findValue id pairs : Except Format Strata.SMT.Term :=
+    match pairs.find? (fun p => p.fst == id) with
+    | none => .error f!"Cannot find model for id: {id}"
+    | some p => .ok p.snd
+
 /--
 Interprets the output of SMT solver.
 
+Both the verdict line (sat/unsat/unknown) and the model (when sat) are
+parsed using the SMTResponse DDM dialect. The verdict is parsed as a full
+`Command`, while the model is parsed by targeting the
+`SMTResponse.GetValueResponse` category directly via
+`parseCategoryFromDialect`.
+
 When `reachCheck` is `true`, the solver output contains two verdict lines:
 the first is the reachability check result (are the path-condition assumptions
 satisfiable?), and the second is the proof check result. The reachability
@@ -128,46 +198,45 @@ def solverResult {P : PureExpr} [ToFormat P.Ident]
     (vars : List P.TypedIdent) (output : IO.Process.Output)
     (E : Strata.SMT.EncoderState) (smtsolver : String)
     (reachCheck : Bool := false)
-    : Except Format (Option (Result P.Ident) × Result P.Ident) := do
+    : IO (Except Format (Option (Result P.Ident) × Result P.Ident)) := do
   let stdout := output.stdout
-  -- When reachCheck is true, the first line of stdout is the reachability
-  -- verdict; strip it and parse it separately.
-  let (reachResult, proofStdout) := if reachCheck then
-    let pos := stdout.find (· == '\n')
-    let reachVerdictStr := (stdout.extract stdout.startPos pos).trimAscii
-    let reachResult : Result P.Ident := match reachVerdictStr with
-      | "sat" => .sat []
-      | "unsat" => .unsat
-      | _ => .unknown
-    let remaining := (stdout.extract pos stdout.endPos).drop 1
-    (some reachResult, remaining)
+  -- Split the next line from the remaining stdout, returning (line, rest).
+  let splitLine (s : String) : String × String :=
+    let pos := s.find (· == '\n')
+    let line := (s.extract s.startPos pos).trimAscii.toString
+    let rest := s.extract pos s.endPos
+    (line, rest)
+  -- When reachCheck is true, the first line is the reachability verdict.
+  let (reachResult, proofStdout) ← if reachCheck then do
+    let (reachLine, remaining) := splitLine stdout
+    let reachResult : Result P.Ident ← do
+      match ← parseVerdict reachLine with
+      | some (.sat _) => pure (.sat [])
+      | some .unsat   => pure .unsat
+      | _             => pure .unknown
+    pure (some reachResult, remaining.drop 1 |>.toString)
   else
-    (none, stdout)
-  -- Parse the proof verdict from the (possibly trimmed) stdout
-  let pos := proofStdout.find (· == '\n')
-  let verdict := proofStdout.extract proofStdout.startPos pos |>.trimAscii
-  let rest := proofStdout.extract pos proofStdout.endPos
-  match verdict with
-  | "sat"     =>
-    let rawModel ← getModel rest
-    -- We suppress any model processing errors.
-    -- Likely, these would be because of the suboptimal implementation
-    -- of the model parser, which shouldn't hold back useful
-    -- feedback (i.e., problem was `sat`) from the user.
-    match (processModel typedVarToSMTFn vars rawModel E) with
-    | .ok model => .ok (reachResult, .sat model)
-    | .error _model_err => .ok (reachResult, .sat [])
-  | "unsat"   => .ok (reachResult, .unsat)
-  | "unknown" => .ok (reachResult, .unknown)
-  | _     =>
+    pure (none, stdout)
+  -- Parse the proof verdict from the (possibly trimmed) stdout.
+  let (verdictStr, rest) := splitLine proofStdout
+  match ← parseVerdict verdictStr with
+  | some (.sat _) =>
+    -- Parse model via SMTDDM targeting GetValueResponse category directly.
+    let pairs ← parseModelDDM rest
+    match processModel typedVarToSMTFn vars pairs E with
+    | .ok model => return .ok (reachResult, .sat model)
+    | .error _  => return .ok (reachResult, .sat [])
+  | some .unsat   => return .ok (reachResult, .unsat)
+  | some .unknown => return .ok (reachResult, .unknown)
+  | _ =>
     let stderr := output.stderr
     let hasExecError := stderr.contains "could not execute external process"
     let hasFileError := stderr.contains "No such file or directory"
     let suggestion :=
       if (hasExecError || hasFileError) && smtsolver == Core.defaultSolver then
         s!" \nEnsure {Core.defaultSolver} is on your PATH or use --solver to specify another SMT solver."
       else ""
-    .error s!"stderr:{stderr}{suggestion}\nsolver stdout: {output.stdout}\n"
+    return .error s!"stderr:{stderr}{suggestion}\nsolver stdout: {output.stdout}\n"
 
 def addLocationInfo {P : PureExpr} [BEq P.Ident]
   (md : Imperative.MetaData P) (message : String × String)
@@ -211,7 +280,7 @@ def dischargeObligation {P : PureExpr} [ToFormat P.Ident] [BEq P.Ident]
   if printFilename then IO.println s!"Wrote problem to {filename}."
 
   let solver_output ← runSolver smtsolver (#[filename] ++ solver_options)
-  match solverResult typedVarToSMTFn vars solver_output estate smtsolver (reachCheck := reachCheck) with
+  match ← solverResult typedVarToSMTFn vars solver_output estate smtsolver (reachCheck := reachCheck) with
   | .error e => return .error e
   | .ok (reachDecision, result) => return .ok (reachDecision, result, estate)