[Type check] An and/or whose right operand returns counts as divergent, so a non-Unit function can end without a value #3929

"""A non-Unit function can end without a value: `and`/`or` with a `return` on the right counts as divergent

Type-checks each Q# program below (it should be rejected), then runs it in the Q#
simulator and compiles it to QIR for each target profile, running that QIR with qir-runner. Exits 0 while the bug reproduces.

Requires: pip install qdk qirrunner
"""

import os
import re
import sys
import tempfile

import qirrunner
from qdk import Context, TargetProfile

# name -> (Q# source, correct outcome, runs that show the bug). Each program should be
# rejected by the type checker.
CASES = {
    'false and (return 1);': ("""
function F() : Int {
    false and (return 1);
}
operation Main() : Int { F() }
""", 'rejected', ['Simulator', 'Adaptive_RI', 'Adaptive_RIF', 'Adaptive']),
    'true or (return 1);': ("""
function F() : Int {
    true or (return 1);
}
operation Main() : Int { F() }
""", 'rejected', ['Simulator', 'Adaptive_RI', 'Adaptive_RIF', 'Adaptive']),
    'set b and= (return 1);': ("""
function F() : Int {
    mutable b = false;
    set b and= (return 1);
}
operation Main() : Int { F() }
""", 'rejected', ['Simulator', 'Adaptive_RI', 'Adaptive_RIF', 'Adaptive']),
}

ENTRY = "Main()"
SHOTS = 1
PROFILES = ['Base', 'Adaptive_RI', 'Adaptive_RIF', 'Adaptive']


def fmt(value):
    """Formats a simulator value the same way as qir-runner output below."""
    if value is None:
        return "()"
    if type(value).__name__ == "Result":
        return "1" if str(value).endswith("One") else "0"
    if isinstance(value, bool):
        return str(value).lower()
    if isinstance(value, float):
        return repr(value)
    if isinstance(value, (tuple, list)):
        inner = ", ".join(fmt(v) for v in value)
        return f"({inner})" if isinstance(value, tuple) else f"[{inner}]"
    return str(value)


def parse_shots(text):
    """Turns qir-runner's output records into one formatted value per shot."""
    shots, records = [], None
    for line in text.splitlines():
        fields = line.split("\t")
        if fields[0] == "START":
            records = []
        elif fields[0] == "OUTPUT" and records is not None:
            records.append((fields[1], fields[2]))
        elif fields[0] == "END" and records is not None:
            shots.append(to_text(iter(records)) if records else "()")
            if fields[1] != "0":
                shots[-1] = f"exit code {fields[1]}"
            records = None
    return shots


def to_text(records):
    kind, value = next(records)
    if kind in ("TUPLE", "ARRAY"):
        inner = ", ".join(to_text(records) for _ in range(int(value)))
        return f"({inner})" if kind == "TUPLE" else f"[{inner}]"
    if kind == "DOUBLE":
        return repr(float(value))
    if kind == "BOOL":
        return value.lower()
    return value


def run_qir(qir):
    with tempfile.NamedTemporaryFile("w", suffix=".ll", delete=False) as file:
        file.write(qir)
    handler = qirrunner.OutputHandler()
    try:
        qirrunner.run(file.name, shots=SHOTS, rng_seed=1, output_fn=handler.handle)
    except Exception as error:  # qir-runner rejects the module
        message = str(error).strip().splitlines()[0]
        return {"invalid QIR: " + re.sub(r"^\S+\.ll:\d+:\d+: (error: )?", "", message)}
    finally:
        os.unlink(file.name)
    return set(parse_shots(handler.get_output()))


def compile_qir(source, profile):
    """Returns (qir, None) on success or (None, description) when compilation fails."""
    ctx = Context(target_profile=getattr(TargetProfile, profile))
    try:
        ctx.eval(source)
        return str(ctx.compile(ENTRY)), None
    except (KeyboardInterrupt, SystemExit):
        raise
    except BaseException as error:  # pyo3 surfaces Rust panics as PanicException (a BaseException)
        kind = "COMPILER PANIC" if type(error).__name__ == "PanicException" else "rejected"
        return None, f"{kind}: {str(error).strip().splitlines()[0]}"


def outputs(source):
    rows = {}
    ctx = Context(target_profile=TargetProfile.Unrestricted)
    try:
        ctx.eval(source)
    except Exception as error:  # the type checker rejects the program, as it should
        rows["Simulator"] = {"rejected: " + str(error).strip().splitlines()[0]}
    else:
        rows["Simulator"] = {fmt(v) for v in ctx.run(ENTRY, SHOTS, seed=1)}
    for profile in PROFILES:
        qir, failure = compile_qir(source, profile)
        rows[profile] = {failure} if failure else run_qir(qir)
    return rows


def show(name, rows, expected):
    print(f"--- {name} (expected {expected})")
    width = max(map(len, rows)) + 2
    for row, values in rows.items():
        print(f"  {row:<{width}}{' | '.join(sorted(values))}")


def main():
    reproduced = []
    for name, (source, expected, affected) in CASES.items():
        rows = outputs(source)
        show(name, rows, expected)
        wrong = [run for run in affected if not all(v.startswith("rejected") for v in rows[run])]
        if wrong:
            reproduced.append(f"{name} ({', '.join(wrong)})")
    if reproduced:
        print(f"REPRODUCED: {'; '.join(reproduced)}")
        return 0
    print("NOT REPRODUCED")
    return 1


if __name__ == "__main__":
    sys.exit(main())