/*
 * Copyright Cedar Contributors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use std::{collections::BTreeMap, sync::Arc};

use super::{op::Uuf, term::Term, term_type::TermType};

/// Unary defined function
#[derive(Clone, Debug, PartialEq, Eq, Ord, PartialOrd)]
pub struct Udf {
    pub arg: TermType,
    pub out: TermType,
    pub table: Arc<BTreeMap<Term, Term>>,
    pub default: Term,
}

impl Udf {
    pub fn is_literal(&self) -> bool {
        self.default.is_literal()
            && self
                .table
                .iter()
                .all(|(k, v)| v.is_literal() && k.is_literal())
    }
}

/// Represents unary functions over terms. These functions can be uninterpreted
/// symbols (`Uuf`) or interpretations of these symbols that are generated by an
/// SMT solver (`Udf`).
///
/// We assume that SMT function interpretations take a fixed format --
/// effectively, a lookup table with a fixed number of entries, along with a
/// default value that is returned for all elements of the domain that are not
/// found in the table. The SMT solver represents these tables as nested `ite`
/// expressions, where each condition checks if the input is equal to a value in
/// the domain, and if so returns the corresponding output. Our target SMT
/// solver (CVC5) always returns interpretations of this form.
#[derive(Clone, Debug, PartialEq, Eq, Ord, PartialOrd)]
pub enum UnaryFunction {
    Uuf(Arc<Uuf>),
    Udf(Arc<Udf>),
}

impl UnaryFunction {
    pub fn arg_type(&self) -> &TermType {
        match self {
            UnaryFunction::Uuf(f) => &f.arg,
            UnaryFunction::Udf(f) => &f.arg,
        }
    }

    pub fn out_type(&self) -> &TermType {
        match self {
            UnaryFunction::Uuf(f) => &f.out,
            UnaryFunction::Udf(f) => &f.out,
        }
    }

    pub fn is_literal(&self) -> bool {
        match self {
            UnaryFunction::Uuf(_) => false,
            UnaryFunction::Udf(f) => f.is_literal(),
        }
    }
}
