litex-lang 0.9.65-beta

use crate::obj::field_access_to_string;
use crate::prelude::*;
use std::collections::HashMap;
use std::rc::Rc;

pub struct Runtime {
    pub module_manager: ModuleManager,
    pub environment_stack: Vec<Box<Environment>>,
    pub parsing_time_name_scope_stack: Vec<HashMap<String, LineFile>>,
}

impl Runtime {
    pub fn new() -> Self {
        let module_manager = ModuleManager::new_empty_module_manager(BUILTIN_CODE_PATH);
        let new_environment = Box::new(Environment::new_empty_env());

        Runtime {
            module_manager,
            environment_stack: vec![new_environment],
            parsing_time_name_scope_stack: vec![HashMap::new()],
        }
    }
}

impl Runtime {
    fn push_parsing_time_name_scope(&mut self) {
        self.parsing_time_name_scope_stack.push(HashMap::new());
    }

    pub fn validate_name(
        &mut self,
        name: &str,
        current_line_file: LineFile,
    ) -> Result<(), RuntimeError> {
        if let Err(invalid_name_message) = is_valid_litex_name(name) {
            return Err(ParseRuntimeError(RuntimeErrorStruct::new(
                None,
                invalid_name_message,
                default_line_file(),
                None,
                vec![],
            ))
            .into());
        }

        for names_in_scope in self.parsing_time_name_scope_stack.iter().rev() {
            if let Some(_) = names_in_scope.get(name) {
                return Err(ParseRuntimeError(RuntimeErrorStruct::new(
                    None,
                    format!("name `{}` is already used", name),
                    current_line_file,
                    None,
                    vec![],
                ))
                .into());
            }
        }

        if self.is_name_used(name) {
            return Err(ParseRuntimeError(RuntimeErrorStruct::new(
                None,
                format!("name `{}` is already used", name),
                current_line_file,
                None,
                vec![],
            ))
            .into());
        }

        Ok(())
    }

    pub fn validate_name_for_mangled_fn_param(
        &mut self,
        name: &str,
        current_line_file: LineFile,
    ) -> Result<(), RuntimeError> {
        if let Err(invalid_name_message) = is_valid_mangled_fn_param_name(name) {
            return Err(ParseRuntimeError(RuntimeErrorStruct::new(
                None,
                invalid_name_message,
                default_line_file(),
                None,
                vec![],
            ))
            .into());
        }

        for names_in_scope in self.parsing_time_name_scope_stack.iter().rev() {
            if let Some(_) = names_in_scope.get(name) {
                return Err(ParseRuntimeError(RuntimeErrorStruct::new(
                    None,
                    format!("name `{}` is already used", name,),
                    current_line_file,
                    None,
                    vec![],
                ))
                .into());
            }
        }

        if self.is_name_used(name) {
            return Err(ParseRuntimeError(RuntimeErrorStruct::new(
                None,
                format!("name `{}` is already used", name),
                current_line_file,
                None,
                vec![],
            ))
            .into());
        }

        Ok(())
    }

    pub fn validate_name_and_insert_mangled_fn_param(
        &mut self,
        name: &str,
        (line, path): LineFile,
    ) -> Result<(), RuntimeError> {
        self.validate_name_for_mangled_fn_param(name, (line, path.clone()))?;
        if let Some(names_in_top_scope) = self.parsing_time_name_scope_stack.last_mut() {
            names_in_top_scope.insert(name.to_string(), (line, path.clone()));
        }
        Ok(())
    }

    pub fn register_collected_mangled_fn_param_names_for_def_parse(
        &mut self,
        names: &Vec<String>,
        line_file: LineFile,
    ) -> Result<(), RuntimeError> {
        for name in names {
            self.validate_name_and_insert_mangled_fn_param(name, line_file.clone())
                .map_err(|e| {
                    RuntimeError::from(ParseRuntimeError(RuntimeErrorStruct::new(
                        None,
                        String::new(),
                        line_file.clone(),
                        Some(e),
                        vec![],
                    )))
                })?;
        }
        Ok(())
    }

