from typing import TypeGuard, cast

import z3


class Relation(z3.FuncDeclRef):
    def __call__(self, *args: z3.ExprRef) -> z3.BoolRef:
        return cast(z3.BoolRef, cast(z3.FuncDeclRef, super())(*args))


def is_relation(decl: z3.FuncDeclRef) -> TypeGuard[Relation]:
    return decl.range() == z3.BoolSort()


def relation(name: str, *sig: z3.SortRef) -> Relation:
    return cast(Relation, z3.Function(name, *sig, z3.BoolSort()))


def fresh_relation(*sig: z3.SortRef) -> Relation:
    return cast(Relation, z3.FreshFunction(*sig, z3.BoolSort()))
