litex-lang 0.9.75-beta

use crate::{common::keywords::GREATER_EQUAL, prelude::*, stmt::parameter_def::ParamDefWithType};
use std::fmt;

#[derive(Clone)]
pub struct HaveFnByInducNestedCase {
    pub case_fact: AndChainAtomicFact,
    pub equal_to: Obj,
}

#[derive(Clone)]
pub enum HaveFnByInducLastCase {
    EqualTo(Obj),
    NestedCases(Vec<HaveFnByInducNestedCase>),
}

// have fn by induc from 0: f(x Z: x >= 0) R:
//     case 0: 1
//     case 1: 1
//     case >= 2:
//         case x % 2 = 0: f(x / 2)
//         case x % 2 = 1: f(x / 2) + f(x / 2 + 1)
#[derive(Clone)]
pub struct HaveFnByInducStmt {
    pub induc_from: Obj,
    pub name: String,
    pub param: String,
    pub ret_set: Obj,
    pub special_cases_equal_tos: Vec<Obj>,
    pub last_case: HaveFnByInducLastCase,
    pub line_file: LineFile,
}

#[derive(Clone)]
pub struct DefAbstractPropStmt {
    pub name: String,
    pub params: Vec<String>,
    pub line_file: LineFile,
}

impl DefAbstractPropStmt {
    pub fn new(name: String, params: Vec<String>, line_file: LineFile) -> Self {
        DefAbstractPropStmt {
            name,
            params,
            line_file,
        }
    }
}

#[derive(Clone)]
pub struct DefFamilyStmt {
    pub name: String,
    pub params_def_with_type: ParamDefWithType,
    pub dom_facts: Vec<OrAndChainAtomicFact>,
    pub equal_to: Obj,
    pub line_file: LineFile,
}

/// `have fn` `{ … }` piece: parameter names match binders in dom/ret; build stored `Obj::FnSet` with [`Runtime::fn_set_from_fn_set_clause`].
#[derive(Clone)]
pub struct FnSetClause {
    pub params_def_with_set: Vec<ParamGroupWithSet>,
    pub dom_facts: Vec<OrAndChainAtomicFact>,
    pub ret_set: Obj,
}

impl FnSetClause {
    pub fn new(
        params_def_with_set: Vec<ParamGroupWithSet>,
        dom_facts: Vec<OrAndChainAtomicFact>,
        ret_set: Obj,
    ) -> Self {
        FnSetClause {
            params_def_with_set,
            dom_facts,
            ret_set,
        }
    }

    /// Outer `{...}` binders first, then each nested function return `fn` layer, in order.
    /// Used when parsing `have fn ... = rhs` so RHS identifiers match nested return fn-set params.
    pub fn collect_all_param_names_including_nested_ret_fn_sets(&self) -> Vec<String> {
        let mut names = ParamGroupWithSet::collect_param_names(&self.params_def_with_set);
        let mut ret_set = self.ret_set.clone();
        while let Obj::FnSet(inner) = ret_set {
            names.extend(ParamGroupWithSet::collect_param_names(
                &inner.body.params_def_with_set,
            ));
            ret_set = (*inner.body.ret_set).clone();
        }
        names
    }
}

#[derive(Clone)]
pub struct HaveFnEqualCaseByCaseStmt {
    pub name: String,
    pub fn_set_clause: FnSetClause,
    pub cases: Vec<AndChainAtomicFact>,
    pub equal_tos: Vec<Obj>,
    pub line_file: LineFile,
}

#[derive(Clone)]
pub struct HaveFnEqualStmt {
    pub name: String,
    pub equal_to_anonymous_fn: AnonymousFn,
    pub line_file: LineFile,
}

#[derive(Clone)]
pub struct HaveFnByForallExistUniqueStmt {
    pub fn_name: String,
    pub forall: ForallFact,
    pub line_file: LineFile,
}

// have by exist a R st {$p(a)}: a
#[derive(Clone)]
pub struct HaveByExistStmt {
    pub equal_tos: Vec<String>,
    pub exist_fact_in_have_obj_st: ExistFactEnum,
    pub line_file: LineFile,
}

#[derive(Clone)]
pub struct HaveObjEqualStmt {
    pub param_def: ParamDefWithType,
    pub objs_equal_to: Vec<Obj>,
    pub line_file: LineFile,
}

#[derive(Clone)]
pub struct HaveObjInNonemptySetOrParamTypeStmt {
    pub param_def: ParamDefWithType,
    pub line_file: LineFile,
}