    pub fn register_mangled_fn_param_binding(
        &mut self,
        user_written_names: &[String],
        line_file: LineFile,
    ) -> Result<(Vec<String>, HashMap<String, Obj>), RuntimeError> {
        // 虽然本质上存的是 __param，但param本身也要符合litex命名规则，比如你不能让 param 是 1
        for name in user_written_names {
            if let Err(e) = is_valid_litex_name(name) {
                return Err(ParseRuntimeError(RuntimeErrorStruct::new(
                    None,
                    e,
                    line_file.clone(),
                    None,
                    vec![],
                ))
                .into());
            }
        }

        let (mangled, map) =
            mangled_fn_param_binding(user_written_names, DEFAULT_MANGLED_FN_PARAM_PREFIX);
        self.register_collected_mangled_fn_param_names_for_def_parse(&mangled, line_file)?;
        Ok((mangled, map))
    }

    fn pop_parsing_time_name_scope(&mut self) {
        self.parsing_time_name_scope_stack.pop();
    }

    pub fn validate_names_and_insert_into_top_parsing_time_name_scope(
        &mut self,
        names: &Vec<String>,
        line_file: LineFile,
    ) -> Result<(), RuntimeError> {
        for name in names {
            self.validate_name_and_insert_into_top_parsing_time_name_scope(
                name,
                line_file.clone(),
            )?;
        }
        Ok(())
    }

    pub fn validate_name_and_insert_into_top_parsing_time_name_scope(
        &mut self,
        name: &str,
        (line, path): LineFile,
    ) -> Result<(), RuntimeError> {
        self.validate_name(name, (line, path.clone()))?;
        self.register_name_into_name_scope(name, (line, path.clone()))?;
        Ok(())
    }

    pub fn register_name_into_name_scope(
        &mut self,
        name: &str,
        (line, path): LineFile,
    ) -> Result<(), RuntimeError> {
        if let Some(names_in_top_scope) = self.parsing_time_name_scope_stack.last_mut() {
            names_in_top_scope.insert(name.to_string(), (line, path.clone()));
        }
        Ok(())
    }
}

impl Runtime {
    pub fn new_file_path_new_env_new_name_scope(&mut self, path: &str) {
        let path_rc: Rc<str> = Rc::from(path);
        self.module_manager.run_file_paths.push(path_rc.clone());
        self.module_manager.current_file_index += 1;
        self.module_manager.display_entry_rc = Some(path_rc);
        self.module_manager.entry_path = path.to_string();
        self.push_parsing_time_name_scope();
        self.push_env();
    }

    pub fn is_name_used(&self, name: &str) -> bool {
        self.parsing_time_name_scope_stack
            .iter()
            .any(|scope| scope.contains_key(name))
    }
}

impl Runtime {
    pub fn top_level_env(&mut self) -> &mut Environment {
        let result = self.environment_stack.last_mut();
        match result {
            Some(environment) => environment,
            None => unreachable!("no top level environment"),
        }
    }
}

impl Runtime {
    fn push_env(&mut self) {
        let new_env = Box::new(Environment::new_empty_env());
        self.environment_stack.push(new_env);
    }

    fn pop_env(&mut self) {
        let last_env = self.environment_stack.last();

        match last_env {
            None => {
                unreachable!("no top level environment")
            }
            Some(_) => {
                self.environment_stack.pop();
            }
        }
    }

    /// 在临时子环境中执行闭包：`push_env` → `f` → `pop_env`；`Ok`/`Err` 都会弹出。
    /// 与手写 `push`/`pop` 等价；若闭包 panic，栈不会恢复（与手写相同）。
    pub fn run_in_local_env<T, E, F>(&mut self, f: F) -> Result<T, E>
    where
        F: FnOnce(&mut Self) -> Result<T, E>,
    {
        self.push_env();
        let result = f(self);
        self.pop_env();
        result
    }

    /// Same contract as `run_in_local_env`, but for `parsing_time_name_scope_stack`.
    pub fn run_in_local_parsing_time_name_scope<T, E, F>(&mut self, f: F) -> Result<T, E>
    where
        F: FnOnce(&mut Self) -> Result<T, E>,
    {
        self.push_parsing_time_name_scope();
        let result = f(self);
        self.pop_parsing_time_name_scope();
        result
    }
}

impl Runtime {
    pub fn is_name_used_for_identifier_and_field_access(&self, name: &str) -> bool {
        if is_builtin_identifier_name(name) {
            return true;
        }

        for env in self.iter_environments_from_top() {
            if env.defined_identifiers.contains_key(name) {
                return true;
            }
        }

        false
    }

