Repository: sympy/sympy
File: sympy/logic/algorithms/z3_wrapper.py (Python)

Code before:
    for var in z3_model:
        var_name = var.name()
        # Only process variables that start with 'd' followed by digits
        if var_name.startswith('d') and var_name[1:].isdigit():
            var_id = int(var_name[1:])
            # Only include if this variable ID exists in our reverse encoding
            if var_id in rev_enc:
                result[rev_enc[var_id]] = bool(z3_model[var])
    return result


def clause_to_assertion(clause):
    """Converts a single CNF clause to SMT-LIB assertion format."""
    clause_strings = [f"d{abs(lit)}" if lit > 0 else f"(not d{abs(lit)})" for lit in clause]
    return "(assert (or " + " ".join(clause_strings) + "))"


def find_uninterpreted_functions(expr):
    uf_arity = {}
    # Handle bare function heads (symbols that are functions but not called)
    for head in expr.atoms(UndefinedFunction):
        uf_arity.setdefault(str(head), None)
    # Handle function calls like f(a), g(x,y)
    if AppliedUndef is not None:
        calls = expr.atoms(AppliedUndef)
    else:
        calls = [c for c in expr.atoms(Function) if isinstance(c.func, UndefinedFunction)]

    for call in calls:
        name = str(call.func)
        arity = len(call.args)
        prev = uf_arity.get(name)
        uf_arity[name] = arity if prev is None else max(prev or 0, arity)

    return uf_arity


def collect_all_uninterpreted_functions(enc_cnf):
    """Returns a dictionary, mapping uninterpreted function names to their maximum arity."""
    uf_table = {}
    for pred, enc in enc_cnf.encoding.items():
        if not isinstance(pred, AppliedPredicate) or pred.function not in supported_predicates:
            continue
        # Update the table of uninterpreted functions.
        pred_ufs = find_uninterpreted_functions(pred)
        # If a function appears multiple times, store the largest number of arguments it takes.
        for name, arity in pred_ufs.items():
            if name in uf_table:
                if uf_table[name] is None or arity is None:
                    uf_table[name] = max(uf_table[name] or 0, arity or 0)
                else:

Review comment (suggestion):
They shouldn't have multiple possible arities. If that happens you should probably raise an exception.

Code after:

        symbols |= pred.free_symbols
        pred = pred_str
        assertion = f"(assert  (implies d{enc} {pred}))"
        assertions.append(assertion)

    for sym in symbols:
        declarations.append(f"(declare-const {sym} Real)")

    for f, arity in uninterpreted_functions_to_arity.items():
        # Rather than defining new types for the domain and range
        # we assume they are real to avoid unnecesary complexity.
        arg_types = " ".join(["Real"] * arity)
        declarations.append(f"(declare-fun {f} ({arg_types}) Real)")

    declarations = "\n".join(declarations)
    assertions = "\n".join(assertions)
    s.from_string(declarations)
    s.from_string(assertions)

    return s


known_functions = {
            Add: '+',
            Mul: '*',

            Equality: '=',
            LessThan: '<=',
            GreaterThan: '>=',
            StrictLessThan: '<',
            StrictGreaterThan: '>',

            EqualityPredicate(): '=',
            LessThanPredicate(): '<=',
            GreaterThanPredicate(): '>=',
            StrictLessThanPredicate(): '<',
            StrictGreaterThanPredicate(): '>',

            Abs: 'abs',
            Min: 'min',
            Max: 'max',
            Pow: '^',

            And: 'and',
            Or: 'or',
            Xor: 'xor',
            Not: 'not',
            ITE: 'ite',
            Implies: '=>',
        }