nixjit/src/ty/internal/attrset.rs

use std::rc::Rc;
use std::collections::HashSet;

use derive_more::Constructor;
use rpds::HashTrieMap;

use crate::error::Result;
use crate::vm::{Env, VM};
use super::super::public as p;
use super::Value;

#[repr(C)]
#[derive(Debug, Constructor, Clone, PartialEq)]
pub struct AttrSet<'vm> {
    data: HashTrieMap<usize, Value<'vm>>,
}

impl<'vm> AttrSet<'vm> {
    pub fn empty() -> Self {
        AttrSet {
            data: HashTrieMap::new(),
        }
    }

    pub fn push_attr_force(&mut self, sym: usize, val: Value<'vm>) {
        self.data.insert_mut(sym, val);
    }

    pub fn push_attr(&mut self, sym: usize, val: Value<'vm>) {
        if self.data.get_mut(&sym).is_some() {
            todo!()
        }
        self.data.insert_mut(sym, val);
    }

    pub fn select(&self, sym: usize) -> Option<Value<'vm>> {
        self.data.get(&sym).cloned()
    }

    pub fn has_attr(&self, sym: usize) -> bool {
        self.data.get(&sym).is_some()
    }

    pub fn capture(&mut self, env: &Rc<Env<'vm>>) {
        self
            .data
            .iter()
            .for_each(|(_, v)| {
                    match v.clone() {
                        Value::Thunk(ref thunk) => {
                            thunk.capture(env.clone());
                        }
                        _ => ()
                    }
            })
    }

    pub fn update(&mut self, other: &AttrSet<'vm>) {
        for (k, v) in other.data.iter() {
            self.push_attr_force(k.clone(), v.clone())
        }
    }

    pub fn update_rec(&mut self, other: &RecAttrSet<'vm>) {
        for (k, v) in other.data.map.borrow().iter() {
            self.push_attr_force(k.clone(), v.clone())
        }
    }

    pub fn as_inner(&self) -> &HashTrieMap<usize, Value<'vm>> {
        &self.data
    }

    pub fn force_deep(&mut self, vm: &'vm VM<'_>) -> Result<()> {
        let mut map: Vec<_> = self
            .data
            .into_iter()
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();
        for (_, v) in map.iter_mut() {
            v.force_deep(vm)?;
        }
        self.data = map.into_iter().collect();
        Ok(())
    }

    pub fn to_public(&self, vm: &'vm VM, seen: &mut HashSet<Value<'vm>>) -> p::Value {
        p::Value::AttrSet(p::AttrSet::new(
            self.data
                .iter()
                .map(|(&sym, value)| (vm.get_sym(sym), value.clone().to_public(vm, seen)))
                .collect(),
        ))
    }
}

#[derive(Debug, Constructor, Clone, PartialEq)]
pub struct RecAttrSet<'vm> {
    data: Rc<Env<'vm>>,
}

impl<'vm> RecAttrSet<'vm> {
    pub fn empty() -> Self {
        RecAttrSet {
            data: Rc::default(),
        }
    }

    pub fn push_attr_force(&mut self, sym: usize, val: Value<'vm>) {
        self.data.insert(sym, val);
    }

    pub fn push_attr(&mut self, sym: usize, val: Value<'vm>) {
        if self.data.lookup(sym).is_some() {
            todo!()
        }
        self.data.insert(sym, val);
    }

    pub fn select(&self, sym: usize) -> Option<Value<'vm>> {
        self.data.lookup(sym)
    }

    pub fn has_attr(&self, sym: usize) -> bool {
        self.data.lookup(sym).is_some()
    }

    pub fn update(&mut self, other: RecAttrSet<'vm>) {
        for (k, v) in other.data.map.borrow().iter() {
            self.push_attr_force(k.clone(), v.clone())
        }
    }

    pub fn update_normal(&self, other: &AttrSet<'vm>) -> AttrSet<'vm> {
        let map = self
            .data
            .map
            .borrow()
            .into_iter()
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();
        let mut new = AttrSet::new(map);
        for (k, v) in other.data.iter() {
            new.push_attr_force(k.clone(), v.clone())
        }
        new
    }

    pub fn into_inner(self) -> Rc<Env<'vm>> {
        self.data
    }

    pub fn from_inner(data: Rc<Env<'vm>>) -> Self {
        RecAttrSet { data }
    }

    pub fn force_deep(&mut self, vm: &'vm VM<'_>) -> Result<()> {
        let mut map: Vec<_> = self
            .data
            .map
            .borrow()
            .into_iter()
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();
        map.iter_mut()
            .map(|(_, v)| v.force_deep(vm).map(|_| ()))
            .find(|v| v.is_err())
            .map_or(Ok(()), |err| err)?;
        *self.data.map.borrow_mut() = map.into_iter().collect();
        Ok(())
    }

    pub fn to_public(&self, vm: &'vm VM, seen: &mut HashSet<Value<'vm>>) -> p::Value {
        p::Value::AttrSet(p::AttrSet::new(
            self.data
                .map
                .borrow()
                .iter()
                .map(|(&sym, value)| (vm.get_sym(sym), value.clone().to_public(vm, seen)))
                .collect(),
        ))
    }
}