    pub fn is_name_used_for_prop(&self, name: &str) -> bool {
        return self.get_prop_definition_by_name(name).is_some();
    }

    pub fn is_name_used_for_abstract_prop(&self, name: &str) -> bool {
        if is_builtin_predicate(name) {
            return true;
        }

        return self.get_abstract_prop_definition_by_name(name).is_some();
    }

    pub fn is_name_used_for_struct(&self, name: &str) -> bool {
        return self.get_definition_of_struct_by_name(name).is_some();
    }

    pub fn is_name_used_for_family(&self, name: &str) -> bool {
        return self.get_family_definition_by_name(name).is_some();
    }

    pub fn is_name_used_for_algo(&self, name: &str) -> bool {
        return self.get_algo_definition_by_name(name).is_some();
    }
}

impl Runtime {
    pub fn new_file_and_update_runtime_with_file_content(&mut self, path: &str) {
        let path_rc: Rc<str> = Rc::from(path);
        self.module_manager.run_file_paths.push(path_rc.clone());
        self.module_manager.current_file_index = self.module_manager.run_file_paths.len() - 1;
        self.module_manager.display_entry_rc = Some(path_rc);
        self.module_manager.entry_path = path.to_string();
    }

    pub fn change_file_index_to(&mut self, file_index: usize) {
        self.module_manager.current_file_index = file_index;
    }
}

impl Runtime {
    pub fn store_tuple_obj_and_cart(
        &mut self,
        name: &str,
        tuple: Option<Tuple>,
        cart: Option<Cart>,
        line_file: LineFile,
    ) {
        let known_tuple_objs = &mut self.top_level_env().known_objs_equal_to_tuple;
        let old_tuple_and_cart = known_tuple_objs.get(name).cloned();

        let merged_tuple = match (tuple, old_tuple_and_cart.as_ref()) {
            (Some(new_tuple), _) => Some(new_tuple),
            (None, Some((old_tuple, _, _))) => old_tuple.clone(),
            (None, None) => None,
        };
        let merged_cart = match (cart, old_tuple_and_cart.as_ref()) {
            (Some(new_cart), _) => Some(new_cart),
            (None, Some((_, old_cart, _))) => old_cart.clone(),
            (None, None) => None,
        };
        let merged_line_file = line_file;

        known_tuple_objs.insert(
            name.to_string(),
            (merged_tuple, merged_cart, merged_line_file),
        );
    }

    pub fn store_known_cart_obj(&mut self, name: &str, cart: Cart, line_file: LineFile) {
        self.top_level_env()
            .known_objs_equal_to_cart
            .insert(name.to_string(), (cart, line_file));
    }

    pub fn store_known_finite_seq_list_obj(
        &mut self,
        name: &str,
        list: FiniteSeqListObj,
        member_of_finite_seq_set: Option<FiniteSeqSet>,
        line_file: LineFile,
    ) {
        let map = &mut self.top_level_env().known_objs_equal_to_finite_seq_list;
        let old = map.get(name).cloned();
        let merged_member = match (member_of_finite_seq_set, old.as_ref()) {
            (Some(new_s), _) => Some(new_s),
            (None, Some((_, Some(old_s), _))) => Some(old_s.clone()),
            (None, _) => None,
        };
        map.insert(
            name.to_string(),
            (list, merged_member, line_file),
        );
    }

    pub fn store_known_matrix_list_obj(
        &mut self,
        name: &str,
        matrix: MatrixListObj,
        member_of_matrix_set: Option<MatrixSet>,
        line_file: LineFile,
    ) {
        let map = &mut self.top_level_env().known_objs_equal_to_matrix_list;
        let old = map.get(name).cloned();
        let merged_member = match (member_of_matrix_set, old.as_ref()) {
            (Some(new_s), _) => Some(new_s),
            (None, Some((_, Some(old_s), _))) => Some(old_s.clone()),
            (None, _) => None,
        };
        map.insert(
            name.to_string(),
            (matrix, merged_member, line_file),
        );
    }