#[derive(Clone)]
pub struct DefLetStmt {
    pub param_def: ParamDefWithType,
    pub facts: Vec<Fact>,
    pub line_file: LineFile,
}

#[derive(Clone)]
pub struct DefPropStmt {
    pub name: String,
    pub params_def_with_type: ParamDefWithType,
    pub iff_facts: Vec<Fact>,
    pub line_file: LineFile,
}

impl fmt::Display for DefAbstractPropStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{} {}{}{}{}",
            ABSTRACT_PROP,
            self.name,
            LEFT_BRACE,
            vec_to_string_join_by_comma(&self.params),
            RIGHT_BRACE
        )
    }
}

impl DefPropStmt {
    pub fn new(
        name: String,
        params_def_with_type: ParamDefWithType,
        iff_facts: Vec<Fact>,
        line_file: LineFile,
    ) -> Self {
        DefPropStmt {
            name,
            params_def_with_type,
            iff_facts,
            line_file,
        }
    }
}

impl fmt::Display for DefPropStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.iff_facts.len() {
            0 => write!(
                f,
                "{} {}{}",
                PROP,
                self.name,
                braced_string(&self.params_def_with_type)
            ),
            _ => write!(
                f,
                "{} {}{}{}\n{}",
                PROP,
                self.name,
                braced_string(&self.params_def_with_type),
                COLON,
                vec_to_string_add_four_spaces_at_beginning_of_each_line(&self.iff_facts, 1)
            ),
        }
    }
}

impl DefLetStmt {
    pub fn new(param_def: ParamDefWithType, facts: Vec<Fact>, line_file: LineFile) -> Self {
        DefLetStmt {
            param_def,
            facts,
            line_file,
        }
    }
}

impl fmt::Display for DefLetStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let param_str = self.param_def.to_string();
        match self.facts.len() {
            0 => write!(f, "{} {}", LET, param_str),
            _ => write!(
                f,
                "{} {}{}\n{}",
                LET,
                param_str,
                COLON,
                vec_to_string_add_four_spaces_at_beginning_of_each_line(&self.facts, 1)
            ),
        }
    }
}

impl HaveObjInNonemptySetOrParamTypeStmt {
    pub fn new(param_def: ParamDefWithType, line_file: LineFile) -> Self {
        HaveObjInNonemptySetOrParamTypeStmt {
            param_def,
            line_file,
        }
    }
}

impl fmt::Display for HaveObjInNonemptySetOrParamTypeStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} {}", HAVE, self.param_def.to_string())
    }
}

impl HaveObjEqualStmt {
    pub fn new(param_def: ParamDefWithType, objs_equal_to: Vec<Obj>, line_file: LineFile) -> Self {
        HaveObjEqualStmt {
            param_def,
            objs_equal_to,
            line_file,
        }
    }
}

impl fmt::Display for HaveObjEqualStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{} {} {} {}",
            HAVE,
            self.param_def.to_string(),
            EQUAL,
            vec_to_string_join_by_comma(&self.objs_equal_to)
        )
    }
}

impl HaveByExistStmt {
    pub fn new(
        equal_tos: Vec<String>,
        exist_fact_in_have_obj_st: ExistFactEnum,
        line_file: LineFile,
    ) -> Self {
        HaveByExistStmt {
            equal_tos,
            exist_fact_in_have_obj_st,
            line_file,
        }
    }
}

impl fmt::Display for HaveByExistStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{} {} {} {} {}",
            HAVE,
            BY,
            self.exist_fact_in_have_obj_st,
            COLON,
            vec_to_string_join_by_comma(&self.equal_tos),
        )
    }
}

impl HaveFnEqualStmt {
    pub fn new(name: String, equal_to_anonymous_fn: AnonymousFn, line_file: LineFile) -> Self {
        HaveFnEqualStmt {
            name,
            equal_to_anonymous_fn,
            line_file,
        }
    }
}

impl fmt::Display for HaveFnEqualStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let fn_set_clause = FnSetClause::new(
            self.equal_to_anonymous_fn.body.params_def_with_set.clone(),
            self.equal_to_anonymous_fn.body.dom_facts.clone(),
            (*self.equal_to_anonymous_fn.body.ret_set).clone(),
        );
        write!(
            f,
            "{} {} {}{} {} {}",
            HAVE,
            FN_LOWER_CASE,
            self.name,
            brace_vec_colon_vec_to_string(
                &fn_set_clause.params_def_with_set,
                &fn_set_clause.dom_facts
            ),
            EQUAL,
            self.equal_to_anonymous_fn.equal_to
        )
    }
}

