z3/.github/skills/encode/scripts/encode.py

#!/usr/bin/env python3
"""
encode.py: validate and format SMT-LIB2 encodings.

Usage:
    python encode.py --validate formula.smt2
    python encode.py --validate formula.smt2 --debug
"""

import argparse
import re
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent / "shared"))
from z3db import Z3DB, run_z3, setup_logging


VALIDATION_TIMEOUT = 5


def read_input(path_or_stdin: str) -> str:
    """Read formula from a file path or stdin (when path is '-')."""
    if path_or_stdin == "-":
        return sys.stdin.read()
    p = Path(path_or_stdin)
    if not p.exists():
        print(f"file not found: {p}", file=sys.stderr)
        sys.exit(1)
    return p.read_text()


def find_errors(output: str) -> list:
    """Extract (error ...) messages from Z3 output."""
    return re.findall(r'\(error\s+"([^"]+)"\)', output)


def validate(formula: str, z3_bin: str = None, debug: bool = False) -> dict:
    """
    Validate an SMT-LIB2 formula by piping it through z3 -in.
    Returns a dict with 'valid' (bool), 'errors' (list), and 'raw' output.
    """
    result = run_z3(
        formula, z3_bin=z3_bin, timeout=VALIDATION_TIMEOUT, debug=debug,
    )
    errors = find_errors(result["stdout"]) + find_errors(result["stderr"])

    if result["result"] == "timeout":
        # Timeout during validation is not a syntax error: the formula
        # parsed successfully but solving exceeded the limit.  That counts
        # as syntactically valid.
        return {"valid": True, "errors": [], "raw": result}

    if errors or result["exit_code"] != 0:
        return {"valid": False, "errors": errors, "raw": result}

    return {"valid": True, "errors": [], "raw": result}


def report_errors(errors: list, formula: str):
    """Print each syntax error with surrounding context."""
    lines = formula.splitlines()
    print(f"validation failed: {len(errors)} error(s)", file=sys.stderr)
    for err in errors:
        print(f"  : {err}", file=sys.stderr)
    if len(lines) <= 20:
        print("formula:", file=sys.stderr)
        for i, line in enumerate(lines, 1):
            print(f"  {i:3d}  {line}", file=sys.stderr)


def write_output(formula: str, output_path: str, fmt: str):
    """Write the validated formula to a file or stdout."""
    if fmt == "python":
        print("python format output is generated by the agent, "
              "not by this script", file=sys.stderr)
        sys.exit(1)

    if output_path:
        Path(output_path).write_text(formula)
        print(f"written to {output_path}")
    else:
        print(formula)


def main():
    parser = argparse.ArgumentParser(prog="encode")
    parser.add_argument(
        "--validate",
        metavar="FILE",
        help="path to .smt2 file to validate, or '-' for stdin",
    )
    parser.add_argument(
        "--format",
        choices=["smtlib2", "python"],
        default="smtlib2",
        help="output format (default: smtlib2)",
    )
    parser.add_argument(
        "--output",
        metavar="FILE",
        default=None,
        help="write result to file instead of stdout",
    )
    parser.add_argument("--z3", default=None, help="path to z3 binary")
    parser.add_argument("--db", default=None)
    parser.add_argument("--debug", action="store_true")
    args = parser.parse_args()

    setup_logging(args.debug)

    if not args.validate:
        parser.error("provide --validate FILE")
        return

    formula = read_input(args.validate)

    db = Z3DB(args.db)
    run_id = db.start_run("encode", formula)

    result = validate(formula, z3_bin=args.z3, debug=args.debug)

    if result["valid"]:
        db.log_formula(run_id, formula, "valid")
        db.finish_run(run_id, "valid", result["raw"]["duration_ms"], 0)
        write_output(formula, args.output, args.format)
        db.close()
        sys.exit(0)
    else:
        db.log_formula(run_id, formula, "error")
        for err in result["errors"]:
            db.log_finding(run_id, "syntax", err, severity="error")
        db.finish_run(
            run_id, "error",
            result["raw"]["duration_ms"],
            result["raw"]["exit_code"],
        )
        report_errors(result["errors"], formula)
        db.close()
        sys.exit(1)


if __name__ == "__main__":
    main()