    pub fn matrix_set_to_fn_set(&self, ms: &MatrixSet, line_file: LineFile) -> FnSet {
        let pair = self.generate_random_unused_names(2);
        let p1 = format!("{}{}", DEFAULT_MANGLED_FN_PARAM_PREFIX, pair[0]);
        let p2 = format!("{}{}", DEFAULT_MANGLED_FN_PARAM_PREFIX, pair[1]);
        FnSet::new(
            vec![
                ParamGroupWithSet::new(vec![p1.clone()], StandardSet::NPos.into()),
                ParamGroupWithSet::new(vec![p2.clone()], StandardSet::NPos.into()),
            ],
            vec![
                AtomicFact::from(LessEqualFact::new(
                    p1.into(),
                    (*ms.row_len).clone(),
                    line_file.clone(),
                ))
                .into(),
                AtomicFact::from(LessEqualFact::new(
                    p2.into(),
                    (*ms.col_len).clone(),
                    line_file.clone(),
                ))
                .into(),
            ],
            (*ms.set).clone(),
        )
    }

    pub fn finite_seq_set_to_fn_set(&self, fs: &FiniteSeqSet, line_file: LineFile) -> FnSet {
        let param = format!(
            "{}{}",
            DEFAULT_MANGLED_FN_PARAM_PREFIX,
            self.generate_random_unused_name()
        );
        FnSet::new(
            vec![ParamGroupWithSet::new(
                vec![param.clone()],
                StandardSet::NPos.into(),
            )],
            vec![
                AtomicFact::from(LessEqualFact::new(param.into(), (*fs.n).clone(), line_file))
                    .into(),
            ],
            (*fs.set).clone(),
        )
    }

    pub fn seq_set_to_fn_set(&self, ss: &SeqSet, _line_file: LineFile) -> FnSet {
        let param = format!(
            "{}{}",
            DEFAULT_MANGLED_FN_PARAM_PREFIX,
            self.generate_random_unused_name()
        );
        FnSet::new(
            vec![ParamGroupWithSet::new(
                vec![param.clone()],
                StandardSet::NPos.into(),
            )],
            vec![],
            (*ss.set).clone(),
        )
    }

    /// Same mangling as `parse_fn_set`: surface `x` becomes stored `__x`, etc.
    pub fn finite_seq_set_to_fn_set_from_surface_dom_param(
        &self,
        fs: &FiniteSeqSet,
        line_file: LineFile,
        surface_dom_param: &str,
    ) -> Result<FnSet, RuntimeError> {
        let params = vec![ParamGroupWithSet::new(
            vec![surface_dom_param.to_string()],
            StandardSet::NPos.into(),
        )];
        let dom_facts: Vec<OrAndChainAtomicFact> = vec![OrAndChainAtomicFact::AtomicFact(
            LessEqualFact::new(
                Identifier::new(surface_dom_param.to_string()).into(),
                (*fs.n).clone(),
                line_file,
            )
            .into(),
        )];
        self.new_fn_set_and_add_mangled_prefix(params, dom_facts, (*fs.set).clone())
    }

    pub fn store_well_defined_obj_cache(&mut self, obj: &Obj) {
        self.top_level_env()
            .cache_well_defined_obj
            .insert(obj.to_string(), ());
    }
}

impl Runtime {
    pub fn new_fn_set_and_add_mangled_prefix(
        &self,
        params_and_their_sets: Vec<ParamGroupWithSet>,
        dom_facts: Vec<OrAndChainAtomicFact>,
        ret_set: Obj,
    ) -> Result<FnSet, RuntimeError> {
        let names = ParamGroupWithSet::collect_param_names(&params_and_their_sets);
        let (new_param_names, param_arg_map) =
            mangled_fn_param_binding(&names, DEFAULT_MANGLED_FN_PARAM_PREFIX);
        let mut flat_stored_idx: usize = 0;
        let mut new_def_with_set: Vec<ParamGroupWithSet> = Vec::new();
        for param_group in &params_and_their_sets {
            let mut new_params: Vec<String> = Vec::new();
            for _ in 0..param_group.params.len() {
                new_params.push(new_param_names[flat_stored_idx].clone());
                flat_stored_idx += 1;
            }
            new_def_with_set.push(ParamGroupWithSet::new(new_params, param_group.set.clone()));
        }
        let mut dom_stored = Vec::with_capacity(dom_facts.len());
        for d in &dom_facts {
            dom_stored.push(self.inst_or_and_chain_atomic_fact(d, &param_arg_map)?);
        }
        let ret_stored = self.inst_obj(&ret_set, &param_arg_map)?;
        Ok(FnSet::new(new_def_with_set, dom_stored, ret_stored))
    }