impl HaveFnByForallExistUniqueStmt {
    pub fn new(fn_name: String, forall: ForallFact, line_file: LineFile) -> Self {
        HaveFnByForallExistUniqueStmt {
            fn_name,
            forall,
            line_file,
        }
    }
}

impl fmt::Display for HaveFnByForallExistUniqueStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} {} {} {}", HAVE, FN_LOWER_CASE, self.fn_name, BY)?;
        write!(f, " {}", self.forall)
    }
}

impl fmt::Display for HaveFnEqualCaseByCaseStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let cases_and_proofs = self
            .cases
            .iter()
            .enumerate()
            .map(|(i, case)| {
                to_string_and_add_four_spaces_at_beginning_of_each_line(
                    &format!(
                        "{} {}{} {}{} {} {}",
                        CASE,
                        case,
                        COMMA,
                        self.name,
                        braced_vec_to_string(&ParamGroupWithSet::collect_param_names(
                            &self.fn_set_clause.params_def_with_set,
                        )),
                        EQUAL,
                        self.equal_tos[i]
                    ),
                    1,
                )
            })
            .collect::<Vec<String>>();

        write!(
            f,
            "{} {} {}{} {} {}\n{}",
            HAVE,
            FN_LOWER_CASE,
            self.name,
            brace_vec_colon_vec_to_string(
                &self.fn_set_clause.params_def_with_set,
                &self.fn_set_clause.dom_facts
            ),
            EQUAL,
            COLON,
            vec_to_string_with_sep(&cases_and_proofs, "\n".to_string())
        )
    }
}

impl HaveFnEqualCaseByCaseStmt {
    pub fn new(
        name: String,
        fn_set_clause: FnSetClause,
        cases: Vec<AndChainAtomicFact>,
        equal_tos: Vec<Obj>,
        line_file: LineFile,
    ) -> Self {
        HaveFnEqualCaseByCaseStmt {
            name,
            fn_set_clause,
            cases,
            equal_tos,
            line_file,
        }
    }
}

pub fn induc_obj_plus_offset(induc_from: &Obj, offset: usize) -> Obj {
    if offset == 0 {
        induc_from.clone()
    } else {
        Add::new(induc_from.clone(), Number::new(offset.to_string()).into()).into()
    }
}

fn flatten_and_chain_to_atomic_facts(c: &AndChainAtomicFact) -> Vec<AtomicFact> {
    match c {
        AndChainAtomicFact::AtomicFact(a) => vec![a.clone()],
        AndChainAtomicFact::AndFact(af) => af.facts.clone(),
        AndChainAtomicFact::ChainFact(cf) => cf.facts().unwrap(),
    }
}

fn merge_two_and_chain_clauses(
    a: AndChainAtomicFact,
    b: AndChainAtomicFact,
    line_file: LineFile,
) -> AndChainAtomicFact {
    let mut atoms = flatten_and_chain_to_atomic_facts(&a);
    atoms.extend(flatten_and_chain_to_atomic_facts(&b));
    AndChainAtomicFact::AndFact(AndFact::new(atoms, line_file))
}

impl HaveFnByInducStmt {
    /// Source-shaped `fn` block (param names + dom + ret); build stored [`FnSet`] via [`Runtime::fn_set_from_fn_set_clause`].
    pub fn fn_user_fn_set_clause(&self) -> FnSetClause {
        FnSetClause::new(
            vec![ParamGroupWithSet::new(
                vec![self.param.clone()],
                StandardSet::Z.into(),
            )],
            vec![OrAndChainAtomicFact::AtomicFact(
                GreaterEqualFact::new(
                    obj_for_bound_param_in_scope(self.param.clone(), ParamObjType::FnSet),
                    self.induc_from.clone(),
                    self.line_file.clone(),
                )
                .into(),
            )],
            self.ret_set.clone(),
        )
    }

    /// `forall x Z: ...` 里与 `fn` 定义域一致的那一段：标识符用源码 [`Self::param`]，与 [`Self::fn_user_fn_set_clause`] 的 dom 语义相同。
    pub fn forall_fn_base_dom_exist_or_facts(&self) -> Vec<Fact> {
        vec![GreaterEqualFact::new(
            obj_for_bound_param_in_scope(self.param.clone(), ParamObjType::Forall),
            self.induc_from.clone(),
            self.line_file.clone(),
        )
        .into()]
    }