    pub fn add_mangled_prefix_to_fn_set_clause(
        &self,
        clause: &FnSetClause,
        _line_file: LineFile,
    ) -> Result<FnSet, RuntimeError> {
        self.new_fn_set_and_add_mangled_prefix(
            clause.params_def_with_set.clone(),
            clause.dom_facts.clone(),
            clause.ret_set.clone(),
        )
    }
}

impl Runtime {
    /// Like [`ParamDefWithType::param_defs_and_args_to_param_to_arg_map`], but when a
    /// parameter has [`ParamType::Struct`] and the argument is a [`Obj::Tuple`], loads the struct
    /// definition and inserts `param.field_name -> tuple[number_of_params + i]` for each field in
    /// definition order (same layout as [`Runtime::inst_field_access_on_struct_tuple`]), and
    /// registers the struct instance for field access on the tuple root.
    pub fn params_to_arg_map(
        &mut self,
        param_defs: &ParamDefWithType,
        args: &[Obj],
    ) -> Result<HashMap<String, Obj>, RuntimeError> {
        let param_names = param_defs.collect_param_names();
        if param_names.len() != args.len() {
            return Err(InstantiateRuntimeError(RuntimeErrorStruct::new(
                None,
                format!(
                    "params_to_arg_map: expected {} argument(s), got {}",
                    param_names.len(),
                    args.len()
                ),
                default_line_file(),
                None,
                vec![],
            ))
            .into());
        }

        let mut flat_types: Vec<ParamType> = Vec::with_capacity(param_names.len());
        for param_def in param_defs.groups.iter() {
            for _ in param_def.params.iter() {
                flat_types.push(param_def.param_type.clone());
            }
        }

        let mut result: HashMap<String, Obj> = HashMap::new();
        for ((param_name, param_type), arg) in
            param_names.iter().zip(flat_types.iter()).zip(args.iter())
        {
            match param_type {
                ParamType::Struct(struct_ty) => {
                    let struct_name = struct_ty.name.to_string();
                    let Some(def) = self.get_cloned_definition_of_struct(&struct_name) else {
                        result.insert(param_name.clone(), arg.clone());
                        continue;
                    };
                    if let Obj::Tuple(t) = arg {
                        let expected_len = def.number_of_params() + def.fields.len();
                        if t.args.len() != expected_len {
                            return Err(
                                InstantiateRuntimeError(RuntimeErrorStruct::new(
                                    None,
                                    format!(
                                        "params_to_arg_map: tuple for `{}` has {} component(s), struct `{}` expects {}",
                                        param_name,
                                        t.args.len(),
                                        struct_name,
                                        expected_len
                                    ),
                                    default_line_file(),
                                    None,
                                    vec![],
                                ))
                                .into(),
                            );
                        }
                        self.register_param_as_struct_instance(param_name, struct_ty.clone());
                        result.insert(param_name.clone(), arg.clone());
                        for (fi, (field_name, _)) in def.fields.iter().enumerate() {
                            let ti = def.number_of_params() + fi;
                            let component = (*t.args[ti]).clone();
                            result
                                .insert(field_access_to_string(param_name, field_name), component);
                        }
                    } else {
                        result.insert(param_name.clone(), arg.clone());
                    }
                }
                _ => {
                    result.insert(param_name.clone(), arg.clone());
                }
            }
        }
        Ok(result)
    }

    /// [`DefStructStmt::dom_facts`] under type arguments, then [`DefStructStmt::facts`] (`<=>:`) with
    /// [`SELF`] replaced by `param_name`, in source order.
    pub fn instantiated_struct_def_or_and_facts_for_def(
        &self,
        struct_ty: &StructObj,
        def: &DefStructStmt,
        param_name: &str,
    ) -> Result<Vec<OrAndChainAtomicFact>, RuntimeError> {
        let base_map = def
            .param_defs
            .param_defs_and_args_to_param_to_arg_map(struct_ty.args.as_slice());
        let mut out = Vec::new();
        for fact in def.dom_facts.iter() {
            out.push(self.inst_or_and_chain_atomic_fact(fact, &base_map)?);
        }
        let mut map_with_self = base_map.clone();
        map_with_self.insert(SELF.to_string(), param_name.to_string().into());
        for fact in def.facts.iter() {
            out.push(self.inst_or_and_chain_atomic_fact(fact, &map_with_self)?);
        }
        Ok(out)
    }
}