    pub fn new(
        name: String,
        param: String,
        ret_set: Obj,
        induc_from: Obj,
        special_cases_equal_tos: Vec<Obj>,
        last_case: HaveFnByInducLastCase,
        line_file: LineFile,
    ) -> Self {
        HaveFnByInducStmt {
            name,
            param,
            ret_set,
            induc_from,
            special_cases_equal_tos,
            last_case,
            line_file,
        }
    }

    /// 展开为与旧 `HaveFnEqualCaseByCaseStmt` 兼容的平铺 `case` 列表（源码最后一条为 `case >= n:`（n 为特例个数），此处仍展开为 `param = from + n` 与可选子条件的合取）。
    pub fn to_have_fn_equal_case_by_case_stmt(&self) -> HaveFnEqualCaseByCaseStmt {
        let line_file = self.line_file.clone();
        let left_id: Obj = obj_for_bound_param_in_scope(self.param.clone(), ParamObjType::Induc);
        let n = self.special_cases_equal_tos.len();
        let mut cases: Vec<AndChainAtomicFact> = Vec::new();
        let mut equal_tos: Vec<Obj> = Vec::new();
        for i in 0..n {
            let when = AndChainAtomicFact::AtomicFact(
                EqualFact::new(
                    left_id.clone(),
                    induc_obj_plus_offset(&self.induc_from, i),
                    line_file.clone(),
                )
                .into(),
            );
            cases.push(when);
            equal_tos.push(self.special_cases_equal_tos[i].clone());
        }
        let step = AndChainAtomicFact::AtomicFact(
            EqualFact::new(
                left_id.clone(),
                induc_obj_plus_offset(&self.induc_from, n),
                line_file.clone(),
            )
            .into(),
        );
        match &self.last_case {
            HaveFnByInducLastCase::EqualTo(eq) => {
                cases.push(step);
                equal_tos.push(eq.clone());
            }
            HaveFnByInducLastCase::NestedCases(last_pairs) => {
                for nested in last_pairs {
                    let merged = merge_two_and_chain_clauses(
                        step.clone(),
                        nested.case_fact.clone(),
                        line_file.clone(),
                    );
                    cases.push(merged);
                    equal_tos.push(nested.equal_to.clone());
                }
            }
        }
        HaveFnEqualCaseByCaseStmt::new(
            self.name.clone(),
            self.fn_user_fn_set_clause(),
            cases,
            equal_tos,
            line_file,
        )
    }
}

impl fmt::Display for HaveFnByInducStmt {
    /// Display uses the same parameter names as in source.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let n = self.special_cases_equal_tos.len();
        write!(
            f,
            "{} {} {} {} {} {} {}{}",
            HAVE, FN_LOWER_CASE, BY, INDUC, FROM, " ", self.induc_from, COLON,
        )?;
        write!(f, " {} {}{}", self.name, LEFT_BRACE, self.param)?;
        write!(
            f,
            " {} {} {} {} {} {} {} {} {} {} {} {}",
            Z,
            COLON,
            " ",
            self.param,
            " ",
            GREATER_EQUAL,
            " ",
            self.induc_from,
            RIGHT_BRACE,
            " ",
            self.ret_set,
            COLON,
        )?;
        for (i, eq) in self.special_cases_equal_tos.iter().enumerate() {
            writeln!(f)?;
            write!(f, "    {} {}: {}", CASE, i, eq)?;
        }
        writeln!(f)?;
        match &self.last_case {
            HaveFnByInducLastCase::EqualTo(eq) => {
                write!(f, "    {} {} {}: {}", CASE, GREATER_EQUAL, n, eq)?;
            }
            HaveFnByInducLastCase::NestedCases(nested) => {
                write!(f, "    {} {} {}:", CASE, GREATER_EQUAL, n)?;
                for nc in nested {
                    writeln!(f)?;
                    write!(f, "        {} {}: {}", CASE, nc.case_fact, nc.equal_to)?;
                }
            }
        }
        Ok(())
    }
}

impl DefFamilyStmt {
    pub fn new(
        name: String,
        params_def_with_type: ParamDefWithType,
        dom_facts: Vec<OrAndChainAtomicFact>,
        equal_to: Obj,
        line_file: LineFile,
    ) -> Self {
        DefFamilyStmt {
            name,
            params_def_with_type,
            dom_facts,
            equal_to,
            line_file,
        }
    }
}

impl fmt::Display for DefFamilyStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{} {}{}{} {} {}{} {} {}",
            FAMILY,
            self.name,
            LEFT_BRACE,
            self.params_def_with_type.to_string(),
            COLON,
            vec_to_string_join_by_comma(&self.dom_facts),
            RIGHT_BRACE,
            EQUAL,
            self.equal_to
        )
    }
}